import maya.api.OpenMaya as om
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
import cv2
import os

from V2A.core import process

# 註冊資訊
maya_use_NewApi = True

# ARKit 52 個標準表情清單
ARKIT_SHAPES = [
    "browDownLeft", "browDownRight", "browInnerUp", "browOuterUpLeft", "browOuterUpRight",
    "cheekPuff", "cheekSquintLeft", "cheekSquintRight", "eyeBlinkLeft", "eyeBlinkRight",
    "eyeLookDownLeft", "eyeLookDownRight", "eyeLookInLeft", "eyeLookInRight", "eyeLookOutLeft",
    "eyeLookOutRight", "eyeLookUpLeft", "eyeLookUpRight", "eyeSquintLeft", "eyeSquintRight",
    "eyeWideLeft", "eyeWideRight", "jawForward", "jawLeft", "jawOpen", "jawRight",
    "mouthClose", "mouthDimpleLeft", "mouthDimpleRight", "mouthFrownLeft", "mouthFrownRight",
    "mouthFunnel", "mouthLeft", "mouthLowerDownLeft", "mouthLowerDownRight", "mouthPressLeft",
    "mouthPressRight", "mouthPucker", "mouthRight", "mouthRollLower", "mouthRollUpper",
    "mouthShrugLower", "mouthShrugUpper", "mouthSmileLeft", "mouthSmileRight", "mouthStretchLeft",
    "mouthStretchRight", "mouthUpperUpLeft", "mouthUpperUpRight", "noseSneerLeft", "noseSneerRight", "tongueOut"
]

class Video2ARKitNode(om.MPxNode):
    NODE_NAME = "Video2ARKit"
    NODE_ID = om.MTypeId(0x00123456) # 測試用 ID

    # 屬性定義
    aInTime = om.MObject()
    aVideoPath = om.MObject()
    aModelPath = om.MObject()
    aProcessTrigger = om.MObject()
    
    # 輸出屬性字典
    output_attrs = {}

    def __init__(self):
        om.MPxNode.__init__(self)
        self._cache = {} # 格式: { frame_index: { shape_name: value } }

    def compute(self, plug, data_block):
        # 如果正在請求輸出屬性 (52個中的任何一個)
        if plug in self.output_attrs.values():
            
            # 獲取輸入值
            current_time = data_block.inputValue(self.aInTime).asMTime().asUnits(om.MTime.kFilm) # 假設 24fps
            trigger = data_block.inputValue(self.aProcessTrigger).asBool()
            video_path = data_block.inputValue(self.aVideoPath).asString()
            model_path = data_block.inputValue(self.aModelPath).asString()

            # 檢查是否需要重新分析影片
            if trigger and not self._cache:
                self._cache = process.process_video(video_path, model_path)

            # 獲取當前幀數的權重
            frame_idx = int(current_time)
            frame_data = self._cache.get(frame_idx, {})

            # 輸出對應的值
            for name, attr_obj in self.output_attrs.items():
                if plug == attr_obj:
                    val = frame_data.get(name, 0.0)
                    handle = data_block.outputValue(attr_obj)
                    handle.setFloat(val)
                    data_block.setClean(plug)
                    return
                
    def setDependencyDirty(self, plug, affectedPlugs):
        if plug == self.aProcessTrigger:
            self._cache = {}
        super().setDependencyDirty(plug, affectedPlugs)

    def _reload_cache(self):
        current_value = self.aProcessTrigger.asBool()
        self.aProcessTrigger.setBool(current_value)
        om.MGlobal.displayInfo("Reload cache complete. cache populated.")

    @staticmethod
    def creator():
        return Video2ARKitNode()

    @staticmethod
    def initialize():
        nAttr = om.MFnNumericAttribute()
        tAttr = om.MFnTypedAttribute()
        uAttr = om.MFnUnitAttribute()

        # 輸入：時間
        Video2ARKitNode.aInTime = uAttr.create("inTime", "it", om.MFnUnitAttribute.kTime, 0.0)
        om.MPxNode.addAttribute(Video2ARKitNode.aInTime)

        # 輸入：影片路徑
        Video2ARKitNode.aVideoPath = tAttr.create("videoPath", "vp", om.MFnData.kString)
        om.MPxNode.addAttribute(Video2ARKitNode.aVideoPath)

        # 輸入：模型路徑
        data = om.MFnStringData()
        default_model_path = data.create(os.getenv('V2A_MODEL_PATH'))
        Video2ARKitNode.aModelPath = tAttr.create("modelPath", "mp", om.MFnData.kString, default_model_path)
        om.MPxNode.addAttribute(Video2ARKitNode.aModelPath)

        # 輸入：觸發器 (勾選後開始分析)
        Video2ARKitNode.aProcessTrigger = nAttr.create("processTrigger", "pt", om.MFnNumericData.kBoolean, False)
        nAttr.keyable = True
        nAttr.hidden = True
        om.MPxNode.addAttribute(Video2ARKitNode.aProcessTrigger)

        # 輸出：52 個表情係數
        for shape_name in ARKIT_SHAPES:
            attr = nAttr.create(shape_name, shape_name, om.MFnNumericData.kFloat, 0.0)
            nAttr.writable = False
            nAttr.storable = False
            om.MPxNode.addAttribute(attr)
            Video2ARKitNode.output_attrs[shape_name] = attr
            # 建立依賴
            om.MPxNode.attributeAffects(Video2ARKitNode.aVideoPath, attr)
            om.MPxNode.attributeAffects(Video2ARKitNode.aInTime, attr)
            om.MPxNode.attributeAffects(Video2ARKitNode.aProcessTrigger, attr)

