#include "Video2ARKitNode.h"

#include <maya/MPlug.h>
#include <maya/MDataBlock.h>
#include <maya/MDataHandle.h>
#include <maya/MGlobal.h>
#include <maya/MTime.h>
#include <maya/MFnUnitAttribute.h>
#include <maya/MFnTypedAttribute.h>
#include <maya/MFnNumericAttribute.h>
#include <maya/MFnCompoundAttribute.h>

#include <fstream>
#include <streambuf>

// --- Dependencies (Must be installed on your system) ---
#include <opencv2/opencv.hpp>
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/classification.pb.h"
#include "mediapipe/framework/port/parse_text_proto.h"
#include "mediapipe/framework/port/status.h"


// --- Helper Function for Video Processing ---
BlendshapeCache processVideo_impl(const MString& videoPath)
{
    BlendshapeCache cache;
    MGlobal::displayInfo("Processing video: " + videoPath);

    std::string graph_path = "D:/workspace/Video2ARkit/src/face_blendshapes_graph.pbtxt";
    std::ifstream graph_file(graph_path);
    if (!graph_file.is_open()) {
        MGlobal::displayError("Failed to open MediaPipe graph file: " + MString(graph_path.c_str()));
        return cache;
    }
    std::string graph_config_string((std::istreambuf_iterator<char>(graph_file)), std::istreambuf_iterator<char>());

    mediapipe::CalculatorGraphConfig config;
    if (!google::protobuf::TextFormat::ParseFromString(graph_config_string, &config)) {
        MGlobal::displayError("Failed to parse MediaPipe graph config.");
        return cache;
    }

    mediapipe::CalculatorGraph graph;
    auto status = graph.Initialize(config);
    if (!status.ok()) {
        MGlobal::displayError("Failed to initialize MediaPipe graph.");
        return cache;
    }

    auto poller_status = graph.AddOutputStreamPoller("face_blendshapes");
    if (!poller_status.ok()) {
        MGlobal::displayError("Failed to add output stream poller.");
        return cache;
    }
    mediapipe::OutputStreamPoller poller = std::move(poller_status.value());

    status = graph.StartRun({});
    if (!status.ok()) {
        MGlobal::displayError("Failed to start MediaPipe graph.");
        return cache;
    }

    cv::VideoCapture cap(videoPath.asChar());
    if (!cap.isOpened()) {
        MGlobal::displayError("Failed to open video file with OpenCV: " + videoPath);
        return cache;
    }

    long frame_count = 0;
    while (true) {
        cv::Mat frame;
        cap >> frame;
        if (frame.empty()) {
            break;
        }

        auto input_frame = std::make_unique<mediapipe::ImageFrame>(
            mediapipe::ImageFormat::SRGB, frame.cols, frame.rows, mediapipe::ImageFrame::kDefaultAlignmentBoundary);
        cv::cvtColor(frame, input_frame->MatView(), cv::COLOR_BGR2RGB);

        mediapipe::Timestamp timestamp(frame_count);
        status = graph.AddPacketToInputStream("input_video", mediapipe::MakePacket<mediapipe::ImageFrame>(std::move(input_frame)).At(timestamp));
        if (!status.ok()) {
            MGlobal::displayError("Failed to add packet to input stream.");
            break;
        }

        mediapipe::Packet packet;
        if (poller.Next(&packet)) {
            auto& classification_list_vec = packet.Get<std::vector<mediapipe::ClassificationList>>();
            if (!classification_list_vec.empty()) {
                BlendshapeFrame frameData;
                const auto& blendshapes = classification_list_vec[0];
                for (int i = 0; i < blendshapes.classification_size(); ++i) {
                    const auto& blendshape = blendshapes.classification(i);
                    frameData[blendshape.label()] = blendshape.score();
                }
                cache.push_back(frameData);
            }
        }

        frame_count++;
        if (frame_count % 100 == 0) {
            MGlobal::displayInfo("Processed " + MString() + frame_count + " frames...");
        }
    }

    graph.CloseInputStream("input_video");
    graph.WaitUntilDone();
    cap.release();

    MGlobal::displayInfo("Finished processing video. Cached " + MString() + (long)cache.size() + " frames.");
    return cache;
}


// --- Static Member Initialization ---
MTypeId Video2ARKitNode::id(0x00081031);
BlendshapeCache Video2ARKitNode::m_blendshapeCache;
MString Video2ARKitNode::m_cachedVideoPath;

MObject Video2ARKitNode::videoPath, Video2ARKitNode::currentTime, Video2ARKitNode::enable, Video2ARKitNode::outputBlendshapes;
MObject Video2ARKitNode::browDownLeft, Video2ARKitNode::browDownRight, Video2ARKitNode::browInnerUp, 
        Video2ARKitNode::browOuterUpLeft, Video2ARKitNode::browOuterUpRight, Video2ARKitNode::cheekPuff, 
        Video2ARKitNode::cheekSquintLeft, Video2ARKitNode::cheekSquintRight, Video2ARKitNode::eyeBlinkLeft, 
        Video2ARKitNode::eyeBlinkRight, Video2ARKitNode::eyeLookDownLeft, Video2ARKitNode::eyeLookDownRight, 
        Video2ARKitNode::eyeLookInLeft, Video2ARKitNode::eyeLookInRight, Video2ARKitNode::eyeLookOutLeft, 
        Video2ARKitNode::eyeLookOutRight, Video2ARKitNode::eyeLookUpLeft, Video2ARKitNode::eyeLookUpRight, 
        Video2ARKitNode::eyeSquintLeft, Video2ARKitNode::eyeSquintRight, Video2ARKitNode::eyeWideLeft, 
        Video2ARKitNode::eyeWideRight, Video2ARKitNode::jawForward, Video2ARKitNode::jawLeft, 
        Video2ARKitNode::jawOpen, Video2ARKitNode::jawRight, Video2ARKitNode::mouthClose, 
        Video2ARKitNode::mouthDimpleLeft, Video2ARKitNode::mouthDimpleRight, Video2ARKitNode::mouthFrownLeft, 
        Video2ARKitNode::mouthFrownRight, Video2ARKitNode::mouthFunnel, Video2ARKitNode::mouthLowerDownLeft, 
        Video2ARKitNode::mouthLowerDownRight, Video2ARKitNode::mouthPressLeft, Video2ARKitNode::mouthPressRight, 
        Video2ARKitNode::mouthPucker, Video2ARKitNode::mouthRight, Video2ARKitNode::mouthRollLower, 
        Video2ARKitNode::mouthRollUpper, Video2ARKitNode::mouthShrugLower, Video2ARKitNode::mouthShrugUpper, 
        Video2ARKitNode::mouthSmileLeft, Video2ARKitNode::mouthSmileRight, Video2ARKitNode::mouthStretchLeft, 
        Video2ARKitNode::mouthStretchRight, Video2ARKitNode::mouthUpperUpLeft, Video2ARKitNode::mouthUpperUpRight, 
        Video2ARKitNode::noseSneerLeft, Video2ARKitNode::noseSneerRight;


// --- Node Method Implementations ---

Video2ARKitNode::Video2ARKitNode() {}
Video2ARKitNode::~Video2ARKitNode() {}

void* Video2ARKitNode::creator() {
    return new Video2ARKitNode();
}

MStatus Video2ARKitNode::compute(const MPlug& plug, MDataBlock& data) {
    if (plug.parent() != outputBlendshapes) {
        return MS::kUnknownParameter;
    }

    MDataHandle pathHandle = data.inputValue(videoPath);
    MString path = pathHandle.asString();

    if (path != m_cachedVideoPath && path.length() > 0) {
        m_blendshapeCache = processVideo_impl(path);
        m_cachedVideoPath = path;
    }

    MDataHandle outputCompoundHandle = data.outputValue(outputBlendshapes);
    MDataHandle enableHandle = data.inputValue(enable);
    bool isEnabled = enableHandle.asBool();

    if (isEnabled && !m_blendshapeCache.empty()) {
        MDataHandle timeHandle = data.inputValue(currentTime);
        MTime time = timeHandle.asTime();
        int frame = static_cast<int>(time.as(MTime::kFilm));

        if (frame >= 0 && frame < m_blendshapeCache.size()) {
            const auto& frameData = m_blendshapeCache[frame];

            #define SET_OUTPUT_VALUE(attr, name) \
                if (frameData.count(name)) { \
                    outputCompoundHandle.child(attr).setFloat(frameData.at(name)); \
                } else { \
                    outputCompoundHandle.child(attr).setFloat(0.0f); \
                }

            SET_OUTPUT_VALUE(browDownLeft, "browDown_L"); SET_OUTPUT_VALUE(browDownRight, "browDown_R");
            SET_OUTPUT_VALUE(browInnerUp, "browInnerUp"); SET_OUTPUT_VALUE(browOuterUpLeft, "browOuterUp_L");
            SET_OUTPUT_VALUE(browOuterUpRight, "browOuterUp_R"); SET_OUTPUT_VALUE(cheekPuff, "cheekPuff");
            SET_OUTPUT_VALUE(cheekSquintLeft, "cheekSquint_L"); SET_OUTPUT_VALUE(cheekSquintRight, "cheekSquint_R");
            SET_OUTPUT_VALUE(eyeBlinkLeft, "eyeBlink_L"); SET_OUTPUT_VALUE(eyeBlinkRight, "eyeBlink_R");
            SET_OUTPUT_VALUE(eyeLookDownLeft, "eyeLookDown_L"); SET_OUTPUT_VALUE(eyeLookDownRight, "eyeLookDown_R");
            SET_OUTPUT_VALUE(eyeLookInLeft, "eyeLookIn_L"); SET_OUTPUT_VALUE(eyeLookInRight, "eyeLookIn_R");
            SET_OUTPUT_VALUE(eyeLookOutLeft, "eyeLookOut_L"); SET_OUTPUT_VALUE(eyeLookOutRight, "eyeLookOut_R");
            SET_OUTPUT_VALUE(eyeLookUpLeft, "eyeLookUp_L"); SET_OUTPUT_VALUE(eyeLookUpRight, "eyeLookUp_R");
            SET_OUTPUT_VALUE(eyeSquintLeft, "eyeSquint_L"); SET_OUTPUT_VALUE(eyeSquintRight, "eyeSquint_R");
            SET_OUTPUT_VALUE(eyeWideLeft, "eyeWide_L"); SET_OUTPUT_VALUE(eyeWideRight, "eyeWide_R");
            SET_OUTPUT_VALUE(jawForward, "jawForward"); SET_OUTPUT_VALUE(jawLeft, "jawLeft");
            SET_OUTPUT_VALUE(jawOpen, "jawOpen"); SET_OUTPUT_VALUE(jawRight, "jawRight");
            SET_OUTPUT_VALUE(mouthClose, "mouthClose"); SET_OUTPUT_VALUE(mouthDimpleLeft, "mouthDimple_L");
            SET_OUTPUT_VALUE(mouthDimpleRight, "mouthDimple_R"); SET_OUTPUT_VALUE(mouthFrownLeft, "mouthFrown_L");
            SET_OUTPUT_VALUE(mouthFrownRight, "mouthFrown_R"); SET_OUTPUT_VALUE(mouthFunnel, "mouthFunnel");
            SET_OUTPUT_VALUE(mouthLowerDownLeft, "mouthLowerDown_L"); SET_OUTPUT_VALUE(mouthLowerDownRight, "mouthLowerDown_R");
            SET_OUTPUT_VALUE(mouthPressLeft, "mouthPress_L"); SET_OUTPUT_VALUE(mouthPressRight, "mouthPress_R");
            SET_OUTPUT_VALUE(mouthPucker, "mouthPucker"); SET_OUTPUT_VALUE(mouthRight, "mouthRight");
            SET_OUTPUT_VALUE(mouthRollLower, "mouthRollLower"); SET_OUTPUT_VALUE(mouthRollUpper, "mouthRollUpper");
            SET_OUTPUT_VALUE(mouthShrugLower, "mouthShrugLower"); SET_OUTPUT_VALUE(mouthShrugUpper, "mouthShrugUpper");
            SET_OUTPUT_VALUE(mouthSmileLeft, "mouthSmile_L"); SET_OUTPUT_VALUE(mouthSmileRight, "mouthSmile_R");
            SET_OUTPUT_VALUE(mouthStretchLeft, "mouthStretch_L"); SET_OUTPUT_VALUE(mouthStretchRight, "mouthStretch_R");
            SET_OUTPUT_VALUE(mouthUpperUpLeft, "mouthUpperUp_L"); SET_OUTPUT_VALUE(mouthUpperUpRight, "mouthUpperUp_R");
            SET_OUTPUT_VALUE(noseSneerLeft, "noseSneer_L"); SET_OUTPUT_VALUE(noseSneerRight, "noseSneer_R");
        }
    } else {
        #define ZERO_OUTPUT_VALUE(attr) outputCompoundHandle.child(attr).setFloat(0.0f)
        // ... Zero out all 52 attributes ...
    }

    outputCompoundHandle.setClean();
    data.setClean(plug);
    return MS::kSuccess;
}

MStatus Video2ARKitNode::initialize() {
    MFnNumericAttribute nAttr;
    MFnTypedAttribute tAttr;
    MFnUnitAttribute uAttr;
    MFnCompoundAttribute cAttr;

    videoPath = tAttr.create("videoPath", "vp", MFnData::kString); tAttr.setStorable(true); addAttribute(videoPath);
    currentTime = uAttr.create("currentTime", "ct", MFnUnitAttribute::kTime, 0.0); uAttr.setStorable(true); addAttribute(currentTime);
    enable = nAttr.create("enable", "en", MFnNumericData::kBoolean, 1); nAttr.setKeyable(true); addAttribute(enable);

    #define CREATE_BLENDSHAPE_ATTR(longName, shortName, attrObject) \
        attrObject = nAttr.create(longName, shortName, MFnNumericData::kFloat, 0.0); \
        nAttr.setWritable(false); nAttr.setStorable(false);

    CREATE_BLENDSHAPE_ATTR("browDownLeft", "bdl", browDownLeft); CREATE_BLENDSHAPE_ATTR("browDownRight", "bdr", browDownRight);
    CREATE_BLENDSHAPE_ATTR("browInnerUp", "biu", browInnerUp); CREATE_BLENDSHAPE_ATTR("browOuterUpLeft", "boul", browOuterUpLeft);
    CREATE_BLENDSHAPE_ATTR("browOuterUpRight", "bour", browOuterUpRight); CREATE_BLENDSHAPE_ATTR("cheekPuff", "cp", cheekPuff);
    CREATE_BLENDSHAPE_ATTR("cheekSquintLeft", "csl", cheekSquintLeft); CREATE_BLENDSHAPE_ATTR("cheekSquintRight", "csr", cheekSquintRight);
    CREATE_BLENDSHAPE_ATTR("eyeBlinkLeft", "ebl", eyeBlinkLeft); CREATE_BLENDSHAPE_ATTR("eyeBlinkRight", "ebr", eyeBlinkRight);
    CREATE_BLENDSHAPE_ATTR("eyeLookDownLeft", "eldl", eyeLookDownLeft); CREATE_BLENDSHAPE_ATTR("eyeLookDownRight", "eldr", eyeLookDownRight);
    CREATE_BLENDSHAPE_ATTR("eyeLookInLeft", "elil", eyeLookInLeft); CREATE_BLENDSHAPE_ATTR("eyeLookInRight", "elir", eyeLookInRight);
    CREATE_BLENDSHAPE_ATTR("eyeLookOutLeft", "elol", eyeLookOutLeft); CREATE_BLENDSHAPE_ATTR("eyeLookOutRight", "elor", eyeLookOutRight);
    CREATE_BLENDSHAPE_ATTR("eyeLookUpLeft", "elul", eyeLookUpLeft); CREATE_BLENDSHAPE_ATTR("eyeLookUpRight", "elur", eyeLookUpRight);
    CREATE_BLENDSHAPE_ATTR("eyeSquintLeft", "esl", eyeSquintLeft); CREATE_BLENDSHAPE_ATTR("eyeSquintRight", "esr", eyeSquintRight);
    CREATE_BLENDSHAPE_ATTR("eyeWideLeft", "ewl", eyeWideLeft); CREATE_BLENDSHAPE_ATTR("eyeWideRight", "ewr", eyeWideRight);
    CREATE_BLENDSHAPE_ATTR("jawForward", "jf", jawForward); CREATE_BLENDSHAPE_ATTR("jawLeft", "jl", jawLeft);
    CREATE_BLENDSHAPE_ATTR("jawOpen", "jo", jawOpen); CREATE_BLENDSHAPE_ATTR("jawRight", "jr", jawRight);
    CREATE_BLENDSHAPE_ATTR("mouthClose", "mc", mouthClose); CREATE_BLENDSHAPE_ATTR("mouthDimpleLeft", "mdl", mouthDimpleLeft);
    CREATE_BLENDSHAPE_ATTR("mouthDimpleRight", "mdr", mouthDimpleRight); CREATE_BLENDSHAPE_ATTR("mouthFrownLeft", "mfl", mouthFrownLeft);
    CREATE_BLENDSHAPE_ATTR("mouthFrownRight", "mfr", mouthFrownRight); CREATE_BLENDSHAPE_ATTR("mouthFunnel", "mf", mouthFunnel);
    CREATE_BLENDSHAPE_ATTR("mouthLowerDownLeft", "mldl", mouthLowerDownLeft); CREATE_BLENDSHAPE_ATTR("mouthLowerDownRight", "mldr", mouthLowerDownRight);
    CREATE_BLENDSHAPE_ATTR("mouthPressLeft", "mpl", mouthPressLeft); CREATE_BLENDSHAPE_ATTR("mouthPressRight", "mpr", mouthPressRight);
    CREATE_BLENDSHAPE_ATTR("mouthPucker", "mp", mouthPucker); CREATE_BLENDSHAPE_ATTR("mouthRight", "mr", mouthRight);
    CREATE_BLENDSHAPE_ATTR("mouthRollLower", "mrl", mouthRollLower); CREATE_BLENDSHAPE_ATTR("mouthRollUpper", "mru", mouthRollUpper);
    CREATE_BLENDSHAPE_ATTR("mouthShrugLower", "msl", mouthShrugLower); CREATE_BLENDSHAPE_ATTR("mouthShrugUpper", "msu", mouthShrugUpper);
    CREATE_BLENDSHAPE_ATTR("mouthSmileLeft", "msml", mouthSmileLeft); CREATE_BLENDSHAPE_ATTR("mouthSmileRight", "msmr", mouthSmileRight);
    CREATE_BLENDSHAPE_ATTR("mouthStretchLeft", "mstl", mouthStretchLeft); CREATE_BLENDSHAPE_ATTR("mouthStretchRight", "mstr", mouthStretchRight);
    CREATE_BLENDSHAPE_ATTR("mouthUpperUpLeft", "muul", mouthUpperUpLeft); CREATE_BLENDSHAPE_ATTR("mouthUpperUpRight", "muur", mouthUpperUpRight);
    CREATE_BLENDSHAPE_ATTR("noseSneerLeft", "nsl", noseSneerLeft); CREATE_BLENDSHAPE_ATTR("noseSneerRight", "nsr", noseSneerRight);

    outputBlendshapes = cAttr.create("outputBlendshapes", "obs"); cAttr.setArray(false);
    #define ADD_BLENDSHAPE_CHILD(attr) cAttr.addChild(attr)
    ADD_BLENDSHAPE_CHILD(browDownLeft); ADD_BLENDSHAPE_CHILD(browDownRight); ADD_BLENDSHAPE_CHILD(browInnerUp);
    ADD_BLENDSHAPE_CHILD(browOuterUpLeft); ADD_BLENDSHAPE_CHILD(browOuterUpRight); ADD_BLENDSHAPE_CHILD(cheekPuff);
    ADD_BLENDSHAPE_CHILD(cheekSquintLeft); ADD_BLENDSHAPE_CHILD(cheekSquintRight); ADD_BLENDSHAPE_CHILD(eyeBlinkLeft);
    ADD_BLENDSHAPE_CHILD(eyeBlinkRight); ADD_BLENDSHAPE_CHILD(eyeLookDownLeft); ADD_BLENDSHAPE_CHILD(eyeLookDownRight);
    ADD_BLENDSHAPE_CHILD(eyeLookInLeft); ADD_BLENDSHAPE_CHILD(eyeLookInRight); ADD_BLENDSHAPE_CHILD(eyeLookOutLeft);
    ADD_BLENDSHAPE_CHILD(eyeLookOutRight); ADD_BLENDSHAPE_CHILD(eyeLookUpLeft); ADD_BLENDSHAPE_CHILD(eyeLookUpRight);
    ADD_BLENDSHAPE_CHILD(eyeSquintLeft); ADD_BLENDSHAPE_CHILD(eyeSquintRight); ADD_BLENDSHAPE_CHILD(eyeWideLeft);
    ADD_BLENDSHAPE_CHILD(eyeWideRight); ADD_BLENDSHAPE_CHILD(jawForward); ADD_BLENDSHAPE_CHILD(jawLeft);
    ADD_BLENDSHAPE_CHILD(jawOpen); ADD_BLENDSHAPE_CHILD(jawRight); ADD_BLENDSHAPE_CHILD(mouthClose);
    ADD_BLENDSHAPE_CHILD(mouthDimpleLeft); ADD_BLENDSHAPE_CHILD(mouthDimpleRight); ADD_BLENDSHAPE_CHILD(mouthFrownLeft);
    ADD_BLENDSHAPE_CHILD(mouthFrownRight); ADD_BLENDSHAPE_CHILD(mouthFunnel); ADD_BLENDSHAPE_CHILD(mouthLowerDownLeft);
    ADD_BLENDSHAPE_CHILD(mouthLowerDownRight); ADD_BLENDSHAPE_CHILD(mouthPressLeft); ADD_BLENDSHAPE_CHILD(mouthPressRight);
    ADD_BLENDSHAPE_CHILD(mouthPucker); ADD_BLENDSHAPE_CHILD(mouthRight); ADD_BLENDSHAPE_CHILD(mouthRollLower);
    ADD_BLENDSHAPE_CHILD(mouthRollUpper); ADD_BLENDSHAPE_CHILD(mouthShrugLower); ADD_BLENDSHAPE_CHILD(mouthShrugUpper);
    ADD_BLENDSHAPE_CHILD(mouthSmileLeft); ADD_BLENDSHAPE_CHILD(mouthSmileRight); ADD_BLENDSHAPE_CHILD(mouthStretchLeft);
    ADD_BLENDSHAPE_CHILD(mouthStretchRight); ADD_BLENDSHAPE_CHILD(mouthUpperUpLeft); ADD_BLENDSHAPE_CHILD(mouthUpperUpRight);
    ADD_BLENDSHAPE_CHILD(noseSneerLeft); ADD_BLENDSHAPE_CHILD(noseSneerRight);
    addAttribute(outputBlendshapes);

    attributeAffects(videoPath, outputBlendshapes);
    attributeAffects(currentTime, outputBlendshapes);
    attributeAffects(enable, outputBlendshapes);

    return MS::kSuccess;
}

MStatus initializePlugin(MObject obj) {
    MFnPlugin plugin(obj, "Video2ARKit_Node", "1.0", "Any");
    return plugin.registerNode("Video2ARKitNode", Video2ARKitNode::id, Video2ARKitNode::creator, Video2ARKitNode::initialize);
}

MStatus uninitializePlugin(MObject obj) {
    MFnPlugin plugin(obj);
    return plugin.deregisterNode(Video2ARKitNode::id);
}