#include "IddBridge.h"
#include <string>
#include <chrono>

IddBridge::IddBridge() = default;

IddBridge::~IddBridge() {
    if (frame_acquired_) {
        ReleaseFrame();
    }
    if (shared_) {
        UnmapViewOfFile(shared_);
        shared_ = nullptr;
    }
    if (hMap_) {
        CloseHandle(hMap_);
        hMap_ = nullptr;
    }
    if (hReadyEvent_) {
        CloseHandle(hReadyEvent_);
        hReadyEvent_ = nullptr;
    }
    if (hConsumedEvent_) {
        CloseHandle(hConsumedEvent_);
        hConsumedEvent_ = nullptr;
    }
}

bool IddBridge::Initialize() {
    D3D_FEATURE_LEVEL featureLevels[] = {
        D3D_FEATURE_LEVEL_11_1,
        D3D_FEATURE_LEVEL_11_0,
    };
    D3D_FEATURE_LEVEL featureLevel;
    HRESULT hr = D3D11CreateDevice(
        nullptr,
        D3D_DRIVER_TYPE_HARDWARE,
        nullptr,
        D3D11_CREATE_DEVICE_BGRA_SUPPORT,
        featureLevels,
        ARRAYSIZE(featureLevels),
        D3D11_SDK_VERSION,
        &device_,
        &featureLevel,
        &context_
    );
    if (FAILED(hr)) {
        hr = D3D11CreateDevice(
            nullptr,
            D3D_DRIVER_TYPE_WARP,
            nullptr,
            D3D11_CREATE_DEVICE_BGRA_SUPPORT,
            featureLevels,
            ARRAYSIZE(featureLevels),
            D3D11_SDK_VERSION,
            &device_,
            &featureLevel,
            &context_
        );
        if (FAILED(hr)) {
            return false;
        }
    }

    std::wstring mapName = L"Global\\DisplayFlowIddFrame";
    std::wstring readyName = L"Global\\DisplayFlowIddReady";
    std::wstring consumedName = L"Global\\DisplayFlowIddConsumed";

    hMap_ = OpenFileMappingW(FILE_MAP_READ, FALSE, mapName.c_str());
    if (!hMap_) {
        // Fallback to Local (Simulation mode)
        mapName = L"Local\\DisplayFlowIddFrame";
        readyName = L"Local\\DisplayFlowIddReady";
        consumedName = L"Local\\DisplayFlowIddConsumed";
        hMap_ = OpenFileMappingW(FILE_MAP_READ, FALSE, mapName.c_str());
    }

    if (!hMap_) {
        // Create for testing; in production UMDF service should create
        shared_size_ = 64 * 1024 * 1024; // 64MB
        hMap_ = CreateFileMappingW(INVALID_HANDLE_VALUE, nullptr, PAGE_READWRITE, 0, (DWORD)shared_size_, mapName.c_str());
        if (!hMap_) {
            return false;
        }
    } else {
        // Unknown size; assume max
        shared_size_ = 64 * 1024 * 1024;
    }

    shared_ = (uint8_t*)MapViewOfFile(hMap_, FILE_MAP_READ, 0, 0, shared_size_);
    if (!shared_) {
        return false;
    }

    hReadyEvent_ = OpenEventW(SYNCHRONIZE, FALSE, readyName.c_str());
    if (!hReadyEvent_) {
        hReadyEvent_ = CreateEventW(nullptr, FALSE, FALSE, readyName.c_str());
    }
    hConsumedEvent_ = OpenEventW(EVENT_MODIFY_STATE, FALSE, consumedName.c_str());
    if (!hConsumedEvent_) {
        hConsumedEvent_ = CreateEventW(nullptr, FALSE, FALSE, consumedName.c_str());
    }

    return true;
}

bool IddBridge::EnsureTexture(int width, int height) {
    if (gpuTexture_ && tex_width_ == width && tex_height_ == height) return true;
    gpuTexture_.Reset();

    D3D11_TEXTURE2D_DESC desc = {};
    desc.Width = width;
    desc.Height = height;
    desc.MipLevels = 1;
    desc.ArraySize = 1;
    desc.Format = DXGI_FORMAT_B8G8R8A8_UNORM;
    desc.SampleDesc.Count = 1;
    desc.Usage = D3D11_USAGE_DEFAULT;
    desc.BindFlags = 0;

    HRESULT hr = device_->CreateTexture2D(&desc, nullptr, &gpuTexture_);
    if (FAILED(hr)) return false;
    tex_width_ = width;
    tex_height_ = height;
    return true;
}

bool IddBridge::CaptureFrame(ComPtr<ID3D11Texture2D>& texture) {
    if (!shared_ || !hReadyEvent_) return false;

    DWORD wait = WaitForSingleObject(hReadyEvent_, 100);
    if (wait != WAIT_OBJECT_0) {
        return false;
    }

    auto* hdr = reinterpret_cast<const IddSharedHeader*>(shared_);
    if (hdr->format != 0) {
        return false;
    }
    int w = (int)hdr->width;
    int h = (int)hdr->height;
    int stride = (int)hdr->stride;
    const uint8_t* pixels = shared_ + sizeof(IddSharedHeader);

    if (!EnsureTexture(w, h)) return false;

    D3D11_BOX box = {};
    box.left = 0; box.right = w;
    box.top = 0; box.bottom = h;
    box.front = 0; box.back = 1;

    context_->UpdateSubresource(gpuTexture_.Get(), 0, nullptr, pixels, stride, 0);

    frame_acquired_ = true;
    texture = gpuTexture_;

    if (hConsumedEvent_) {
        SetEvent(hConsumedEvent_);
    }
    return true;
}

void IddBridge::ReleaseFrame() {
    frame_acquired_ = false;
}

// ============================================================================
// IddProducer Implementation
// ============================================================================

IddProducer::IddProducer() = default;

IddProducer::~IddProducer() {
    if (shared_) UnmapViewOfFile(shared_);
    if (hMap_) CloseHandle(hMap_);
    if (hReadyEvent_) CloseHandle(hReadyEvent_);
    if (hConsumedEvent_) CloseHandle(hConsumedEvent_);
}

bool IddProducer::Initialize() {
    std::wstring mapName = L"Local\\DisplayFlowIddFrame";
    std::wstring readyName = L"Local\\DisplayFlowIddReady";
    std::wstring consumedName = L"Local\\DisplayFlowIddConsumed";

    shared_size_ = 64 * 1024 * 1024; // 64MB
    
    // Producer prefers Creating
    hMap_ = CreateFileMappingW(INVALID_HANDLE_VALUE, nullptr, PAGE_READWRITE, 0, (DWORD)shared_size_, mapName.c_str());
    if (!hMap_) {
        // Maybe already exists
        hMap_ = OpenFileMappingW(FILE_MAP_WRITE, FALSE, mapName.c_str());
        if (!hMap_) return false;
    }

    shared_ = (uint8_t*)MapViewOfFile(hMap_, FILE_MAP_WRITE, 0, 0, shared_size_);
    if (!shared_) return false;

    hReadyEvent_ = CreateEventW(nullptr, FALSE, FALSE, readyName.c_str());
    if (!hReadyEvent_) {
        hReadyEvent_ = OpenEventW(EVENT_MODIFY_STATE, FALSE, readyName.c_str());
        if (!hReadyEvent_) return false;
    }

    hConsumedEvent_ = CreateEventW(nullptr, FALSE, FALSE, consumedName.c_str());
    if (!hConsumedEvent_) {
        hConsumedEvent_ = OpenEventW(SYNCHRONIZE, FALSE, consumedName.c_str());
        // If failed, consumer might create it later.
    }
    
    return true;
}

bool IddProducer::SubmitFrame(const void* data, uint32_t width, uint32_t height, uint32_t stride) {
    if (!shared_ || !hReadyEvent_) return false;
    
    // Ensure consumer created the event if we missed it
    if (!hConsumedEvent_) {
        hConsumedEvent_ = OpenEventW(SYNCHRONIZE, FALSE, L"Local\\DisplayFlowIddConsumed");
    }

    // Optional: Wait for previous frame consumption to avoid tearing/overflow
    // if (hConsumedEvent_) WaitForSingleObject(hConsumedEvent_, 0); 
    // For now, overwrite mode is faster for streaming.

    IddSharedHeader* hdr = reinterpret_cast<IddSharedHeader*>(shared_);
    
    uint32_t needed = sizeof(IddSharedHeader) + stride * height;
    if (needed > shared_size_) return false;

    hdr->width = width;
    hdr->height = height;
    hdr->format = 0;
    hdr->stride = stride;
    hdr->timestamp = (uint64_t)std::chrono::duration_cast<std::chrono::milliseconds>(
                            std::chrono::high_resolution_clock::now().time_since_epoch()).count();
    hdr->frameId = ++frameId_;
    hdr->dataSize = stride * height;
    
    std::memcpy(shared_ + sizeof(IddSharedHeader), data, hdr->dataSize);
    
    SetEvent(hReadyEvent_);
    return true;
}