#include "PCH.h"
#include "ScreenCapture.h"
#include "logger.h"

namespace ScreenCapture {
    
    // Convert string to ImageFormat
    ImageFormat StringToImageFormat(const std::string& str) {
        std::string lower = str;
        std::transform(lower.begin(), lower.end(), lower.begin(), ::tolower);
        
        if (lower == "png") return ImageFormat::PNG;
        if (lower == "jpg" || lower == "jpeg") return ImageFormat::JPEG;
        if (lower == "bmp") return ImageFormat::BMP;
        if (lower == "tif" || lower == "tiff") return ImageFormat::TIF;
        if (lower == "gif") return ImageFormat::GIF;
        if (lower == "gif_multiframe" || lower == "gif_multi") return ImageFormat::GIF_MULTIFRAME;
        if (lower == "dds") return ImageFormat::DDS;
        
        return ImageFormat::PNG; // Default
    }
    
    // Convert string to TiffCompression
    TiffCompression StringToTiffCompression(const std::string& str) {
        std::string lower = str;
        std::transform(lower.begin(), lower.end(), lower.begin(), ::tolower);
        
        if (lower == "none") return TiffCompression::None;
        if (lower == "rle") return TiffCompression::RLE;
        if (lower == "lzw") return TiffCompression::LZW;
        if (lower == "zip") return TiffCompression::ZIP;
        
        return TiffCompression::None; // Default
    }
    
    // Convert string to DDSCompression
    DDSCompression StringToDDSCompression(const std::string& str) {
        std::string lower = str;
        std::transform(lower.begin(), lower.end(), lower.begin(), ::tolower);
        
        if (lower == "bc1") return DDSCompression::BC1;
        if (lower == "bc2") return DDSCompression::BC2;
        if (lower == "bc3") return DDSCompression::BC3;
        if (lower == "bc4") return DDSCompression::BC4;
        if (lower == "bc5") return DDSCompression::BC5;
        if (lower == "bc6h") return DDSCompression::BC6H;
        if (lower == "bc7") return DDSCompression::BC7;
        if (lower == "rgba") return DDSCompression::RGBA;
        
        return DDSCompression::BC7; // Default
    }
    
    // Get WIC codec GUID
    GUID GetWICCodec(ImageFormat format) {
        switch (format) {
            case ImageFormat::PNG: return GUID_ContainerFormatPng;
            case ImageFormat::JPEG: return GUID_ContainerFormatJpeg;
            case ImageFormat::BMP: return GUID_ContainerFormatBmp;
            case ImageFormat::TIF: return GUID_ContainerFormatTiff;
            case ImageFormat::GIF:
            case ImageFormat::GIF_MULTIFRAME: return GUID_ContainerFormatGif;
            default: return GUID_ContainerFormatPng;
        }
    }
    
    // Get file extension
    std::wstring GetFileExtension(ImageFormat format) {
        switch (format) {
            case ImageFormat::PNG: return L".png";
            case ImageFormat::JPEG: return L".jpg";
            case ImageFormat::BMP: return L".bmp";
            case ImageFormat::TIF: return L".tif";
            case ImageFormat::GIF:
            case ImageFormat::GIF_MULTIFRAME: return L".gif";
            case ImageFormat::DDS: return L".dds";
            default: return L".png";
        }
    }
    
    // Get DXGI format for DDS
    DXGI_FORMAT GetDXGIFormat(DDSCompression compression) {
        switch (compression) {
            case DDSCompression::BC1: return DXGI_FORMAT_BC1_UNORM;
            case DDSCompression::BC2: return DXGI_FORMAT_BC2_UNORM;
            case DDSCompression::BC3: return DXGI_FORMAT_BC3_UNORM;
            case DDSCompression::BC4: return DXGI_FORMAT_BC4_UNORM;
            case DDSCompression::BC5: return DXGI_FORMAT_BC5_UNORM;
            case DDSCompression::BC6H: return DXGI_FORMAT_BC6H_UF16;
            case DDSCompression::BC7: return DXGI_FORMAT_BC7_UNORM;
            case DDSCompression::RGBA: return DXGI_FORMAT_R8G8B8A8_UNORM;
            default: return DXGI_FORMAT_BC7_UNORM;
        }
    }
    
    // Helper function to convert RGBA to BGRA
    void ConvertRGBAToBGRA(const DirectX::Image* img, uint8_t* bgraData) {
        const uint8_t* src = img->pixels;
        for (size_t i = 0; i < img->width * img->height; ++i) {
            bgraData[i * 4 + 0] = src[i * 4 + 2]; // B
            bgraData[i * 4 + 1] = src[i * 4 + 1]; // G
            bgraData[i * 4 + 2] = src[i * 4 + 0]; // R
            bgraData[i * 4 + 3] = src[i * 4 + 3]; // A
        }
    }
    
    // Helper function to setup D3D11 and desktop duplication
    HRESULT SetupDesktopDuplication(
        Microsoft::WRL::ComPtr<ID3D11Device>& device,
        Microsoft::WRL::ComPtr<ID3D11DeviceContext>& context,
        Microsoft::WRL::ComPtr<IDXGIOutputDuplication>& duplication
    ) {
        // Create D3D11 device
        HRESULT hr = D3D11CreateDevice(
            nullptr,
            D3D_DRIVER_TYPE_HARDWARE,
            nullptr,
            0,
            nullptr,
            0,
            D3D11_SDK_VERSION,
            &device,
            nullptr,
            &context
        );
        if (FAILED(hr)) return hr;
        
        // Get DXGI device
        Microsoft::WRL::ComPtr<IDXGIDevice> dxgiDevice;
        hr = device.As(&dxgiDevice);
        if (FAILED(hr)) return hr;
        
        // Get DXGI adapter
        Microsoft::WRL::ComPtr<IDXGIAdapter> adapter;
        hr = dxgiDevice->GetAdapter(&adapter);
        if (FAILED(hr)) return hr;
        
        // Get DXGI output (monitor)
        Microsoft::WRL::ComPtr<IDXGIOutput> output;
        hr = adapter->EnumOutputs(0, &output);
        if (FAILED(hr)) return hr;
        
        // Get DXGI output1 for desktop duplication
        Microsoft::WRL::ComPtr<IDXGIOutput1> output1;
        hr = output.As(&output1);
        if (FAILED(hr)) return hr;
        
        // Create desktop duplication
        hr = output1->DuplicateOutput(device.Get(), &duplication);
        return hr;
    }
    
    // Helper function to capture a single frame
    HRESULT CaptureSingleFrame(
        ID3D11Device* device,
        ID3D11DeviceContext* context,
        IDXGIOutputDuplication* duplication,
        DirectX::ScratchImage& frameImage
    ) {
        // Capture frame
        DXGI_OUTDUPL_FRAME_INFO frameInfo;
        Microsoft::WRL::ComPtr<IDXGIResource> resource;
        
        HRESULT hr = duplication->AcquireNextFrame(1000, &frameInfo, &resource);
        if (FAILED(hr)) return hr;
        
        // Get texture from resource
        Microsoft::WRL::ComPtr<ID3D11Texture2D> texture;
        hr = resource.As(&texture);
        if (FAILED(hr)) {
            duplication->ReleaseFrame();
            return hr;
        }
        
        // Get texture description
        D3D11_TEXTURE2D_DESC desc;
        texture->GetDesc(&desc);
        
        // Create staging texture for CPU access
        desc.Usage = D3D11_USAGE_STAGING;
        desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
        desc.BindFlags = 0;
        desc.MiscFlags = 0;
        
        Microsoft::WRL::ComPtr<ID3D11Texture2D> stagingTexture;
        hr = device->CreateTexture2D(&desc, nullptr, &stagingTexture);
        if (FAILED(hr)) {
            duplication->ReleaseFrame();
            return hr;
        }
        
        // Copy texture to staging
        context->CopyResource(stagingTexture.Get(), texture.Get());
        
        // Map staging texture
        D3D11_MAPPED_SUBRESOURCE mapped;
        hr = context->Map(stagingTexture.Get(), 0, D3D11_MAP_READ, 0, &mapped);
        if (FAILED(hr)) {
            duplication->ReleaseFrame();
            return hr;
        }
        
        // Create DirectXTex image
        DirectX::Image image;
        image.width = desc.Width;
        image.height = desc.Height;
        image.format = desc.Format;
        image.rowPitch = mapped.RowPitch;
        image.slicePitch = mapped.RowPitch * desc.Height;
        image.pixels = static_cast<uint8_t*>(mapped.pData);
        
        // Initialize ScratchImage
        hr = frameImage.InitializeFromImage(image);
        if (FAILED(hr)) {
            context->Unmap(stagingTexture.Get(), 0);
            duplication->ReleaseFrame();
            return hr;
        }
        
        // Convert to RGBA if needed
        DirectX::ScratchImage convertedImage;
        if (desc.Format != DXGI_FORMAT_R8G8B8A8_UNORM) {
            hr = DirectX::Convert(
                frameImage.GetImages(), 
                frameImage.GetImageCount(), 
                frameImage.GetMetadata(),
                DXGI_FORMAT_R8G8B8A8_UNORM,
                DirectX::TEX_FILTER_DEFAULT,
                DirectX::TEX_THRESHOLD_DEFAULT,
                convertedImage
            );
            if (SUCCEEDED(hr)) {
                frameImage = std::move(convertedImage);
            }
        }
        
        // Unmap and release frame
        context->Unmap(stagingTexture.Get(), 0);
        duplication->ReleaseFrame();
        
        return hr;
    }
    
    // Multi-frame GIF capture function
    CaptureResult CaptureMultiFrameGIF(const CaptureParams& params) {
        CaptureResult result;
        
        try {
            logger::info("Starting multi-frame GIF capture for {} seconds", params.gifDuration);
            
            // Initialize COM
            HRESULT hr = CoInitializeEx(nullptr, COINIT_APARTMENTTHREADED);
            bool comInitialized = SUCCEEDED(hr);
            
            // Initialize WIC
            Microsoft::WRL::ComPtr<IWICImagingFactory> wicFactory;
            hr = CoCreateInstance(
                CLSID_WICImagingFactory,
                nullptr,
                CLSCTX_INPROC_SERVER,
                IID_PPV_ARGS(&wicFactory)
            );
            if (FAILED(hr)) {
                result.message = "Failed to create WIC factory";
                if (comInitialized) CoUninitialize();
                return result;
            }
            
            // Create GIF encoder
            Microsoft::WRL::ComPtr<IWICBitmapEncoder> encoder;
            hr = wicFactory->CreateEncoder(GUID_ContainerFormatGif, nullptr, &encoder);
            if (FAILED(hr)) {
                result.message = "Failed to create GIF encoder";
                if (comInitialized) CoUninitialize();
                return result;
            }
            
            // Generate filename
            auto now = std::chrono::system_clock::now();
            auto time_t_val = std::chrono::system_clock::to_time_t(now);
            auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(
                now.time_since_epoch()) % 1000;
            
            std::tm timeinfo{};
            localtime_s(&timeinfo, &time_t_val);
            
            std::wstringstream filename;
            filename << params.basePath;
            filename << L"_" << std::put_time(&timeinfo, L"%Y%m%d_%H%M%S");
            filename << L"_" << std::setfill(L'0') << std::setw(3) << ms.count();
            filename << L".gif";
            
            result.filepath = filename.str();
            
            // Create file stream
            Microsoft::WRL::ComPtr<IWICStream> stream;
            hr = wicFactory->CreateStream(&stream);
            if (FAILED(hr)) {
                result.message = "Failed to create stream";
                if (comInitialized) CoUninitialize();
                return result;
            }
            
            hr = stream->InitializeFromFilename(result.filepath.c_str(), GENERIC_WRITE);
            if (FAILED(hr)) {
                result.message = "Failed to initialize file stream";
                if (comInitialized) CoUninitialize();
                return result;
            }
            
            hr = encoder->Initialize(stream.Get(), WICBitmapEncoderNoCache);
            if (FAILED(hr)) {
                result.message = "Failed to initialize encoder";
                if (comInitialized) CoUninitialize();
                return result;
            }
            
            // Setup D3D11 for frame capture
            Microsoft::WRL::ComPtr<ID3D11Device> device;
            Microsoft::WRL::ComPtr<ID3D11DeviceContext> context;
            Microsoft::WRL::ComPtr<IDXGIOutputDuplication> duplication;
            
            hr = SetupDesktopDuplication(device, context, duplication);
            if (FAILED(hr)) {
                result.message = "Failed to setup desktop duplication";
                if (comInitialized) CoUninitialize();
                return result;
            }
            
            // Calculate frame timing
            const float frameRate = 10.0f; // 10 FPS for GIF
            const float frameInterval = 1.0f / frameRate;
            const int totalFrames = static_cast<int>(params.gifDuration * frameRate);
            const int delayTime = static_cast<int>(frameInterval * 100); // Delay in centiseconds
            
            logger::info("Capturing {} frames over {} seconds at {} FPS", 
                        totalFrames, params.gifDuration, frameRate);
            
            int successfulFrames = 0;
            
            // Capture frames in loop
            for (int frameIndex = 0; frameIndex < totalFrames; ++frameIndex) {
                // Check for cancellation
                if (params.cancelFlag && params.cancelFlag->load()) {
                    result.message = "Cancelled";
                    if (comInitialized) CoUninitialize();
                    return result;
                }
                
                // Capture single frame
                DirectX::ScratchImage frameImage;
                hr = CaptureSingleFrame(device.Get(), context.Get(), duplication.Get(), frameImage);
                if (FAILED(hr)) {
                    logger::warn("Failed to capture frame {}, continuing...", frameIndex);
                    // Small delay before retrying
                    std::this_thread::sleep_for(std::chrono::milliseconds(10));
                    continue;
                }
                
                // Convert frame to GIF-compatible format
                Microsoft::WRL::ComPtr<IWICBitmapFrameEncode> frame;
                Microsoft::WRL::ComPtr<IPropertyBag2> props;
                hr = encoder->CreateNewFrame(&frame, &props);
                if (FAILED(hr)) {
                    logger::warn("Failed to create frame {}", frameIndex);
                    continue;
                }
                
                // Set frame delay
                PROPBAG2 delayOption = {};
                delayOption.pstrName = const_cast<LPOLESTR>(L"FrameDelay");
                VARIANT delayValue;
                VariantInit(&delayValue);
                delayValue.vt = VT_UI2;
                delayValue.uiVal = delayTime;
                hr = props->Write(1, &delayOption, &delayValue);
                VariantClear(&delayValue);
                
                if (FAILED(hr)) {
                    logger::warn("Failed to set frame delay for frame {}", frameIndex);
                }
                
                hr = frame->Initialize(props.Get());
                if (FAILED(hr)) {
                    logger::warn("Failed to initialize frame {}", frameIndex);
                    continue;
                }
                
                // Write frame data
                const DirectX::Image* img = frameImage.GetImage(0, 0, 0);
                if (img) {
                    hr = frame->SetSize(static_cast<UINT>(img->width), static_cast<UINT>(img->height));
                    if (SUCCEEDED(hr)) {
                        WICPixelFormatGUID pixelFormat = GUID_WICPixelFormat32bppBGRA;
                        hr = frame->SetPixelFormat(&pixelFormat);
                        
                        if (SUCCEEDED(hr)) {
                            // Convert RGBA to BGRA and write
                            std::vector<uint8_t> bgraData(img->width * img->height * 4);
                            ConvertRGBAToBGRA(img, bgraData.data());
                            
                            hr = frame->WritePixels(
                                static_cast<UINT>(img->height),
                                static_cast<UINT>(img->width * 4),
                                static_cast<UINT>(bgraData.size()),
                                bgraData.data()
                            );
                            
                            if (SUCCEEDED(hr)) {
                                hr = frame->Commit();
                                if (SUCCEEDED(hr)) {
                                    successfulFrames++;
                                }
                            }
                        }
                    }
                }
                
                if (FAILED(hr)) {
                    logger::warn("Failed to write frame {} data", frameIndex);
                }
                
                // Wait for next frame (except for last frame)
                if (frameIndex < totalFrames - 1) {
                    std::this_thread::sleep_for(
                        std::chrono::milliseconds(static_cast<int>(frameInterval * 1000))
                    );
                }
                
                if (frameIndex % 10 == 0) { // Log progress every 10 frames
                    logger::debug("Captured frame {} of {} (successful: {})", 
                                frameIndex + 1, totalFrames, successfulFrames);
                }
            }
            
            // Commit encoder
            hr = encoder->Commit();
            if (FAILED(hr)) {
                result.message = "Failed to commit GIF encoder";
                if (comInitialized) CoUninitialize();
                return result;
            }
            
            if (successfulFrames > 0) {
                result.success = true;
                result.message = "Success";
                logger::info("Multi-frame GIF capture completed: {} successful frames", successfulFrames);
            } else {
                result.message = "No frames were successfully captured";
                logger::error("Multi-frame GIF capture failed: no successful frames");
            }
            
            if (comInitialized) {
                CoUninitialize();
            }
            
        } catch (const std::exception& e) {
            result.message = std::string("Exception: ") + e.what();
            logger::error("Multi-frame GIF capture exception: {}", e.what());
        }
        
        return result;
    }
    
    // Main capture function
    CaptureResult CaptureScreen(const CaptureParams& params) {
        // Handle multi-frame GIF separately
        if (params.format == ImageFormat::GIF_MULTIFRAME) {
            return CaptureMultiFrameGIF(params);
        }
        
        CaptureResult result;
        
        try {
            // Initialize COM
            HRESULT hr = CoInitializeEx(nullptr, COINIT_APARTMENTTHREADED);
            bool comInitialized = SUCCEEDED(hr);
            
            // Setup D3D11 and desktop duplication
            Microsoft::WRL::ComPtr<ID3D11Device> device;
            Microsoft::WRL::ComPtr<ID3D11DeviceContext> context;
            Microsoft::WRL::ComPtr<IDXGIOutputDuplication> duplication;
            
            hr = SetupDesktopDuplication(device, context, duplication);
            if (FAILED(hr)) {
                result.message = "Failed to setup desktop duplication";
                if (comInitialized) CoUninitialize();
                return result;
            }
            
            // Capture single frame
            DirectX::ScratchImage frameImage;
            hr = CaptureSingleFrame(device.Get(), context.Get(), duplication.Get(), frameImage);
            if (FAILED(hr)) {
                result.message = "Failed to capture frame";
                if (comInitialized) CoUninitialize();
                return result;
            }
            
            // Convert to RGBA if needed (already handled in CaptureSingleFrame)
            DirectX::ScratchImage convertedImage;
            const DirectX::Image* img = frameImage.GetImage(0, 0, 0);
            if (!img) {
                result.message = "Failed to get image from frame";
                if (comInitialized) CoUninitialize();
                return result;
            }
            
            // Copy the image data to avoid issues with mapped memory
            hr = convertedImage.InitializeFromImage(*img);
            if (FAILED(hr)) {
                result.message = "Failed to copy image";
                if (comInitialized) CoUninitialize();
                return result;
            }
            
            // Generate filename with timestamp
            auto now = std::chrono::system_clock::now();
            auto time_t_val = std::chrono::system_clock::to_time_t(now);
            auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(
                now.time_since_epoch()) % 1000;
            
            // Use localtime_s for safety
            std::tm timeinfo{};
            localtime_s(&timeinfo, &time_t_val);
            
            std::wstringstream filename;
            filename << params.basePath;
            filename << L"_" << std::put_time(&timeinfo, L"%Y%m%d_%H%M%S");
            filename << L"_" << std::setfill(L'0') << std::setw(3) << ms.count();
            filename << GetFileExtension(params.format);
            
            result.filepath = filename.str();
            
            // Save image based on format
            if (params.format == ImageFormat::DDS) {
                // Save as DDS using DirectXTex
                DXGI_FORMAT targetFormat = GetDXGIFormat(params.ddsMode);
                
                DirectX::ScratchImage compressedImage;
                if (DirectX::IsCompressed(targetFormat)) {
                    hr = DirectX::Compress(
                        convertedImage.GetImages(),
                        convertedImage.GetImageCount(),
                        convertedImage.GetMetadata(),
                        targetFormat,
                        DirectX::TEX_COMPRESS_DEFAULT,
                        DirectX::TEX_THRESHOLD_DEFAULT,
                        compressedImage
                    );
                    if (FAILED(hr)) {
                        result.message = "Failed to compress DDS image";
                        if (comInitialized) CoUninitialize();
                        return result;
                    }
                    
                    hr = DirectX::SaveToDDSFile(
                        compressedImage.GetImages(),
                        compressedImage.GetImageCount(),
                        compressedImage.GetMetadata(),
                        DirectX::DDS_FLAGS_NONE,
                        result.filepath.c_str()
                    );
                } else {
                    hr = DirectX::SaveToDDSFile(
                        convertedImage.GetImages(),
                        convertedImage.GetImageCount(),
                        convertedImage.GetMetadata(),
                        DirectX::DDS_FLAGS_NONE,
                        result.filepath.c_str()
                    );
                }
                
                if (FAILED(hr)) {
                    result.message = "Failed to save DDS file";
                    if (comInitialized) CoUninitialize();
                    return result;
                }
            } else {
                // Save using WIC for other formats
                hr = SaveToWICFile(convertedImage, params, result.filepath);
                if (FAILED(hr)) {
                    result.message = "Failed to save WIC file";
                    if (comInitialized) CoUninitialize();
                    return result;
                }
            }
            
            result.success = true;
            result.message = "Success";
            
            if (comInitialized) {
                CoUninitialize();
            }
            
        } catch (const std::exception& e) {
            result.message = std::string("Exception: ") + e.what();
            logger::error("ScreenCapture exception: {}", e.what());
        } catch (...) {
            result.message = "Unknown exception occurred";
            logger::error("Unknown exception in ScreenCapture");
        }
        
        return result;
    }
    
    // Helper function to save using WIC
    HRESULT SaveToWICFile(const DirectX::ScratchImage& image, const CaptureParams& params, const std::wstring& filepath) {
        // Initialize WIC
        Microsoft::WRL::ComPtr<IWICImagingFactory> wicFactory;
        HRESULT hr = CoCreateInstance(
            CLSID_WICImagingFactory,
            nullptr,
            CLSCTX_INPROC_SERVER,
            IID_PPV_ARGS(&wicFactory)
        );
        if (FAILED(hr)) return hr;
        
        // Create encoder
        Microsoft::WRL::ComPtr<IWICBitmapEncoder> encoder;
        GUID containerFormat = GetWICCodec(params.format);
        hr = wicFactory->CreateEncoder(containerFormat, nullptr, &encoder);
        if (FAILED(hr)) return hr;
        
        // Create file stream
        Microsoft::WRL::ComPtr<IWICStream> stream;
        hr = wicFactory->CreateStream(&stream);
        if (FAILED(hr)) return hr;
        
        hr = stream->InitializeFromFilename(filepath.c_str(), GENERIC_WRITE);
        if (FAILED(hr)) return hr;
        
        hr = encoder->Initialize(stream.Get(), WICBitmapEncoderNoCache);
        if (FAILED(hr)) return hr;
        
        // Create frame
        Microsoft::WRL::ComPtr<IWICBitmapFrameEncode> frame;
        Microsoft::WRL::ComPtr<IPropertyBag2> props;
        hr = encoder->CreateNewFrame(&frame, &props);
        if (FAILED(hr)) return hr;
        
        // Set properties for JPEG quality or TIFF compression
        if (params.format == ImageFormat::JPEG) {
            PROPBAG2 option = {};
            option.pstrName = const_cast<LPOLESTR>(L"ImageQuality");
            VARIANT varValue;
            VariantInit(&varValue);
            varValue.vt = VT_R4;
            varValue.fltVal = params.jpegQuality / 100.0f;
            props->Write(1, &option, &varValue);
            VariantClear(&varValue);
        } else if (params.format == ImageFormat::TIF) {
            PROPBAG2 option = {};
            option.pstrName = const_cast<LPOLESTR>(L"TiffCompressionMethod");
            VARIANT varValue;
            VariantInit(&varValue);
            varValue.vt = VT_UI1;
            varValue.bVal = static_cast<BYTE>(params.tiffMode);
            props->Write(1, &option, &varValue);
            VariantClear(&varValue);
        }
        
        hr = frame->Initialize(props.Get());
        if (FAILED(hr)) return hr;
        
        // Get image data
        const DirectX::Image* img = image.GetImage(0, 0, 0);
        if (!img) return E_FAIL;
        
        // Set frame size and format
        hr = frame->SetSize(static_cast<UINT>(img->width), static_cast<UINT>(img->height));
        if (FAILED(hr)) return hr;
        
        WICPixelFormatGUID pixelFormat = GUID_WICPixelFormat32bppBGRA;
        hr = frame->SetPixelFormat(&pixelFormat);
        if (FAILED(hr)) return hr;
        
        // Convert RGBA to BGRA if needed
        if (img->format == DXGI_FORMAT_R8G8B8A8_UNORM) {
            // Convert RGBA to BGRA
            size_t imageSize = img->width * img->height * 4;
            std::vector<uint8_t> bgraData(imageSize);
            
            const uint8_t* src = img->pixels;
            uint8_t* dst = bgraData.data();
            
            for (size_t i = 0; i < img->width * img->height; ++i) {
                dst[i * 4 + 0] = src[i * 4 + 2]; // B
                dst[i * 4 + 1] = src[i * 4 + 1]; // G
                dst[i * 4 + 2] = src[i * 4 + 0]; // R
                dst[i * 4 + 3] = src[i * 4 + 3]; // A
            }
            
            // Write pixels
            hr = frame->WritePixels(
                static_cast<UINT>(img->height),
                static_cast<UINT>(img->rowPitch),
                static_cast<UINT>(imageSize),
                bgraData.data()
            );
        } else {
            // Write pixels directly
            hr = frame->WritePixels(
                static_cast<UINT>(img->height),
                static_cast<UINT>(img->rowPitch),
                static_cast<UINT>(img->slicePitch),
                img->pixels
            );
        }
        
        if (FAILED(hr)) return hr;
        
        hr = frame->Commit();
        if (FAILED(hr)) return hr;
        
        hr = encoder->Commit();
        return hr;
    }
}