#include "ScreenCapture.h"
#include <iostream>

ScreenCapture::ScreenCapture() = default;

ScreenCapture::~ScreenCapture() {
    if (frame_acquired_) {
        ReleaseFrame();
    }
}

bool ScreenCapture::Initialize() {
    HRESULT hr = S_OK;

    // Create D3D11 Device and Context
    D3D_FEATURE_LEVEL featureLevels[] = {
        D3D_FEATURE_LEVEL_11_1,
        D3D_FEATURE_LEVEL_11_0,
    };
    D3D_FEATURE_LEVEL featureLevel;

    hr = D3D11CreateDevice(
        nullptr,
        D3D_DRIVER_TYPE_HARDWARE,
        nullptr,
        D3D11_CREATE_DEVICE_BGRA_SUPPORT, // Needed for GDI compatibility if used, but generally good for D2D
        featureLevels,
        ARRAYSIZE(featureLevels),
        D3D11_SDK_VERSION,
        &device_,
        &featureLevel,
        &context_
    );

    if (FAILED(hr)) {
        std::cerr << "Failed to create D3D11 device: " << std::hex << hr << std::endl;
        return false;
    }

    // Get DXGI Device
    ComPtr<IDXGIDevice> dxgiDevice;
    hr = device_.As(&dxgiDevice);
    if (FAILED(hr)) return false;

    // Get DXGI Adapter
    ComPtr<IDXGIAdapter> dxgiAdapter;
    hr = dxgiDevice->GetAdapter(&dxgiAdapter);
    if (FAILED(hr)) return false;

    // Get DXGI Output (Monitor 0)
    ComPtr<IDXGIOutput> dxgiOutput;
    hr = dxgiAdapter->EnumOutputs(0, &dxgiOutput);
    if (FAILED(hr)) {
        std::cerr << "Failed to get DXGI output (monitor connected?)" << std::endl;
        return false;
    }

    // QI for Output1 to support Duplication
    ComPtr<IDXGIOutput1> dxgiOutput1;
    hr = dxgiOutput.As(&dxgiOutput1);
    if (FAILED(hr)) return false;

    // Create Desktop Duplication
    hr = dxgiOutput1->DuplicateOutput(device_.Get(), &duplication_);
    if (FAILED(hr)) {
        std::cerr << "Failed to duplicate output. Error: " << std::hex << hr << std::endl;
        // Common errors: E_ACCESSDENIED (already duplicated), DXGI_ERROR_UNSUPPORTED (switchable graphics)
        return false;
    }

    return true;
}

bool ScreenCapture::InitializeWithOutputIndex(int index) {
    HRESULT hr = S_OK;
    D3D_FEATURE_LEVEL featureLevels[] = {
        D3D_FEATURE_LEVEL_11_1,
        D3D_FEATURE_LEVEL_11_0,
    };
    D3D_FEATURE_LEVEL featureLevel;
    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)) {
        return false;
    }
    ComPtr<IDXGIDevice> dxgiDevice;
    hr = device_.As(&dxgiDevice);
    if (FAILED(hr)) return false;
    ComPtr<IDXGIAdapter> dxgiAdapter;
    hr = dxgiDevice->GetAdapter(&dxgiAdapter);
    if (FAILED(hr)) return false;
    ComPtr<IDXGIOutput> dxgiOutput;
    hr = dxgiAdapter->EnumOutputs(index, &dxgiOutput);
    if (FAILED(hr)) {
        return false;
    }
    ComPtr<IDXGIOutput1> dxgiOutput1;
    hr = dxgiOutput.As(&dxgiOutput1);
    if (FAILED(hr)) return false;
    hr = dxgiOutput1->DuplicateOutput(device_.Get(), &duplication_);
    if (FAILED(hr)) {
        return false;
    }
    return true;
}
bool ScreenCapture::CaptureFrame(ComPtr<ID3D11Texture2D>& texture) {
    if (frame_acquired_) {
        ReleaseFrame();
    }

    DXGI_OUTDUPL_FRAME_INFO frameInfo;
    ComPtr<IDXGIResource> resource;
    
    // Timeout 100ms
    HRESULT hr = duplication_->AcquireNextFrame(100, &frameInfo, &resource);
    
    if (hr == DXGI_ERROR_WAIT_TIMEOUT) {
        return false; // No new frame
    }
    
    if (FAILED(hr)) {
        // Maybe device lost or resolution changed
        std::cerr << "AcquireNextFrame failed: " << std::hex << hr << std::endl;
        return false;
    }

    frame_acquired_ = true;

    // Only process if we have a desktop image update
    if (frameInfo.LastPresentTime.QuadPart == 0) {
        return false; // Only cursor moved or something else, no image update? 
        // Actually AcquireNextFrame returns even if only cursor updated.
        // But for video stream we might want to send anyway or skip. 
        // If resource is null, it means no desktop image update.
    }
    
    if (!resource) {
        return false;
    }

    hr = resource.As(&texture);
    if (FAILED(hr)) return false;

    return true;
}

void ScreenCapture::ReleaseFrame() {
    if (frame_acquired_ && duplication_) {
        duplication_->ReleaseFrame();
        frame_acquired_ = false;
    }
}