// MediaImagePlane.cpp
// Implementation of MediaImagePlane - Custom MPxImagePlane with FFmpeg video decoding

#include "MediaImagePlane.h"
#include "FFmpegVideoDecoder.h"
#include "FrameCache.h"

#include <maya/MImage.h>

#include <cstring>

// Static member initialization
const MString MediaImagePlane::kNodeName = "MediaImagePlane";
const MTypeId MediaImagePlane::kNodeId = 0x0013A5F1;

// Attribute objects
MObject MediaImagePlane::aVideoFile;
MObject MediaImagePlane::aPlaybackRate;
MObject MediaImagePlane::aUseMayaFrameRate;
MObject MediaImagePlane::aLoop;
MObject MediaImagePlane::aPostEffectCrop;
MObject MediaImagePlane::aPostEffectResize;
MObject MediaImagePlane::aPostEffectFlip;
MObject MediaImagePlane::aCacheSize;
MObject MediaImagePlane::aClearCache;

// Output attributes
MObject MediaImagePlane::aOutFrameWidth;
MObject MediaImagePlane::aOutFrameHeight;
MObject MediaImagePlane::aOutFrameCount;
MObject MediaImagePlane::aOutCacheHitRatio;

// ============================================================================
// Constructor
// ============================================================================
MediaImagePlane::MediaImagePlane()
    : MPxImagePlane()
{
    // Initialize FFmpeg decoder and frame cache
    m_decoder = std::make_unique<MediaPlane::FFmpegVideoDecoder>();
    m_frameCache = std::make_unique<MediaPlane::FrameCache>(m_maxCacheMemory, m_maxCacheFrames);
    
    m_videoWidth = 0;
    m_videoHeight = 0;
    m_videoFrameRate = 0.0;
    m_videoFrameCount = 0;
    
    m_lastFrameIndex = -1;
    m_imageDirty = true;
    
    // Register time change callback for playback updates
    m_timeChangedCallbackId = MEventMessage::addEventCallback(
        "timeChanged", timeChangedCallback, this);
}

// ============================================================================
// Destructor
// ============================================================================
MediaImagePlane::~MediaImagePlane()
{
    // Remove time change callback
    if (m_timeChangedCallbackId != 0) {
        MMessage::removeCallback(m_timeChangedCallbackId);
        m_timeChangedCallbackId = 0;
    }
    
    // Close video file
    closeVideoFile();
}

// ============================================================================
// Creator function
// ============================================================================
void* MediaImagePlane::creator()
{
    return new MediaImagePlane();
}

// ============================================================================
// Initialize node attributes
// ============================================================================
MStatus MediaImagePlane::initialize()
{
    MStatus status;
    
    MFnNumericAttribute numAttr;
    MFnTypedAttribute typedAttr;

    // Input: Video File Path
    aVideoFile = typedAttr.create("videoFile", "vf", MFnData::kString, MObject::kNullObj, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    typedAttr.setUsedAsFilename(true);
    typedAttr.setStorable(true);
    addAttribute(aVideoFile);

    // Input: Playback Rate (0.25 to 4.0)
    aPlaybackRate = numAttr.create("playbackRate", "pr", MFnNumericData::kDouble, 1.0, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    numAttr.setStorable(true);
    numAttr.setKeyable(true);
    numAttr.setMin(0.25);
    numAttr.setMax(4.0);
    addAttribute(aPlaybackRate);

    // Input: Use Maya Frame Rate
    aUseMayaFrameRate = numAttr.create("useMayaFrameRate", "umf", MFnNumericData::kBoolean, true, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    numAttr.setStorable(true);
    numAttr.setKeyable(true);
    addAttribute(aUseMayaFrameRate);

    // Input: Loop Playback
    aLoop = numAttr.create("loop", "lp", MFnNumericData::kBoolean, true, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    numAttr.setStorable(true);
    numAttr.setKeyable(true);
    addAttribute(aLoop);

    // Input: Post-effect Crop (x, y, width, height)
    aPostEffectCrop = numAttr.create("postEffectCrop", "pec", MFnNumericData::k4Double, 0.0, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    numAttr.setStorable(true);
    numAttr.setKeyable(false);
    addAttribute(aPostEffectCrop);

    // Input: Post-effect Resize (width, height)
    aPostEffectResize = numAttr.create("postEffectResize", "per", MFnNumericData::k2Double, 0.0, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    numAttr.setStorable(true);
    numAttr.setKeyable(false);
    addAttribute(aPostEffectResize);

    // Input: Post-effect Flip (horizontal, vertical as bitmask)
    aPostEffectFlip = numAttr.create("postEffectFlip", "pef", MFnNumericData::k2Long, 0, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    numAttr.setStorable(true);
    numAttr.setKeyable(false);
    addAttribute(aPostEffectFlip);

    // Input: Cache Size (in MB)
    aCacheSize = numAttr.create("cacheSize", "cs", MFnNumericData::kInt, 256, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    numAttr.setStorable(true);
    numAttr.setKeyable(false);
    numAttr.setMin(16);
    numAttr.setMax(2048);
    addAttribute(aCacheSize);

    // Input: Clear Cache (trigger)
    aClearCache = numAttr.create("clearCache", "cc", MFnNumericData::kBoolean, false, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    numAttr.setStorable(false);
    numAttr.setKeyable(false);
    addAttribute(aClearCache);

    // Output: Frame Width
    aOutFrameWidth = numAttr.create("outFrameWidth", "ofw", MFnNumericData::kInt, 0, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    numAttr.setStorable(false);
    numAttr.setReadable(true);
    addAttribute(aOutFrameWidth);

    // Output: Frame Height
    aOutFrameHeight = numAttr.create("outFrameHeight", "ofh", MFnNumericData::kInt, 0, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    numAttr.setStorable(false);
    numAttr.setReadable(true);
    addAttribute(aOutFrameHeight);

    // Output: Frame Count
    aOutFrameCount = numAttr.create("outFrameCount", "ofc", MFnNumericData::kInt64, 0, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    numAttr.setStorable(false);
    numAttr.setReadable(true);
    addAttribute(aOutFrameCount);

    // Output: Cache Hit Ratio
    aOutCacheHitRatio = numAttr.create("outCacheHitRatio", "och", MFnNumericData::kDouble, 0.0, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    numAttr.setStorable(false);
    numAttr.setReadable(true);
    addAttribute(aOutCacheHitRatio);

    // Set up attribute affects
    attributeAffects(aVideoFile, aOutFrameWidth);
    attributeAffects(aVideoFile, aOutFrameHeight);
    attributeAffects(aVideoFile, aOutFrameCount);
    attributeAffects(aVideoFile, aOutCacheHitRatio);
    
    attributeAffects(aPlaybackRate, aOutCacheHitRatio);
    attributeAffects(aUseMayaFrameRate, aOutCacheHitRatio);
    attributeAffects(aLoop, aOutCacheHitRatio);
    attributeAffects(aCacheSize, aOutCacheHitRatio);
    attributeAffects(aClearCache, aOutCacheHitRatio);
    attributeAffects(aPostEffectCrop, aOutCacheHitRatio);
    attributeAffects(aPostEffectResize, aOutCacheHitRatio);
    attributeAffects(aPostEffectFlip, aOutCacheHitRatio);

    return MStatus::kSuccess;
}

// ============================================================================
// Load image map - called by Maya to get the image for a specific frame
// ============================================================================
MStatus MediaImagePlane::loadImageMap(const MString& fileName, int frame, MImage& image)
{
    MStatus status;
    
    // Get current time from Maya
    MTime currentTime = MAnimControl::currentTime();
    double currentFrame = currentTime.asUnits(MTime::kFilm);  // Film = 24 fps equivalent
    
    // Calculate target frame index
    int64_t targetFrame = getCurrentFrameIndex();
    
    // Check if we need to reopen video file (if videoFile attribute changed)
    MObject thisObj = thisMObject();
    MPlug videoFilePlug(thisObj, aVideoFile);
    MString videoFilePath;
    videoFilePlug.getValue(videoFilePath);
    
    if (videoFilePath.length() > 0) {
        std::string path = videoFilePath.asUTF8();
        if (path != m_currentVideoFile.asUTF8() || !m_decoder->isOpen()) {
            if (!openVideoFile(path)) {
                MGlobal::displayError("Failed to open video file: " + videoFilePath);
                return MStatus::kFailure;
            }
        }
    }
    
    // Check if frame is valid
    if (!m_decoder->isOpen() || targetFrame < 0 || targetFrame >= m_videoFrameCount) {
        return MStatus::kFailure;
    }
    
    // Decode the frame
    return decodeFrame(targetFrame, image);
}

// ============================================================================
// Get internal value - retrieve custom attribute values
// ============================================================================
bool MediaImagePlane::getInternalValue(const MPlug& plug, MDataHandle& handle)
{
    if (plug == aVideoFile) {
        handle.set(m_currentVideoFile);
        return true;
    }
    if (plug == aPlaybackRate) {
        handle.set(m_playbackRate);
        return true;
    }
    if (plug == aUseMayaFrameRate) {
        handle.set(m_useMayaFrameRate);
        return true;
    }
    if (plug == aLoop) {
        handle.set(m_loop);
        return true;
    }
    if (plug == aPostEffectCrop) {
        double cropData[4] = { m_cropX, m_cropY, m_cropW, m_cropH };
        handle.set(cropData);
        return true;
    }
    if (plug == aPostEffectResize) {
        double resizeData[2] = { m_resizeW, m_resizeH };
        handle.set(resizeData);
        return true;
    }
    if (plug == aPostEffectFlip) {
        int flipInt = (m_flipH ? 1 : 0) | (m_flipV ? 2 : 0);
        handle.set(flipInt);
        return true;
    }
    if (plug == aCacheSize) {
        handle.set(static_cast<int>(m_maxCacheMemory / (1024 * 1024)));
        return true;
    }
    
    return MPxImagePlane::getInternalValue(plug, handle);
}

// ============================================================================
// Set internal value - handle custom attribute changes
// ============================================================================
bool MediaImagePlane::setInternalValue(const MPlug& plug, const MDataHandle& handle)
{
    if (plug == aVideoFile) {
        MString newFile = handle.asString();
        if (newFile != m_currentVideoFile) {
            m_currentVideoFile = newFile;
            m_imageDirty = true;
            setImageDirty();
            
            // Open the new video file
            if (newFile.length() > 0) {
                std::string path = newFile.asUTF8();
                openVideoFile(path);
            }
        }
        return true;
    }
    if (plug == aPlaybackRate) {
        m_playbackRate = handle.asDouble();
        if (m_playbackRate < 0.25) m_playbackRate = 0.25;
        if (m_playbackRate > 4.0) m_playbackRate = 4.0;
        m_imageDirty = true;
        return true;
    }
    if (plug == aUseMayaFrameRate) {
        m_useMayaFrameRate = handle.asBool();
        m_imageDirty = true;
        return true;
    }
    if (plug == aLoop) {
        m_loop = handle.asBool();
        return true;
    }
    if (plug == aPostEffectCrop) {
        const double* cropData = handle.asDouble4();
        m_cropX = cropData[0];
        m_cropY = cropData[1];
        m_cropW = cropData[2];
        m_cropH = cropData[3];
        m_imageDirty = true;
        return true;
    }
    if (plug == aPostEffectResize) {
        const double* resizeData = handle.asDouble2();
        m_resizeW = resizeData[0];
        m_resizeH = resizeData[1];
        m_imageDirty = true;
        return true;
    }
    if (plug == aPostEffectFlip) {
        int flipInt = handle.asInt();
        m_flipH = (flipInt & 1) != 0;
        m_flipV = (flipInt & 2) != 0;
        m_imageDirty = true;
        return true;
    }
    if (plug == aCacheSize) {
        int cacheSizeMB = handle.asInt();
        m_maxCacheMemory = cacheSizeMB * 1024 * 1024;
        m_frameCache->setMaxMemorySize(m_maxCacheMemory);
        return true;
    }
    if (plug == aClearCache) {
        bool clearCache = handle.asBool();
        if (clearCache) {
            m_frameCache->clear();
            m_imageDirty = true;
        }
        return true;
    }
    
    return MPxImagePlane::setInternalValue(plug, handle);
}

// ============================================================================
// Open video file
// ============================================================================
bool MediaImagePlane::openVideoFile(const std::string& filePath)
{
    std::lock_guard<std::mutex> lock(m_decoderMutex);
    
    // Close existing file
    closeVideoFile();
    
    m_currentVideoFile = filePath.c_str();
    
    if (!m_decoder->open(filePath)) {
        MGlobal::displayError(MString("Failed to open video: ") + m_decoder->getLastError().c_str());
        return false;
    }
    
    // Get video info
    const auto& videoInfo = m_decoder->getVideoInfo();
    m_videoWidth = videoInfo.width;
    m_videoHeight = videoInfo.height;
    m_videoFrameRate = videoInfo.frameRate;
    m_videoFrameCount = videoInfo.frameCount;
    
    // Clear cache when video changes
    m_frameCache->invalidate();
    
    m_imageDirty = true;
    
    MGlobal::displayInfo("MediaImagePlane: Opened video - " + 
        MString(std::to_string(m_videoWidth).c_str()) + "x" + 
        MString(std::to_string(m_videoHeight).c_str()) + " @" + 
        MString(std::to_string(m_videoFrameRate).c_str()) + " fps");
    
    return true;
}

// ============================================================================
// Close video file
// ============================================================================
void MediaImagePlane::closeVideoFile()
{
    std::lock_guard<std::mutex> lock(m_decoderMutex);
    
    if (m_decoder && m_decoder->isOpen()) {
        m_decoder->close();
    }
    
    m_currentVideoFile = "";
    m_videoWidth = 0;
    m_videoHeight = 0;
    m_videoFrameRate = 0.0;
    m_videoFrameCount = 0;
    m_lastFrameIndex = -1;
}

// ============================================================================
// Get current frame index based on Maya timeline
// ============================================================================
int64_t MediaImagePlane::getCurrentFrameIndex() const
{
    // Get current time from Maya
    MTime currentTime = MAnimControl::currentTime();
    double currentFrameValue = currentTime.asUnits(MTime::kFilm);
    
    // Determine effective frame rate
    double effectiveFrameRate = m_videoFrameRate;
    
    if (m_useMayaFrameRate) {
        // Get Maya's current frame rate
        MTime::Unit timeUnit = currentTime.uiUnit();
        // Convert to fps (approximate)
        switch (timeUnit) {
            case MTime::kHours: effectiveFrameRate = 3600.0; break;
            case MTime::kMinutes: effectiveFrameRate = 60.0; break;
            case MTime::kSeconds: effectiveFrameRate = 1.0; break;
            case MTime::kMilliseconds: effectiveFrameRate = 1000.0; break;
            case MTime::kGames: effectiveFrameRate = 15.0; break;
            case MTime::kFilm: effectiveFrameRate = 24.0; break;
            case MTime::kNTSCFrame: effectiveFrameRate = 30.0; break;
            case MTime::kNTSCField: effectiveFrameRate = 60.0; break;
            case MTime::kPALFrame: effectiveFrameRate = 25.0; break;
            case MTime::kPALField: effectiveFrameRate = 50.0; break;
            case MTime::kShowScan: effectiveFrameRate = 48.0; break;
            default: effectiveFrameRate = 24.0; break;
        }
    }
    
    if (effectiveFrameRate <= 0) {
        effectiveFrameRate = 24.0;
    }
    
    // Calculate frame index
    double effectiveRate = effectiveFrameRate * m_playbackRate;
    int64_t frameIndex = static_cast<int64_t>(currentFrameValue * effectiveRate);
    
    // Apply looping
    if (m_loop && m_videoFrameCount > 0) {
        frameIndex = frameIndex % m_videoFrameCount;
    }
    
    // Clamp to valid range
    if (frameIndex < 0) frameIndex = 0;
    if (!m_loop && m_videoFrameCount > 0 && frameIndex >= m_videoFrameCount) {
        frameIndex = m_videoFrameCount - 1;
    }
    
    return frameIndex;
}

// ============================================================================
// Decode frame and populate MImage
// ============================================================================
MStatus MediaImagePlane::decodeFrame(int64_t frameIndex, MImage& image)
{
    if (!m_decoder->isOpen()) {
        return MStatus::kFailure;
    }
    
    // Check if frame changed
    if (frameIndex != m_lastFrameIndex) {
        m_lastFrameIndex = frameIndex;
        m_imageDirty = true;
    }
    
    // Get frame from decoder
    auto frameData = m_decoder->getFrame(frameIndex);
    
    if (!frameData.data) {
        return MStatus::kFailure;
    }
    
    // Get image dimensions
    unsigned int width = static_cast<unsigned int>(frameData.width);
    unsigned int height = static_cast<unsigned int>(frameData.height);
    
    // Create MImage with the frame data
    // MImage expects RGBA format
    image.create(width, height, 4, MImage::kByte);
    
    // Copy pixel data to MImage
    unsigned char* pixels = image.pixels();
    if (pixels && frameData.data) {
        // FFmpeg decoder outputs RGBA, copy directly
        memcpy(pixels, frameData.data, width * height * 4);
    }
    
    // Apply post-effects if needed
    applyPostEffects(image);
    
    return MStatus::kSuccess;
}

// ============================================================================
// Apply post-effects to the image
// ============================================================================
void MediaImagePlane::applyPostEffects(MImage& image)
{
    unsigned int width, height;
    image.getSize(width, height);
    
    // Apply flip if needed
    if (m_flipH || m_flipV) {
        // Create temporary buffer
        std::vector<unsigned char> tempPixels(width * height * 4);
        unsigned char* pixels = image.pixels();
        
        for (unsigned int y = 0; y < height; y++) {
            for (unsigned int x = 0; x < width; x++) {
                unsigned int srcY = m_flipV ? (height - 1 - y) : y;
                unsigned int srcX = m_flipH ? (width - 1 - x) : x;
                
                unsigned int srcIdx = (srcY * width + srcX) * 4;
                unsigned int dstIdx = (y * width + x) * 4;
                
                tempPixels[dstIdx + 0] = pixels[srcIdx + 0];
                tempPixels[dstIdx + 1] = pixels[srcIdx + 1];
                tempPixels[dstIdx + 2] = pixels[srcIdx + 2];
                tempPixels[dstIdx + 3] = pixels[srcIdx + 3];
            }
        }
        
        memcpy(pixels, tempPixels.data(), width * height * 4);
    }
    
    // Note: Crop and resize would require more complex implementation
    // For now, Maya's image plane handles these through its native attributes
}

// ============================================================================
// Set image dirty flag
// ============================================================================
void MediaImagePlane::setImageDirty()
{
    // This is handled by Maya's internal mechanism for image planes
    // Setting the image dirty forces Maya to call loadImageMap again
    m_imageDirty = true;
}

// ============================================================================
// Time changed callback
// ============================================================================
void MediaImagePlane::timeChangedCallback(void* clientData)
{
    MediaImagePlane* node = static_cast<MediaImagePlane*>(clientData);
    if (node) {
        node->setImageDirty();
    }
}

// Note: Plugin initialization and uninitialization are handled in Plugin.cpp
// This file only contains the node implementation