#include "AiFrameConverter.h"
#include <wincodec.h>
#include <algorithm>
#include <sstream>
#include <cstring>
#include <utility>
#include <oleauto.h>

#pragma comment(lib, "windowscodecs.lib")

static int AiAbsStride(int v) { return v < 0 ? -v : v; }

const char* AiPixelFormatName(HWS_PIXEL_FORMAT fmt)
{
    switch (fmt) {
    case HWS_PIXFMT_YUY2: return "YUY2";
    case HWS_PIXFMT_NV12: return "NV12";
    case HWS_PIXFMT_RGB24: return "RGB24";
    case HWS_PIXFMT_RGB32: return "RGB32";
    case HWS_PIXFMT_MJPG: return "MJPG";
    case HWS_PIXFMT_H264: return "H264";
    default: return "UNKNOWN";
    }
}

bool AiCopySdkFrame(const HWS_VIDEO_FRAME* src, int slot, AiCapturedFrame* outFrame)
{
    if (!src || !outFrame || !src->data[0] || src->width <= 0 || src->height <= 0) return false;

    AiCapturedFrame f;
    f.slot = slot;
    f.width = src->width;
    f.height = src->height;
    f.pixelFormat = src->pixelFormat;
    f.stride0 = src->stride[0];
    f.stride1 = src->stride[1];
    f.timestamp100ns = src->timestamp100ns;
    f.frameIndex = src->frameIndex;

    const int stride0 = AiAbsStride(src->stride[0]);
    const int stride1 = AiAbsStride(src->stride[1]);
    if (stride0 <= 0) return false;

    size_t bytes0 = 0;
    size_t bytes1 = 0;

    switch (src->pixelFormat) {
    case HWS_PIXFMT_MJPG:
        // The current SDK uses stride[0] as compressed byte size for MJPG/H264.
        bytes0 = (size_t)stride0;
        break;
    case HWS_PIXFMT_RGB32:
    case HWS_PIXFMT_RGB24:
    case HWS_PIXFMT_YUY2:
        bytes0 = (size_t)stride0 * (size_t)src->height;
        break;
    case HWS_PIXFMT_NV12:
        bytes0 = (size_t)stride0 * (size_t)src->height;
        bytes1 = (size_t)(stride1 > 0 ? stride1 : stride0) * (size_t)((src->height + 1) / 2);
        break;
    default:
        return false;
    }

    if (bytes0 == 0) return false;
    f.plane0.resize(bytes0);
    memcpy(f.plane0.data(), src->data[0], bytes0);

    if (src->pixelFormat == HWS_PIXFMT_NV12) {
        f.stride1 = (stride1 > 0 ? stride1 : stride0);
        f.plane1.resize(bytes1);
        const unsigned char* uv = src->data[1];
        if (!uv) uv = src->data[0] + bytes0;
        memcpy(f.plane1.data(), uv, bytes1);
    }

    *outFrame = std::move(f);
    return true;
}

static unsigned char AiClipByte(int v)
{
    if (v < 0) return 0;
    if (v > 255) return 255;
    return (unsigned char)v;
}

static void AiYuvToBgra(unsigned char y, unsigned char u, unsigned char v, unsigned char* bgra)
{
    int c = (int)y - 16;
    int d = (int)u - 128;
    int e = (int)v - 128;
    if (c < 0) c = 0;
    int r = (298 * c + 409 * e + 128) >> 8;
    int g = (298 * c - 100 * d - 208 * e + 128) >> 8;
    int b = (298 * c + 516 * d + 128) >> 8;
    bgra[0] = AiClipByte(b);
    bgra[1] = AiClipByte(g);
    bgra[2] = AiClipByte(r);
    bgra[3] = 255;
}

static bool AiConvertToBgra(const AiCapturedFrame& frame, std::vector<unsigned char>* outBgra)
{
    if (!outBgra || frame.width <= 0 || frame.height <= 0 || frame.plane0.empty()) return false;
    const int w = frame.width;
    const int h = frame.height;
    outBgra->assign((size_t)w * (size_t)h * 4, 0);
    unsigned char* dstBase = outBgra->data();

    switch (frame.pixelFormat) {
    case HWS_PIXFMT_RGB32:
    {
        const int srcStride = AiAbsStride(frame.stride0);
        if (srcStride < w * 4) return false;
        for (int y = 0; y < h; ++y) {
            const unsigned char* src = frame.plane0.data() + (size_t)y * srcStride;
            unsigned char* dst = dstBase + (size_t)y * w * 4;
            memcpy(dst, src, (size_t)w * 4);
            for (int x = 0; x < w; ++x) dst[x * 4 + 3] = 255;
        }
        return true;
    }
    case HWS_PIXFMT_RGB24:
    {
        const int srcStride = AiAbsStride(frame.stride0);
        if (srcStride < w * 3) return false;
        for (int y = 0; y < h; ++y) {
            const unsigned char* src = frame.plane0.data() + (size_t)y * srcStride;
            unsigned char* dst = dstBase + (size_t)y * w * 4;
            for (int x = 0; x < w; ++x) {
                // DirectShow RGB24 is usually BGR byte order.
                dst[x * 4 + 0] = src[x * 3 + 0];
                dst[x * 4 + 1] = src[x * 3 + 1];
                dst[x * 4 + 2] = src[x * 3 + 2];
                dst[x * 4 + 3] = 255;
            }
        }
        return true;
    }
    case HWS_PIXFMT_YUY2:
    {
        const int srcStride = AiAbsStride(frame.stride0);
        if (srcStride < w * 2) return false;
        for (int y = 0; y < h; ++y) {
            const unsigned char* src = frame.plane0.data() + (size_t)y * srcStride;
            unsigned char* dst = dstBase + (size_t)y * w * 4;
            for (int x = 0; x < w; x += 2) {
                unsigned char y0 = src[x * 2 + 0];
                unsigned char u0 = src[x * 2 + 1];
                unsigned char y1 = src[x * 2 + 2];
                unsigned char v0 = src[x * 2 + 3];
                AiYuvToBgra(y0, u0, v0, dst + (size_t)x * 4);
                if (x + 1 < w) AiYuvToBgra(y1, u0, v0, dst + (size_t)(x + 1) * 4);
            }
        }
        return true;
    }
    case HWS_PIXFMT_NV12:
    {
        const int yStride = AiAbsStride(frame.stride0);
        const int uvStride = AiAbsStride(frame.stride1 > 0 ? frame.stride1 : frame.stride0);
        if (yStride < w || uvStride < w || frame.plane1.empty()) return false;
        for (int y = 0; y < h; ++y) {
            const unsigned char* yRow = frame.plane0.data() + (size_t)y * yStride;
            const unsigned char* uvRow = frame.plane1.data() + (size_t)(y / 2) * uvStride;
            unsigned char* dst = dstBase + (size_t)y * w * 4;
            for (int x = 0; x < w; ++x) {
                const int uvx = (x / 2) * 2;
                unsigned char u = uvRow[uvx + 0];
                unsigned char v = uvRow[uvx + 1];
                AiYuvToBgra(yRow[x], u, v, dst + (size_t)x * 4);
            }
        }
        return true;
    }
    default:
        return false;
    }
}

static bool AiCreateWicFactory(IWICImagingFactory** factory)
{
    if (!factory) return false;
    *factory = NULL;
    HRESULT hr = CoCreateInstance(CLSID_WICImagingFactory2, NULL, CLSCTX_INPROC_SERVER,
                                  IID_PPV_ARGS(factory));
    if (FAILED(hr)) {
        hr = CoCreateInstance(CLSID_WICImagingFactory, NULL, CLSCTX_INPROC_SERVER,
                              IID_PPV_ARGS(factory));
    }
    return SUCCEEDED(hr) && *factory;
}

bool AiEncodeFrameToJpeg(const AiCapturedFrame& frame, int jpegQuality, std::vector<unsigned char>* outJpeg, std::wstring* errorText)
{
    if (!outJpeg) return false;
    outJpeg->clear();

    if (frame.pixelFormat == HWS_PIXFMT_MJPG && !frame.plane0.empty()) {
        // MJPG is already JPEG-compressed. Forward it directly to the AI service.
        *outJpeg = frame.plane0;
        return true;
    }

    std::vector<unsigned char> bgra;
    if (!AiConvertToBgra(frame, &bgra)) {
        if (errorText) *errorText = L"AI frame conversion failed.";
        return false;
    }

    IWICImagingFactory* factory = NULL;
    IWICBitmapEncoder* encoder = NULL;
    IWICBitmapFrameEncode* frameEncode = NULL;
    IPropertyBag2* propertyBag = NULL;
    IWICBitmap* bitmap = NULL;
    IWICFormatConverter* converter = NULL;
    IStream* stream = NULL;
    HGLOBAL hGlobal = NULL;
    bool ok = false;

    do {
        if (!AiCreateWicFactory(&factory)) { if (errorText) *errorText = L"Create WIC factory failed."; break; }
        if (FAILED(CreateStreamOnHGlobal(NULL, TRUE, &stream)) || !stream) { if (errorText) *errorText = L"Create stream failed."; break; }
        if (FAILED(factory->CreateEncoder(GUID_ContainerFormatJpeg, NULL, &encoder))) { if (errorText) *errorText = L"Create JPEG encoder failed."; break; }
        if (FAILED(encoder->Initialize(stream, WICBitmapEncoderNoCache))) { if (errorText) *errorText = L"Initialize JPEG encoder failed."; break; }
        if (FAILED(encoder->CreateNewFrame(&frameEncode, &propertyBag))) { if (errorText) *errorText = L"Create JPEG frame failed."; break; }

        if (propertyBag) {
            PROPBAG2 option = {};
            option.pstrName = const_cast<LPOLESTR>(L"ImageQuality");
            VARIANT value;
            VariantInit(&value);
            value.vt = VT_R4;
            int q = jpegQuality;
            if (q < 1) q = 1;
            if (q > 100) q = 100;
            value.fltVal = (float)q / 100.0f;
            propertyBag->Write(1, &option, &value);
            VariantClear(&value);
        }

        if (FAILED(frameEncode->Initialize(propertyBag))) { if (errorText) *errorText = L"Initialize JPEG frame failed."; break; }
        if (FAILED(frameEncode->SetSize((UINT)frame.width, (UINT)frame.height))) { if (errorText) *errorText = L"Set JPEG size failed."; break; }
        const UINT stride = (UINT)frame.width * 4;
        const UINT imageSize = stride * (UINT)frame.height;
        if (FAILED(factory->CreateBitmapFromMemory((UINT)frame.width, (UINT)frame.height,
                                                   GUID_WICPixelFormat32bppBGRA,
                                                   stride, imageSize, bgra.data(), &bitmap))) {
            if (errorText) *errorText = L"Create WIC bitmap failed.";
            break;
        }
        if (FAILED(factory->CreateFormatConverter(&converter))) { if (errorText) *errorText = L"Create WIC converter failed."; break; }
        if (FAILED(converter->Initialize(bitmap, GUID_WICPixelFormat24bppBGR,
                                         WICBitmapDitherTypeNone, NULL, 0.0,
                                         WICBitmapPaletteTypeCustom))) {
            if (errorText) *errorText = L"Initialize WIC converter failed.";
            break;
        }
        WICPixelFormatGUID pixelFormat = GUID_WICPixelFormat24bppBGR;
        if (FAILED(frameEncode->SetPixelFormat(&pixelFormat))) { if (errorText) *errorText = L"Set JPEG pixel format failed."; break; }
        if (FAILED(frameEncode->WriteSource(converter, NULL))) { if (errorText) *errorText = L"Write JPEG source failed."; break; }
        if (FAILED(frameEncode->Commit())) { if (errorText) *errorText = L"Commit JPEG frame failed."; break; }
        if (FAILED(encoder->Commit())) { if (errorText) *errorText = L"Commit JPEG encoder failed."; break; }

        if (FAILED(GetHGlobalFromStream(stream, &hGlobal)) || !hGlobal) { if (errorText) *errorText = L"Get JPEG memory failed."; break; }
        SIZE_T size = GlobalSize(hGlobal);
        void* data = GlobalLock(hGlobal);
        if (!data || size == 0) { if (errorText) *errorText = L"Lock JPEG memory failed."; break; }
        outJpeg->assign((unsigned char*)data, (unsigned char*)data + size);
        GlobalUnlock(hGlobal);
        ok = true;
    } while (false);

    if (converter) converter->Release();
    if (bitmap) bitmap->Release();
    if (propertyBag) propertyBag->Release();
    if (frameEncode) frameEncode->Release();
    if (encoder) encoder->Release();
    if (stream) stream->Release();
    if (factory) factory->Release();
    return ok;
}
