#include "capture.h"
#include "d3d_device.h"
#include "png_encoder.h"
#include "common.h"

#include <dwmapi.h>
#include <dxgi.h>

// DwmGetDxSharedSurface — undocumented but battle-tested API used by OBS and AHK.
// Returns a HANDLE to the DWM shared surface for a given HWND.
// No yellow border, no WM_PRINT, reads directly from DWM composition.
//
// Signature from reverse engineering / OBS source:
//   BOOL DwmGetDxSharedSurface(HWND hwnd, HANDLE* phSurface, LUID* pAdapterLuid,
//                               ULONG* pFmtWindow, ULONG* pPresentFlags, ULONGLONG* pWin32kUpdateId)
//
// The function is exported from user32.dll (not dwmapi.dll).
typedef BOOL(WINAPI* PFN_DwmGetDxSharedSurface)(
    HWND hwnd,
    HANDLE* phSurface,
    LUID* pAdapterLuid,
    ULONG* pFmtWindow,
    ULONG* pPresentFlags,
    ULONGLONG* pWin32kUpdateId);

static PFN_DwmGetDxSharedSurface g_DwmGetDxSharedSurface = nullptr;
static bool g_lookupDone = false;

static bool EnsureDwmFunction() {
    if (g_lookupDone) return g_DwmGetDxSharedSurface != nullptr;
    g_lookupDone = true;

    HMODULE hUser32 = GetModuleHandleW(L"user32.dll");
    if (!hUser32) {
        DWM_LOG("ERROR: user32.dll not loaded");
        return false;
    }

    g_DwmGetDxSharedSurface = reinterpret_cast<PFN_DwmGetDxSharedSurface>(
        GetProcAddress(hUser32, "DwmGetDxSharedSurface"));

    if (!g_DwmGetDxSharedSurface) {
        DWM_LOG("ERROR: DwmGetDxSharedSurface not found in user32.dll");
        return false;
    }

    DWM_LOG("DwmGetDxSharedSurface resolved from user32.dll");
    return true;
}

// --------------------------------------------------------------------------
// Constructor / Promise plumbing
// --------------------------------------------------------------------------

CaptureWorker::CaptureWorker(Napi::Env env, HWND hwnd,
                             Microsoft::WRL::ComPtr<ID3D11Device> device,
                             Microsoft::WRL::ComPtr<ID3D11DeviceContext> context)
    : Napi::AsyncWorker(env)
    , hwnd_(hwnd)
    , device_(device)
    , context_(context)
    , deferred_(Napi::Promise::Deferred::New(env))
{
}

Napi::Promise CaptureWorker::GetPromise() {
    return deferred_.Promise();
}

// --------------------------------------------------------------------------
// Execute — runs on libuv worker thread (no NAPI types allowed here)
// --------------------------------------------------------------------------

// Inner implementation with C++ destructors (ComPtr). Separated from SEH wrapper.
static bool ExecuteDwmCaptureInner(
    HWND hwnd,
    ID3D11Device* device,
    ID3D11DeviceContext* context,
    std::vector<uint8_t>& pngDataOut,
    std::string& errorOut)
{
    // ------------------------------------------------------------------
    // Step 1: Get the DWM shared surface handle for this window
    // ------------------------------------------------------------------
    HANDLE sharedHandle = nullptr;
    LUID adapterLuid = {};
    ULONG fmtWindow = 0;
    ULONG presentFlags = 0;
    ULONGLONG win32kUpdateId = 0;

    BOOL ok = g_DwmGetDxSharedSurface(
        hwnd, &sharedHandle, &adapterLuid, &fmtWindow, &presentFlags, &win32kUpdateId);

    if (!ok || !sharedHandle) {
        errorOut = "DwmGetDxSharedSurface failed — window may not have a DWM surface";
        return false;
    }

    DWM_LOG("Shared surface handle: %p, format: %lu", sharedHandle, fmtWindow);

    // ------------------------------------------------------------------
    // Step 2: Open the shared surface as a D3D11 texture
    // ------------------------------------------------------------------
    Microsoft::WRL::ComPtr<ID3D11Texture2D> sharedTexture;
    HRESULT hr = device->OpenSharedResource(
        sharedHandle,
        IID_PPV_ARGS(&sharedTexture));

    if (FAILED(hr) || !sharedTexture) {
        errorOut = "Failed to open DWM shared surface as D3D11 texture";
        return false;
    }

    D3D11_TEXTURE2D_DESC desc;
    sharedTexture->GetDesc(&desc);
    DWM_LOG("Shared texture: %ux%u, format=%u", desc.Width, desc.Height, desc.Format);

    if (desc.Width == 0 || desc.Height == 0) {
        errorOut = "DWM shared surface has zero dimensions";
        return false;
    }

    // ------------------------------------------------------------------
    // Step 3: Copy to staging texture for CPU read
    // ------------------------------------------------------------------
    D3D11_TEXTURE2D_DESC stagingDesc = {};
    stagingDesc.Width = desc.Width;
    stagingDesc.Height = desc.Height;
    stagingDesc.MipLevels = 1;
    stagingDesc.ArraySize = 1;
    stagingDesc.Format = desc.Format;
    stagingDesc.SampleDesc.Count = 1;
    stagingDesc.Usage = D3D11_USAGE_STAGING;
    stagingDesc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;

    Microsoft::WRL::ComPtr<ID3D11Texture2D> staging;
    hr = device->CreateTexture2D(&stagingDesc, nullptr, &staging);
    if (FAILED(hr)) {
        errorOut = "Failed to create staging texture";
        return false;
    }

    context->CopyResource(staging.Get(), sharedTexture.Get());

    D3D11_MAPPED_SUBRESOURCE mapped;
    hr = context->Map(staging.Get(), 0, D3D11_MAP_READ, 0, &mapped);
    if (FAILED(hr)) {
        errorOut = "Failed to map staging texture for CPU read";
        return false;
    }

    // ------------------------------------------------------------------
    // Step 4: Encode BGRA pixel data to PNG (per D-13, D-14)
    // ------------------------------------------------------------------
    DWM_LOG("Encoding PNG: %ux%u, stride=%u", desc.Width, desc.Height, mapped.RowPitch);
    pngDataOut = EncodeBgraToPng(
        static_cast<const uint8_t*>(mapped.pData),
        static_cast<int>(desc.Width),
        static_cast<int>(desc.Height),
        static_cast<int>(mapped.RowPitch));

    context->Unmap(staging.Get(), 0);

    if (pngDataOut.empty()) {
        errorOut = "PNG encoding produced empty output";
        return false;
    }

    DWM_LOG("Capture complete: %zu bytes PNG", pngDataOut.size());
    return true;
}

// SEH wrapper — MSVC forbids __try/__except in functions with C++ destructors.
// This thin wrapper has no such objects; it forwards to the inner function.
static bool ExecuteDwmCapture(
    HWND hwnd,
    ID3D11Device* device,
    ID3D11DeviceContext* context,
    std::vector<uint8_t>& pngDataOut,
    std::string& errorOut)
{
    __try {
        return ExecuteDwmCaptureInner(hwnd, device, context, pngDataOut, errorOut);
    } __except(EXCEPTION_EXECUTE_HANDLER) {
        errorOut = "DWM capture crashed (SEH caught access violation)";
        return false;
    }
}

void CaptureWorker::Execute() {
    // Validate HWND before capture (per D-12)
    if (!IsWindow(hwnd_)) {
        SetError("Invalid or destroyed HWND");
        return;
    }

    // Ensure we have the DwmGetDxSharedSurface function pointer
    if (!EnsureDwmFunction()) {
        SetError("DwmGetDxSharedSurface not available on this Windows version");
        return;
    }

    // Run the DWM shared surface pipeline in a SEH-wrapped function (per D-11)
    std::string error;
    bool ok = ExecuteDwmCapture(
        hwnd_,
        device_.Get(),
        context_.Get(),
        pngData_,
        error);

    if (!ok) {
        SetError(error);
    }
}

// --------------------------------------------------------------------------
// OnOK / OnError — runs on main thread (NAPI types safe here)
// --------------------------------------------------------------------------

void CaptureWorker::OnOK() {
    auto buffer = Napi::Buffer<uint8_t>::Copy(Env(), pngData_.data(), pngData_.size());
    deferred_.Resolve(buffer);
}

void CaptureWorker::OnError(const Napi::Error& error) {
    deferred_.Reject(error.Value());
}