// ScreenPreview.cpp
#include "stdafx.h"
#include "ScreenPreview.h"
#include "../common/commands.h"   // ScreenPreviewStatus

#include <windows.h>
#include <objidl.h>     // IUnknown / IStream — gdiplus.h 依赖它们已声明
#include <gdiplus.h>
#pragma comment(lib, "gdiplus.lib")

using namespace Gdiplus;

namespace {

// GDI+ 进程级初始化（与 Bmp2Video 互不冲突；Startup 可重入计数）
struct GdiplusBoot {
    ULONG_PTR token = 0;
    bool ok = false;
    GdiplusBoot()
    {
        GdiplusStartupInput in;
        ok = (GdiplusStartup(&token, &in, NULL) == Ok);
    }
    ~GdiplusBoot()
    {
        if (ok) GdiplusShutdown(token);
    }
};
static GdiplusBoot g_boot;

int GetJpegEncoderClsid(CLSID& clsid)
{
    UINT num = 0, size = 0;
    GetImageEncodersSize(&num, &size);
    if (size == 0) return -1;

    std::vector<BYTE> buf(size);
    ImageCodecInfo* info = reinterpret_cast<ImageCodecInfo*>(buf.data());
    GetImageEncoders(num, size, info);
    for (UINT i = 0; i < num; ++i) {
        if (wcscmp(info[i].MimeType, L"image/jpeg") == 0) {
            clsid = info[i].Clsid;
            return 0;
        }
    }
    return -1;
}

// 抓主屏到 24bpp Bitmap，目标尺寸已等比换算。
// 返回新分配的 Bitmap*，失败返回 nullptr。调用者负责 delete。
Bitmap* GrabPrimaryScaled(int targetW, int targetH)
{
    HDC hScreen = GetDC(NULL);
    if (!hScreen) return nullptr;

    int srcX = GetSystemMetrics(SM_XVIRTUALSCREEN);   // 主屏左上 — 仅取主屏时用 0,0
    int srcY = GetSystemMetrics(SM_YVIRTUALSCREEN);
    (void)srcX; (void)srcY;

    int srcW = GetSystemMetrics(SM_CXSCREEN);
    int srcH = GetSystemMetrics(SM_CYSCREEN);
    if (srcW <= 0 || srcH <= 0) {
        ReleaseDC(NULL, hScreen);
        return nullptr;
    }

    HDC hMem = CreateCompatibleDC(hScreen);
    HBITMAP hBmp = CreateCompatibleBitmap(hScreen, targetW, targetH);
    if (!hMem || !hBmp) {
        if (hBmp) DeleteObject(hBmp);
        if (hMem) DeleteDC(hMem);
        ReleaseDC(NULL, hScreen);
        return nullptr;
    }

    HGDIOBJ oldBmp = SelectObject(hMem, hBmp);

    // 高质量缩放：HALFTONE 内插
    SetStretchBltMode(hMem, HALFTONE);
    SetBrushOrgEx(hMem, 0, 0, NULL);
    BOOL bb = StretchBlt(hMem, 0, 0, targetW, targetH,
                         hScreen, 0, 0, srcW, srcH, SRCCOPY | CAPTUREBLT);

    SelectObject(hMem, oldBmp);

    Bitmap* out = nullptr;
    if (bb) {
        // 拷贝 HBITMAP 到 GDI+ Bitmap，避免后续释放设备 DC 影响图像
        Bitmap tmp(hBmp, NULL);
        if (tmp.GetLastStatus() == Ok) {
            out = tmp.Clone(0, 0, targetW, targetH, PixelFormat24bppRGB);
            if (out && out->GetLastStatus() != Ok) {
                delete out;
                out = nullptr;
            }
        }
    }

    DeleteObject(hBmp);
    DeleteDC(hMem);
    ReleaseDC(NULL, hScreen);
    return out;
}

} // namespace

int CaptureAndEncodePreview(int maxWidth, int quality,
                            std::vector<unsigned char>& out,
                            int& outWidth, int& outHeight)
{
    out.clear();
    outWidth = outHeight = 0;

    if (!g_boot.ok) return SCREEN_PREVIEW_NOT_SUPPORTED;
    if (maxWidth < 64) maxWidth = 64;
    if (maxWidth > 1920) maxWidth = 1920;
    if (quality < 1) quality = 1;
    if (quality > 100) quality = 100;

    int srcW = GetSystemMetrics(SM_CXSCREEN);
    int srcH = GetSystemMetrics(SM_CYSCREEN);
    if (srcW <= 0 || srcH <= 0) return SCREEN_PREVIEW_CAPTURE_FAILED;

    // 等比缩放，禁止放大
    int targetW = (srcW <= maxWidth) ? srcW : maxWidth;
    int targetH = (int)((double)srcH * targetW / srcW + 0.5);
    if (targetH <= 0) targetH = 1;
    // 偶数对齐，JPEG 编码更高效
    targetW &= ~1;
    targetH &= ~1;
    if (targetW < 2) targetW = 2;
    if (targetH < 2) targetH = 2;

    Bitmap* bmp = GrabPrimaryScaled(targetW, targetH);
    if (!bmp) return SCREEN_PREVIEW_CAPTURE_FAILED;

    CLSID clsid;
    if (GetJpegEncoderClsid(clsid) != 0) {
        delete bmp;
        return SCREEN_PREVIEW_ENCODE_FAILED;
    }

    EncoderParameters params;
    params.Count = 1;
    params.Parameter[0].Guid = EncoderQuality;
    params.Parameter[0].Type = EncoderParameterValueTypeLong;
    params.Parameter[0].NumberOfValues = 1;
    ULONG q = (ULONG)quality;
    params.Parameter[0].Value = &q;

    IStream* stream = nullptr;
    if (FAILED(CreateStreamOnHGlobal(NULL, TRUE, &stream))) {
        delete bmp;
        return SCREEN_PREVIEW_ENCODE_FAILED;
    }

    Status st = bmp->Save(stream, &clsid, &params);
    delete bmp;

    if (st != Ok) {
        stream->Release();
        return SCREEN_PREVIEW_ENCODE_FAILED;
    }

    HGLOBAL hMem = NULL;
    if (FAILED(GetHGlobalFromStream(stream, &hMem)) || !hMem) {
        stream->Release();
        return SCREEN_PREVIEW_ENCODE_FAILED;
    }
    SIZE_T sz = GlobalSize(hMem);
    if (sz == 0) {
        stream->Release();
        return SCREEN_PREVIEW_ENCODE_FAILED;
    }

    void* p = GlobalLock(hMem);
    if (!p) {
        stream->Release();
        return SCREEN_PREVIEW_ENCODE_FAILED;
    }
    out.assign((unsigned char*)p, (unsigned char*)p + sz);
    GlobalUnlock(hMem);
    stream->Release();

    outWidth = targetW;
    outHeight = targetH;
    return SCREEN_PREVIEW_OK;
}