#include "NetworkSender.h"
#include "ScreenCapture.h"
#include "IddBridge.h"
#include "VideoEncoder.h"
#include "TcpServer.h"
#include <iostream>
#include <sstream>
#include <vector>
#include <thread>
#include <chrono>
#include <fstream>
#include <windows.h>

int main(int argc, char* argv[]) {
    std::string ipStr = "127.0.0.1";
    int port = 8888;
    std::string outputFileName = "";
    std::string source = "screen";
    int outputIndex = -1;
    bool iddProducer = false;
    int producerOutputIndex = -1;
    std::string fileToSend = "";
    std::string folderToSend = "";

    if (argc > 1) ipStr = argv[1];
    if (argc > 2) port = std::stoi(argv[2]);
    if (argc > 3) outputFileName = argv[3];
    for (int i = 4; i < argc; ++i) {
        std::string arg = argv[i];
        if (arg.rfind("--source=", 0) == 0) {
            source = arg.substr(std::string("--source=").size());
        } else if (arg.rfind("--output=", 0) == 0) {
            try {
                outputIndex = std::stoi(arg.substr(std::string("--output=").size()));
            } catch (...) {}
        } else if (arg == "--idd-producer") {
            iddProducer = true;
        } else if (arg.rfind("--producer-output=", 0) == 0) {
            try {
                producerOutputIndex = std::stoi(arg.substr(std::string("--producer-output=").size()));
            } catch (...) {}
        } else if (arg.rfind("--send-file=", 0) == 0) {
            fileToSend = arg.substr(std::string("--send-file=").size());
        } else if (arg.rfind("--send-folder=", 0) == 0) {
            folderToSend = arg.substr(std::string("--send-folder=").size());
        }
    }

    // Parse IPs
    std::vector<std::string> ips;
    std::stringstream ss(ipStr);
    std::string item;
    while (std::getline(ss, item, ',')) {
        // Trim spaces might be good but let's assume no spaces for simplicity or user should handle
        if (!item.empty()) {
            ips.push_back(item);
        }
    }
    if (ips.empty()) ips.push_back("127.0.0.1");

    std::cout << "Starting Windows Sender Demo..." << std::endl;
    std::cout << "Targets: ";
    for (const auto& ip : ips) std::cout << ip << " ";
    std::cout << ", Port: " << port << std::endl;
    std::cout << "Source: " << source << std::endl;
    if (source == "screen" && outputIndex >= 0) {
        std::cout << "Output Index: " << outputIndex << std::endl;
    }
    if (!outputFileName.empty()) {
        std::cout << "Output File: " << outputFileName << std::endl;
    }
    if (source == "idd" && iddProducer) {
        std::cout << "IDD Producer enabled" << std::endl;
        if (producerOutputIndex >= 0) {
            std::cout << "Producer Output Index: " << producerOutputIndex << std::endl;
        }
    }

    // Start TCP Server
    TcpServer tcpServer;
    if (tcpServer.Start(8889)) {
        // Wait for client to connect if we need to send a file immediately
        if (!fileToSend.empty()) {
            std::cout << "Waiting for client to connect to send file: " << fileToSend << "..." << std::endl;
            // Simple polling wait (in a real app, use events/condition variables)
            // But TcpServer::SendFile checks if client is connected. 
            // We'll try to send in a loop or thread.
            std::thread([&tcpServer, fileToSend]() {
                // Wait for up to 30 seconds for a client
                for (int i = 0; i < 300; ++i) {
                    if (tcpServer.SendFile(fileToSend)) {
                        break;
                    }
                    std::this_thread::sleep_for(std::chrono::milliseconds(100));
                }
            }).detach();
        } else if (!folderToSend.empty()) {
            std::cout << "Waiting for client to connect to send folder: " << folderToSend << "..." << std::endl;
            std::thread([&tcpServer, folderToSend]() {
                for (int i = 0; i < 300; ++i) {
                    if (tcpServer.SendFolder(folderToSend)) {
                        break;
                    }
                    std::this_thread::sleep_for(std::chrono::milliseconds(100));
                }
            }).detach();
        }
    } else {
        std::cerr << "Failed to start TCP Server on port 8889" << 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;
        }
    }

    std::thread producerThread;
    if (source == "idd" && iddProducer) {
        producerThread = std::thread([producerOutputIndex]() {
            ScreenCapture cap;
            bool okCap = false;
            if (producerOutputIndex >= 0) okCap = cap.InitializeWithOutputIndex(producerOutputIndex);
            else okCap = cap.Initialize();
            if (!okCap) return;
            D3D11_TEXTURE2D_DESC desc = {};
            ComPtr<ID3D11Texture2D> frame;
            while (!cap.CaptureFrame(frame)) {
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
            }
            frame->GetDesc(&desc);
            cap.ReleaseFrame();
            D3D11_TEXTURE2D_DESC sdesc = {};
            sdesc.Width = desc.Width;
            sdesc.Height = desc.Height;
            sdesc.MipLevels = 1;
            sdesc.ArraySize = 1;
            sdesc.Format = DXGI_FORMAT_B8G8R8A8_UNORM;
            sdesc.SampleDesc.Count = 1;
            sdesc.Usage = D3D11_USAGE_STAGING;
            sdesc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
            ComPtr<ID3D11Texture2D> staging;
            cap.GetDevice()->CreateTexture2D(&sdesc, nullptr, &staging);

            IddProducer producer;
            if (!producer.Initialize()) return;

            while (true) {
                ComPtr<ID3D11Texture2D> tex;
                if (cap.CaptureFrame(tex)) {
                    cap.GetContext()->CopyResource(staging.Get(), tex.Get());
                    D3D11_MAPPED_SUBRESOURCE mapped = {};
                    if (SUCCEEDED(cap.GetContext()->Map(staging.Get(), 0, D3D11_MAP_READ, 0, &mapped))) {
                        producer.SubmitFrame(mapped.pData, desc.Width, desc.Height, (uint32_t)mapped.RowPitch);
                        cap.GetContext()->Unmap(staging.Get(), 0);
                    }
                    cap.ReleaseFrame();
                } else {
                    std::this_thread::sleep_for(std::chrono::milliseconds(5));
                }
            }
        });
    }
    bool ok = false;
    ScreenCapture capture;
    IddBridge idd;
    bool useIdd = (source == "idd");
    if (useIdd) {
        ok = idd.Initialize();
    } else {
        if (outputIndex >= 0) {
            ok = capture.InitializeWithOutputIndex(outputIndex);
        } else {
            ok = capture.Initialize();
        }
    }
    if (!ok) {
        std::cerr << "Failed to initialize capture source" << std::endl;
        return 1;
    }

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

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

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

    NetworkSender sender;
    if (!sender.Initialize(ips, 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;
        bool got = false;
        if (useIdd) {
            got = idd.CaptureFrame(texture);
        } else {
            got = capture.CaptureFrame(texture);
        }
        if (got) {
            D3D11_TEXTURE2D_DESC tdesc = {};
            texture->GetDesc(&tdesc);
            if (tdesc.Width != (UINT)width || tdesc.Height != (UINT)height) {
                std::cout << "Resolution changed: " << width << "x" << height << " -> " << tdesc.Width << "x" << tdesc.Height << std::endl;
                width = (int)tdesc.Width;
                height = (int)tdesc.Height;
                if (!encoder.Reinitialize(width, height, 60, 4000000)) {
                    std::cerr << "Failed to reinitialize encoder on resolution change" << std::endl;
                    break;
                }
            }
            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());
                    }
                }
            }
            if (useIdd) {
                idd.ReleaseFrame();
            } else {
                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;
}