#include "NetworkSender.h"
#include "ScreenCapture.h"
#include "VideoEncoder.h"
#include <iostream>
#include <thread>
#include <chrono>
#include <fstream>

int main(int argc, char* argv[]) {
    std::string ip = "127.0.0.1";
    int port = 8888;
    std::string outputFileName = "";

    if (argc > 1) ip = argv[1];
    if (argc > 2) port = std::stoi(argv[2]);
    if (argc > 3) outputFileName = argv[3];

    std::cout << "Starting Windows Sender Demo..." << std::endl;
    std::cout << "Target: " << ip << ":" << port << std::endl;
    if (!outputFileName.empty()) {
        std::cout << "Output File: " << outputFileName << std::endl;
    }

    // Debug: Open file to save H.264 stream if filename is provided
    std::ofstream outFile;
    if (!outputFileName.empty()) {
        outFile.open(outputFileName, std::ios::binary);
        if (outFile.is_open()) {
            std::cout << "Debug: Saving video stream to '" << outputFileName << "'" << std::endl;
        } else {
            std::cerr << "Warning: Failed to open output file '" << outputFileName << "'" << std::endl;
        }
    }

    ScreenCapture capture;
    if (!capture.Initialize()) {
        std::cerr << "Failed to initialize Screen Capture" << std::endl;
        return 1;
    }

    // Get screen size
    D3D11_TEXTURE2D_DESC desc;
    ComPtr<ID3D11Texture2D> frame;
    // Capture one frame to get size
    std::cout << "Waiting for first frame..." << std::endl;
    while (!capture.CaptureFrame(frame)) {
        std::this_thread::sleep_for(std::chrono::milliseconds(10));
    }
    frame->GetDesc(&desc);
    capture.ReleaseFrame();

    int width = desc.Width;
    int height = desc.Height;
    std::cout << "Screen Size: " << width << "x" << height << std::endl;

    VideoEncoder encoder;
    if (!encoder.Initialize(capture.GetDevice(), width, height, 60, 4000000)) { // 4Mbps
        std::cerr << "Failed to initialize Video Encoder" << std::endl;
        return 1;
    }

    NetworkSender sender;
    if (!sender.Initialize(ip, port)) {
        std::cerr << "Failed to initialize Network Sender" << std::endl;
        return 1;
    }

    std::cout << "Streaming started. Press Ctrl+C to stop." << std::endl;

    int frameCount = 0;
    auto lastTime = std::chrono::high_resolution_clock::now();

    while (true) {
        ComPtr<ID3D11Texture2D> texture;
        if (capture.CaptureFrame(texture)) {
            std::vector<uint8_t> encodedData;
            bool isKeyFrame = false;
            
            if (encoder.EncodeFrame(texture.Get(), encodedData, isKeyFrame)) {
                if (!encodedData.empty()) {
                    // Current timestamp in ms
                    auto now = std::chrono::high_resolution_clock::now();
                    uint64_t timestamp = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch()).count();

                    sender.SendFrame(encodedData, timestamp, width, height, isKeyFrame);
                    // std::cout << "Sent frame: " << encodedData.size() << " bytes, Key: " << isKeyFrame << std::endl;
                    
                    if (outFile.is_open()) {
                        outFile.write(reinterpret_cast<const char*>(encodedData.data()), encodedData.size());
                    }
                }
            }
            capture.ReleaseFrame();
            
            frameCount++;
        }

        auto now = std::chrono::high_resolution_clock::now();
        if (std::chrono::duration_cast<std::chrono::seconds>(now - lastTime).count() >= 1) {
            std::cout << "FPS: " << frameCount << std::endl;
            frameCount = 0;
            lastTime = now;
        }

        // Don't sleep too much, CaptureFrame waits. 
        // But if CaptureFrame returns immediately (high refresh rate), we might want to cap it?
        // Desktop Duplication limits itself to screen refresh rate usually.
    }

    return 0;
}