#include <windows.h>
#include <iostream>
#include <string>
#include <vector>
#include <wincodec.h>
#include <combaseapi.h>
#include <dwmapi.h>

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

class WindowScreenshot {
private:
    HWND hwnd;
    HDC source_hdc;
    HDC dest_hdc;
    HBITMAP hbmp;
    BITMAPINFO bmi;
    void* buffer;
    RECT clip_rect;
    int width, height;

public:
    WindowScreenshot() : hwnd(nullptr), source_hdc(nullptr), dest_hdc(nullptr),
        hbmp(nullptr), buffer(nullptr), width(0), height(0) {
    }

    ~WindowScreenshot() {
        cleanup();
    }

    // Find window by title (similar to FFmpeg's title= logic)
    bool findWindowByTitle(const std::string& windowTitle) {
        std::wstring wideTitle(windowTitle.begin(), windowTitle.end());
        hwnd = FindWindowW(NULL, wideTitle.c_str());

        if (!hwnd) {
            std::wcout << L"Can't find window '" << wideTitle.c_str() << L"'" << std::endl;
            return false;
        }

        std::wcout << L"Found window: " << wideTitle.c_str() << std::endl;
        return true;
    }

    // Initialize capture setup (similar to gdigrab_read_header logic)
    bool initialize() {
        if (!hwnd) {
            std::cout << "No window handle found. Call findWindowByTitle first." << std::endl;
            return false;
        }

        // Get device context for the window
        source_hdc = GetDC(hwnd);
        if (!source_hdc) {
            std::cout << "Couldn't get window device context. Error: " << GetLastError() << std::endl;
            return false;
        }

        // Get window dimensions (use GetWindowRect for full window including borders)
        if (!GetWindowRect(hwnd, &clip_rect)) {
            std::cout << "Couldn't get window rectangle. Error: " << GetLastError() << std::endl;
            return false;
        }

        // Convert screen coordinates to relative coordinates
        width = clip_rect.right - clip_rect.left;
        height = clip_rect.bottom - clip_rect.top;

        if (width <= 0 || height <= 0) {
            std::cout << "Invalid window dimensions: " << width << "x" << height << std::endl;
            return false;
        }

        std::cout << "Window size: " << width << "x" << height << std::endl;

        // Create compatible DC for destination
        dest_hdc = CreateCompatibleDC(source_hdc);
        if (!dest_hdc) {
            std::cout << "CreateCompatibleDC failed. Error: " << GetLastError() << std::endl;
            return false;
        }

        // Get bits per pixel
        int bpp = GetDeviceCaps(source_hdc, BITSPIXEL);
        std::cout << "Bits per pixel: " << bpp << std::endl;

        // Create DIB (Device Independent Bitmap) - similar to FFmpeg's logic
        memset(&bmi, 0, sizeof(bmi));
        bmi.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
        bmi.bmiHeader.biWidth = width;
        bmi.bmiHeader.biHeight = -height;  // Negative for top-down bitmap
        bmi.bmiHeader.biPlanes = 1;
        bmi.bmiHeader.biBitCount = bpp;
        bmi.bmiHeader.biCompression = BI_RGB;
        bmi.bmiHeader.biSizeImage = 0;

        // Create DIB section
        hbmp = CreateDIBSection(dest_hdc, &bmi, DIB_RGB_COLORS, &buffer, NULL, 0);
        if (!hbmp) {
            std::cout << "CreateDIBSection failed. Error: " << GetLastError() << std::endl;
            return false;
        }

        // Select bitmap into destination DC
        if (!SelectObject(dest_hdc, hbmp)) {
            std::cout << "SelectObject failed. Error: " << GetLastError() << std::endl;
            return false;
        }

        std::cout << "Screenshot capture initialized successfully." << std::endl;
        return true;
    }

    // Capture screenshot of hidden/covered window using window redraw method
    bool captureHiddenWindow() {
        if (!dest_hdc || !hbmp) {
            std::cout << "Not initialized. Call initialize() first." << std::endl;
            return false;
        }

        // Method 1: Try to force window to redraw itself to our DC
        RECT windowRect;
        if (!GetWindowRect(hwnd, &windowRect)) {
            std::cout << "Failed to get window rectangle. Error: " << GetLastError() << std::endl;
            return false;
        }

        // Clear the destination bitmap first
        RECT clearRect = {0, 0, width, height};
        FillRect(dest_hdc, &clearRect, (HBRUSH)GetStockObject(WHITE_BRUSH));

        // Try to get the window to paint itself to our DC
        // This works by sending paint messages to the window
        HDC window_dc = GetWindowDC(hwnd);
        if (window_dc) {
            // Try copying from window DC
            if (BitBlt(dest_hdc, 0, 0, width, height, window_dc, 0, 0, SRCCOPY)) {
                ReleaseDC(hwnd, window_dc);
                std::cout << "Hidden window captured using Window DC method!" << std::endl;
                return true;
            }
            ReleaseDC(hwnd, window_dc);
        }

        // Method 2: Try sending WM_PAINT message to force redraw
        std::cout << "Window DC failed, trying paint message method..." << std::endl;
        
        // Save current bitmap selection
        HBITMAP old_bmp = (HBITMAP)SelectObject(dest_hdc, hbmp);
        
        // Try to get the window to paint to our DC
        SendMessage(hwnd, WM_PAINT, (WPARAM)dest_hdc, 0);
        
        // Also try WM_PRINTCLIENT which some windows respond to
        if (SendMessage(hwnd, WM_PRINTCLIENT, (WPARAM)dest_hdc, PRF_CLIENT | PRF_CHILDREN)) {
            std::cout << "Hidden window captured using WM_PRINTCLIENT!" << std::endl;
            return true;
        }

        std::cout << "Paint message methods failed." << std::endl;
        return false;
    }

    // Capture screenshot using memory manipulation (for covered windows)
    bool captureWindowMemory() {
        if (!dest_hdc || !hbmp) {
            std::cout << "Not initialized. Call initialize() first." << std::endl;
            return false;
        }

        // Get window's device context
        HDC window_dc = GetDC(hwnd);
        if (!window_dc) {
            std::cout << "Failed to get window DC. Error: " << GetLastError() << std::endl;
            return false;
        }

        // Try to capture from the window's memory
        bool success = false;
        
        // Method 1: Direct BitBlt from window DC (works for some hidden windows)
        if (BitBlt(dest_hdc, 0, 0, width, height, window_dc, 0, 0, SRCCOPY)) {
            success = true;
            std::cout << "Window captured from memory using direct BitBlt!" << std::endl;
        } else {
            // Method 2: Try with different raster operations
            if (BitBlt(dest_hdc, 0, 0, width, height, window_dc, 0, 0, SRCCOPY | CAPTUREBLT)) {
                success = true;
                std::cout << "Window captured using CAPTUREBLT!" << std::endl;
            }
        }

        ReleaseDC(hwnd, window_dc);
        return success;
    }

    // Alternative method: Minimize covering windows temporarily
    bool captureByMinimizingOthers() {
        if (!dest_hdc || !hbmp) {
            std::cout << "Not initialized. Call initialize() first." << std::endl;
            return false;
        }

        std::cout << "Attempting to capture by temporarily managing window visibility..." << std::endl;

        // Get our target window's Z-order position
        HWND hwndAfter = GetWindow(hwnd, GW_HWNDPREV);
        
        // Collect windows that might be covering our target
        std::vector<HWND> coveringWindows;
        HWND currentWnd = GetTopWindow(GetDesktopWindow());
        
        RECT targetRect;
        GetWindowRect(hwnd, &targetRect);
        
        // Find windows that might be covering our target window
        while (currentWnd && currentWnd != hwnd) {
            if (IsWindowVisible(currentWnd) && !IsIconic(currentWnd)) {
                RECT currentRect;
                if (GetWindowRect(currentWnd, &currentRect)) {
                    // Check if this window overlaps with our target
                    RECT intersection;
                    if (IntersectRect(&intersection, &targetRect, &currentRect)) {
                        coveringWindows.push_back(currentWnd);
                    }
                }
            }
            currentWnd = GetWindow(currentWnd, GW_HWNDNEXT);
        }

        std::cout << "Found " << coveringWindows.size() << " potentially covering windows." << std::endl;

        // Temporarily minimize covering windows
        std::vector<HWND> minimizedWindows;
        for (HWND covering : coveringWindows) {
            if (!IsIconic(covering)) {
                ShowWindow(covering, SW_MINIMIZE);
                minimizedWindows.push_back(covering);
            }
        }

        Sleep(500); // Give time for windows to minimize

        // Now try to capture our target window
        bool success = false;
        
        // Bring our target to front
        SetWindowPos(hwnd, HWND_TOP, 0, 0, 0, 0, SWP_NOMOVE | SWP_NOSIZE | SWP_NOACTIVATE);
        Sleep(200);

        // Capture from desktop at window position
        RECT windowRect;
        if (GetWindowRect(hwnd, &windowRect)) {
            HDC desktop_hdc = GetDC(NULL);
            if (desktop_hdc) {
                success = BitBlt(dest_hdc, 0, 0, width, height,
                               desktop_hdc, windowRect.left, windowRect.top, SRCCOPY);
                ReleaseDC(NULL, desktop_hdc);
            }
        }

        // Restore the minimized windows
        for (HWND minimized : minimizedWindows) {
            ShowWindow(minimized, SW_RESTORE);
        }

        if (success) {
            std::cout << "Window captured by temporarily minimizing covering windows!" << std::endl;
        } else {
            std::cout << "Failed to capture even after minimizing covering windows." << std::endl;
        }

        return success;
    }

    // Alternative PrintWindow method with proper sizing (using same approach as captureScreenshot)
    bool captureScreenshotPrintWindowAlt() {
        if (!hwnd) {
            std::cout << "No window handle found." << std::endl;
            return false;
        }

        // Use the same window rectangle approach as captureScreenshot
        RECT windowRect;
        if (!GetWindowRect(hwnd, &windowRect)) {
            std::cout << "Failed to get window rectangle. Error: " << GetLastError() << std::endl;
            return false;
        }

        // Use the same dimensions as our initialized bitmap (width, height)
        // This ensures consistency with the captureScreenshot method
        int win_width = width;  // Use our class member width
        int win_height = height; // Use our class member height

        std::cout << "Using consistent dimensions for PrintWindow: " << win_width << "x" << win_height << std::endl;

        // Create a new bitmap with the same size as our target
        HDC screen_dc = GetDC(NULL);
        HDC temp_dc = CreateCompatibleDC(screen_dc);
        HBITMAP temp_bitmap = CreateCompatibleBitmap(screen_dc, win_width, win_height);
        HBITMAP old_bitmap = (HBITMAP)SelectObject(temp_dc, temp_bitmap);

        // Use PrintWindow to capture - try full window first
        bool success = PrintWindow(hwnd, temp_dc, 0);

        if (!success) {
            std::cout << "PrintWindow (full) failed, trying client area..." << std::endl;
            success = PrintWindow(hwnd, temp_dc, PW_CLIENTONLY);
        }

        if (success) {
            // Copy the captured image directly to our main bitmap (no stretching needed)
            if (dest_hdc && hbmp) {
                BitBlt(dest_hdc, 0, 0, win_width, win_height,
                       temp_dc, 0, 0, SRCCOPY);
                std::cout << "Window screenshot captured with consistent sizing!" << std::endl;
            } else {
                std::cout << "Destination DC not ready." << std::endl;
                success = false;
            }
        } else {
            std::cout << "PrintWindow failed completely. Error: " << GetLastError() << std::endl;
        }

        // Cleanup
        SelectObject(temp_dc, old_bitmap);
        DeleteObject(temp_bitmap);
        DeleteDC(temp_dc);
        ReleaseDC(NULL, screen_dc);

        return success;
    }

    // Capture screenshot using PrintWindow (works even when window is behind others)
    bool captureScreenshotPrintWindow() {
        if (!dest_hdc || !hbmp) {
            std::cout << "Not initialized. Call initialize() first." << std::endl;
            return false;
        }

        // Try PrintWindow without PW_CLIENTONLY first (captures entire window including borders)
        if (PrintWindow(hwnd, dest_hdc, 0)) {
            std::cout << "Window screenshot captured successfully using PrintWindow (full window)!" << std::endl;
            return true;
        }

        std::cout << "PrintWindow (full window) failed, trying client area only..." << std::endl;

        // Try with PW_CLIENTONLY flag
        if (PrintWindow(hwnd, dest_hdc, PW_CLIENTONLY)) {
            std::cout << "Window screenshot captured successfully using PrintWindow (client area)!" << std::endl;
            return true;
        }

        std::cout << "PrintWindow failed. Error: " << GetLastError() << std::endl;
        return false;
    }

    // Capture screenshot (similar to gdigrab_read_packet BitBlt logic)
    bool captureScreenshot() {
        if (!source_hdc || !dest_hdc || !hbmp) {
            std::cout << "Not initialized. Call initialize() first." << std::endl;
            return false;
        }

        // Method 1: Try to capture directly from window DC first
        if (BitBlt(dest_hdc, 0, 0, width, height,
            source_hdc, 0, 0, SRCCOPY)) {
            std::cout << "Window screenshot captured successfully (direct method)!" << std::endl;
            return true;
        }

        std::cout << "Direct capture failed, trying alternative method..." << std::endl;

        // Method 2: Bring window to front temporarily and capture from desktop
        bool wasVisible = IsWindowVisible(hwnd);
        bool wasMinimized = IsIconic(hwnd);
        
        if (wasMinimized) {
            ShowWindow(hwnd, SW_RESTORE);
            Sleep(100); // Give time for window to restore
        }

        // Bring window to front (but don't steal focus)
        SetWindowPos(hwnd, HWND_TOPMOST, 0, 0, 0, 0, SWP_NOMOVE | SWP_NOSIZE | SWP_NOACTIVATE);
        SetWindowPos(hwnd, HWND_NOTOPMOST, 0, 0, 0, 0, SWP_NOMOVE | SWP_NOSIZE | SWP_NOACTIVATE);
        
        Sleep(200); // Give time for window to come to front

        // Get the window's screen coordinates
        RECT windowRect;
        if (!GetWindowRect(hwnd, &windowRect)) {
            std::cout << "Failed to get window rectangle. Error: " << GetLastError() << std::endl;
            return false;
        }

        // Get the desktop DC to capture from screen coordinates
        HDC desktop_hdc = GetDC(NULL);
        if (!desktop_hdc) {
            std::cout << "Failed to get desktop DC. Error: " << GetLastError() << std::endl;
            return false;
        }

        // Capture from the desktop at the window's position
        bool success = BitBlt(dest_hdc, 0, 0, width, height,
            desktop_hdc, windowRect.left, windowRect.top, SRCCOPY);

        ReleaseDC(NULL, desktop_hdc);

        if (!success) {
            std::cout << "BitBlt failed to capture window image. Error: " << GetLastError() << std::endl;
            return false;
        }

        std::cout << "Window screenshot captured successfully (desktop method)!" << std::endl;
        return true;
    }

    // Save screenshot as PNG file using WIC
    bool saveAsPNG(const std::string& filename) {
        if (!buffer || !hbmp) {
            std::cout << "No image data to save." << std::endl;
            return false;
        }

        HRESULT hr = CoInitialize(NULL);
        if (FAILED(hr)) {
            std::cout << "Failed to initialize COM" << std::endl;
            return false;
        }

        IWICImagingFactory* pFactory = NULL;
        IWICBitmapEncoder* pEncoder = NULL;
        IWICBitmapFrameEncode* pFrameEncode = NULL;
        IWICStream* pStream = NULL;
        IPropertyBag2* pPropertybag = NULL;

        bool success = false;

        do {
            // Create WIC factory
            hr = CoCreateInstance(CLSID_WICImagingFactory, NULL, CLSCTX_INPROC_SERVER, 
                                IID_IWICImagingFactory, (LPVOID*)&pFactory);
            if (FAILED(hr)) break;

            // Create stream
            hr = pFactory->CreateStream(&pStream);
            if (FAILED(hr)) break;

            // Initialize stream with filename
            std::wstring wFilename(filename.begin(), filename.end());
            hr = pStream->InitializeFromFilename(wFilename.c_str(), GENERIC_WRITE);
            if (FAILED(hr)) break;

            // Create PNG encoder
            hr = pFactory->CreateEncoder(GUID_ContainerFormatPng, NULL, &pEncoder);
            if (FAILED(hr)) break;

            // Initialize encoder
            hr = pEncoder->Initialize(pStream, WICBitmapEncoderNoCache);
            if (FAILED(hr)) break;

            // Create frame
            hr = pEncoder->CreateNewFrame(&pFrameEncode, &pPropertybag);
            if (FAILED(hr)) break;

            // Initialize frame
            hr = pFrameEncode->Initialize(pPropertybag);
            if (FAILED(hr)) break;

            // Set frame size
            hr = pFrameEncode->SetSize(width, height);
            if (FAILED(hr)) break;

            // Set pixel format - use the actual format from the bitmap
            WICPixelFormatGUID formatGUID;
            if (bmi.bmiHeader.biBitCount == 32) {
                formatGUID = GUID_WICPixelFormat32bppBGRA;
            } else if (bmi.bmiHeader.biBitCount == 24) {
                formatGUID = GUID_WICPixelFormat24bppBGR;
            } else {
                formatGUID = GUID_WICPixelFormat24bppBGR; // Default fallback
            }
            hr = pFrameEncode->SetPixelFormat(&formatGUID);
            if (FAILED(hr)) break;

            // Calculate stride based on actual bit depth
            UINT stride = ((width * bmi.bmiHeader.biBitCount + 31) / 32) * 4;
            UINT imageSize = stride * height;

            // Write pixels
            hr = pFrameEncode->WritePixels(height, stride, imageSize, (BYTE*)buffer);
            if (FAILED(hr)) break;

            // Commit frame
            hr = pFrameEncode->Commit();
            if (FAILED(hr)) break;

            // Commit encoder
            hr = pEncoder->Commit();
            if (FAILED(hr)) break;

            success = true;
            std::cout << "Screenshot saved as PNG: " << filename << std::endl;

        } while (false);

        // Cleanup
        if (pPropertybag) pPropertybag->Release();
        if (pFrameEncode) pFrameEncode->Release();
        if (pEncoder) pEncoder->Release();
        if (pStream) pStream->Release();
        if (pFactory) pFactory->Release();

        CoUninitialize();

        if (!success) {
            std::cout << "Failed to save PNG file. Error code: " << hr << std::endl;
        }

        return success;
    }

    // Save screenshot as BMP file
    bool saveAsBMP(const std::string& filename) {
        if (!buffer || !hbmp) {
            std::cout << "No image data to save." << std::endl;
            return false;
        }

        HANDLE hFile = CreateFileA(filename.c_str(), GENERIC_WRITE, 0, NULL,
            CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
        if (hFile == INVALID_HANDLE_VALUE) {
            std::cout << "Cannot create file: " << filename << ". Error: " << GetLastError() << std::endl;
            return false;
        }

        // Get bitmap info
        BITMAP bmp;
        GetObject(hbmp, sizeof(BITMAP), &bmp);

        // Calculate sizes
        int frame_size = bmp.bmWidthBytes * bmp.bmHeight;
        int header_size = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER);
        int file_size = header_size + frame_size;

        // Create bitmap file header
        BITMAPFILEHEADER bfh;
        bfh.bfType = 0x4D42;  // "BM"
        bfh.bfSize = file_size;
        bfh.bfReserved1 = 0;
        bfh.bfReserved2 = 0;
        bfh.bfOffBits = header_size;

        DWORD written;

        // Write file header
        WriteFile(hFile, &bfh, sizeof(bfh), &written, NULL);

        // Write info header
        WriteFile(hFile, &bmi.bmiHeader, sizeof(bmi.bmiHeader), &written, NULL);

        // Write pixel data
        WriteFile(hFile, buffer, frame_size, &written, NULL);

        CloseHandle(hFile);
        std::cout << "Screenshot saved as: " << filename << std::endl;
        return true;
    }

    // Get window dimensions
    void getWindowSize(int& w, int& h) {
        w = width;
        h = height;
    }

private:
    void cleanup() {
        if (source_hdc && hwnd) {
            ReleaseDC(hwnd, source_hdc);
        }
        if (dest_hdc) {
            DeleteDC(dest_hdc);
        }
        if (hbmp) {
            DeleteObject(hbmp);
        }

        source_hdc = nullptr;
        dest_hdc = nullptr;
        hbmp = nullptr;
        buffer = nullptr;
        hwnd = nullptr;
    }
};

// Demo function to capture desktop and save as PNG
bool captureDesktop(const std::string& filename) {
    std::cout << "\n=== Capturing Desktop ===" << std::endl;

    HDC screen_dc = GetDC(NULL);  // Get desktop DC
    if (!screen_dc) {
        std::cout << "Failed to get desktop DC" << std::endl;
        return false;
    }

    int screen_width = GetSystemMetrics(SM_CXSCREEN);
    int screen_height = GetSystemMetrics(SM_CYSCREEN);

    std::cout << "Desktop size: " << screen_width << "x" << screen_height << std::endl;

    HDC mem_dc = CreateCompatibleDC(screen_dc);
    HBITMAP bitmap = CreateCompatibleBitmap(screen_dc, screen_width, screen_height);
    HBITMAP old_bitmap = (HBITMAP)SelectObject(mem_dc, bitmap);

    // Capture desktop
    if (!BitBlt(mem_dc, 0, 0, screen_width, screen_height, screen_dc, 0, 0, SRCCOPY)) {
        std::cout << "Desktop capture failed" << std::endl;
        SelectObject(mem_dc, old_bitmap);
        DeleteDC(mem_dc);
        DeleteObject(bitmap);
        ReleaseDC(NULL, screen_dc);
        return false;
    }

    // Save desktop capture as PNG
    BITMAPINFO bmi = {};
    bmi.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
    bmi.bmiHeader.biWidth = screen_width;
    bmi.bmiHeader.biHeight = -screen_height;  // Negative for top-down
    bmi.bmiHeader.biPlanes = 1;
    bmi.bmiHeader.biBitCount = 24;
    bmi.bmiHeader.biCompression = BI_RGB;

    // Calculate image size
    int stride = ((screen_width * 24 + 31) / 32) * 4;
    int imageSize = stride * screen_height;
    
    // Allocate buffer for pixel data
    std::vector<BYTE> pixelData(imageSize);
    
    // Get DIB bits
    if (GetDIBits(mem_dc, bitmap, 0, screen_height, pixelData.data(), &bmi, DIB_RGB_COLORS)) {
        // Save using WIC
        HRESULT hr = CoInitialize(NULL);
        if (SUCCEEDED(hr)) {
            IWICImagingFactory* pFactory = NULL;
            IWICBitmapEncoder* pEncoder = NULL;
            IWICBitmapFrameEncode* pFrameEncode = NULL;
            IWICStream* pStream = NULL;
            IPropertyBag2* pPropertybag = NULL;

            bool success = false;

            do {
                hr = CoCreateInstance(CLSID_WICImagingFactory, NULL, CLSCTX_INPROC_SERVER, 
                                    IID_IWICImagingFactory, (LPVOID*)&pFactory);
                if (FAILED(hr)) break;

                hr = pFactory->CreateStream(&pStream);
                if (FAILED(hr)) break;

                std::wstring wFilename(filename.begin(), filename.end());
                hr = pStream->InitializeFromFilename(wFilename.c_str(), GENERIC_WRITE);
                if (FAILED(hr)) break;

                hr = pFactory->CreateEncoder(GUID_ContainerFormatPng, NULL, &pEncoder);
                if (FAILED(hr)) break;

                hr = pEncoder->Initialize(pStream, WICBitmapEncoderNoCache);
                if (FAILED(hr)) break;

                hr = pEncoder->CreateNewFrame(&pFrameEncode, &pPropertybag);
                if (FAILED(hr)) break;

                hr = pFrameEncode->Initialize(pPropertybag);
                if (FAILED(hr)) break;

                hr = pFrameEncode->SetSize(screen_width, screen_height);
                if (FAILED(hr)) break;

                WICPixelFormatGUID formatGUID = GUID_WICPixelFormat24bppBGR;
                hr = pFrameEncode->SetPixelFormat(&formatGUID);
                if (FAILED(hr)) break;

                hr = pFrameEncode->WritePixels(screen_height, stride, imageSize, pixelData.data());
                if (FAILED(hr)) break;

                hr = pFrameEncode->Commit();
                if (FAILED(hr)) break;

                hr = pEncoder->Commit();
                if (FAILED(hr)) break;

                success = true;
                std::cout << "Desktop screenshot saved as: " << filename << std::endl;

            } while (false);

            if (pPropertybag) pPropertybag->Release();
            if (pFrameEncode) pFrameEncode->Release();
            if (pEncoder) pEncoder->Release();
            if (pStream) pStream->Release();
            if (pFactory) pFactory->Release();

            CoUninitialize();

            if (!success) {
                std::cout << "Failed to save desktop PNG. Error: " << hr << std::endl;
            }
        }
    }

    // Cleanup
    SelectObject(mem_dc, old_bitmap);
    DeleteDC(mem_dc);
    DeleteObject(bitmap);
    ReleaseDC(NULL, screen_dc);
    
    std::cout << "Desktop capture completed!" << std::endl;
    return true;
}

int main() {
    std::cout << "=== Window Screenshot Demo ===" << std::endl;
    std::cout << "Based on FFmpeg gdigrab BitBlt logic" << std::endl;
    std::cout << "=====================================" << std::endl;

    WindowScreenshot screenshot;

    // Example 1: Capture a specific window (like "Webex Beta" in your original command)
    std::string windowTitle;
    std::cout << "\nEnter window title to capture (or press Enter for 'Calculator'): ";
    std::getline(std::cin, windowTitle);

    if (windowTitle.empty()) {
        windowTitle = "Calculator";  // Default
    }

    if (screenshot.findWindowByTitle(windowTitle)) {
        if (screenshot.initialize()) {
            // Try hidden window capture methods (no PrintWindow)
            bool captured = false;
            
            /* Method 1: Try window memory capture
            captured = screenshot.captureScreenshotPrintWindowAlt();
            
            // Method 2: If that fails, try hidden window redraw method
            if (!captured) {
                std::cout << "Window memory capture failed, trying hidden window method..." << std::endl;
                captured = screenshot.captureHiddenWindow();
            }
             */        
            // Method 3: Last resort - traditional BitBlt method
            if (!captured) {
                std::cout << "All advanced methods failed, trying traditional BitBlt method..." << std::endl;
                captured = screenshot.captureScreenshot();
            }
            
            if (captured) {
                // Save as PNG with the requested filename
                std::string filename = "output.png";
                screenshot.saveAsPNG(filename);

                int w, h;
                screenshot.getWindowSize(w, h);
                std::cout << "Captured window: " << windowTitle << " (" << w << "x" << h << ")" << std::endl;
            } else {
                std::cout << "Failed to capture the window using any method." << std::endl;
            }
        }
    }
    else {
        std::cout << "Window '" << windowTitle << "' not found." << std::endl;
        std::cout << "Make sure the window exists and is visible." << std::endl;
    }

    std::cout << "\nPress Enter to exit...";
    std::cin.get();
    return 0;
}