#include "NetworkSender.h"
#include <iostream>
#include <algorithm>

struct TransportHeader {
    uint32_t frameId;
    uint16_t fragId;
    uint16_t totalFrags;
};

NetworkSender::NetworkSender() {
    WSADATA wsaData;
    WSAStartup(MAKEWORD(2, 2), &wsaData);
}

NetworkSender::~NetworkSender() {
    if (socket_ != INVALID_SOCKET) {
        closesocket(socket_);
    }
    WSACleanup();
}

bool NetworkSender::Initialize(const std::string& ip, int port) {
    socket_ = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (socket_ == INVALID_SOCKET) return false;

    // Set buffer size
    int sndBuf = 1024 * 1024; // 1MB
    setsockopt(socket_, SOL_SOCKET, SO_SNDBUF, (char*)&sndBuf, sizeof(sndBuf));

    destAddr_.sin_family = AF_INET;
    destAddr_.sin_port = htons(port);
    inet_pton(AF_INET, ip.c_str(), &destAddr_.sin_addr);

    return true;
}

bool NetworkSender::SendFrame(const std::vector<uint8_t>& data, uint64_t timestamp, int width, int height, bool isKeyFrame) {
    // 1. Serialize Frame Info
    PacketHeader header;
    header.timestamp = timestamp;
    header.width = width;
    header.height = height;
    header.frameType = isKeyFrame ? 0 : 1;
    header.dataSize = (uint32_t)data.size();

    std::vector<uint8_t> buffer;
    buffer.resize(sizeof(PacketHeader) + data.size());
    memcpy(buffer.data(), &header, sizeof(PacketHeader));
    memcpy(buffer.data() + sizeof(PacketHeader), data.data(), data.size());

    // 2. Fragment and Send
    const int MTU = 1400; // Safe MTU
    const int HEADER_SIZE = sizeof(TransportHeader);
    const int PAYLOAD_SIZE = MTU - HEADER_SIZE;
    
    size_t totalSize = buffer.size();
    size_t totalFrags = (totalSize + PAYLOAD_SIZE - 1) / PAYLOAD_SIZE;
    
    static uint32_t frameId = 0;
    frameId++;

    for (size_t i = 0; i < totalFrags; ++i) {
        TransportHeader transHeader;
        transHeader.frameId = frameId;
        transHeader.fragId = (uint16_t)i;
        transHeader.totalFrags = (uint16_t)totalFrags;

        size_t offset = i * PAYLOAD_SIZE;
        size_t chunkSize = std::min((size_t)PAYLOAD_SIZE, totalSize - offset);

        std::vector<uint8_t> packet;
        packet.resize(HEADER_SIZE + chunkSize);
        
        memcpy(packet.data(), &transHeader, HEADER_SIZE);
        memcpy(packet.data() + HEADER_SIZE, buffer.data() + offset, chunkSize);

        int sent = sendto(socket_, (const char*)packet.data(), (int)packet.size(), 0, (sockaddr*)&destAddr_, sizeof(destAddr_));
        if (sent < 0) {
            // std::cerr << "Send failed" << std::endl;
            // Continue sending other fragments anyway
        }
    }

    return true;
}