#pragma once

#include <iostream>
#include <string>
#include <thread>
#include <mutex>
#include <vector>
#include <winsock2.h>
#include <ws2tcpip.h>
#include <opencv2/opencv.hpp>

// Ensure Windows socket libraries are linked
#pragma comment(lib, "Ws2_32.lib")

class NetworkCamera {
public:
    NetworkCamera();
    ~NetworkCamera();

    // Establishes connection to the Raspberry Pi
    bool connectToDevice(const std::string& ipAddress, int port);
    
    // Disconnects and shuts down background threads cleanly
    void disconnect();

    // Call this after your settle delay to get the absolute freshest frame
    cv::Mat getLatestFrame();

    bool isConnected() const { return m_connected; }

private:
    // Background thread function that continuously drains the socket network buffer
    void captureLoop();

    // Network handles
    SOCKET m_connectSocket;
    bool m_connected;
    bool m_running;

    // Background streaming thread
    std::thread m_captureThread;

    // Double-buffering architecture
    cv::Mat m_bufferA;
    cv::Mat m_bufferB;
    cv::Mat* m_writePtr; // Where the network thread decodes the incoming image
    cv::Mat* m_readPtr;  // Where the main loop safely reads from
    
    std::mutex m_bufferMutex;
};

#include "NetworkCamera.h"

NetworkCamera::NetworkCamera() 
    : m_connectSocket(INVALID_SOCKET), m_connected(false), m_running(false) {
    // Point our buffers to our local matrices
    m_writePtr = &m_bufferA;
    m_readPtr = &m_bufferB;
}

NetworkCamera::~NetworkCamera() {
    disconnect();
}

bool NetworkCamera::connectToDevice(const std::string& ipAddress, int port) {
    // Initialize Winsock
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        std::cerr << "Winsock initialization failed.\n";
        return false;
    }

    // Create Socket
    m_connectSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (m_connectSocket == INVALID_SOCKET) {
        std::cerr << "Socket creation failed.\n";
        WSACleanup();
        return false;
    }

    // Setup hint structure
    sockaddr_in clientService;
    clientService.sin_family = AF_INET;
    clientService.sin_port = htons(port);
    inet_pton(AF_INET, ipAddress.c_str(), &clientService.sin_addr);

    // Connect to Raspberry Pi
    if (connect(m_connectSocket, (SOCKADDR*)&clientService, sizeof(clientService)) == SOCKET_ERROR) {
        std::cerr << "Failed to connect to Pi camera server.\n";
        closesocket(m_connectSocket);
        WSACleanup();
        return false;
    }

    m_connected = true;
    m_running = true;

    // Spin up the background network thread
    m_captureThread = std::thread(&NetworkCamera::captureLoop, this);
    return true;
}

void NetworkCamera::disconnect() {
    if (m_running) {
        m_running = false;
        
        // Close socket to unblock any pending recv() calls in the thread
        if (m_connectSocket != INVALID_SOCKET) {
            closesocket(m_connectSocket);
            m_connectSocket = INVALID_SOCKET;
        }

        if (m_captureThread.joinable()) {
            m_captureThread.join();
        }
        
        WSACleanup();
        m_connected = false;
    }
}

void NetworkCamera::captureLoop() {
    while (m_running) {
        uint32_t imageSize = 0;

        // 1. Read the 4-byte header telling us how big the upcoming image is
        int result = recv(m_connectSocket, reinterpret_cast<char*>(&imageSize), sizeof(imageSize), 0);
        if (result <= 0 || !m_running) break; 

        // Convert from network byte order to host byte order (just in case)
        imageSize = ntohl(imageSize); 

        // 2. Read the actual image payload bytes (looping until full size reached)
        std::vector<uchar> byteBuffer(imageSize);
        int totalBytesRead = 0;
        
        while (totalBytesRead < static_cast<int>(imageSize) && m_running) {
            int bytesLeft = imageSize - totalBytesRead;
            int received = recv(m_connectSocket, reinterpret_cast<char*>(byteBuffer.data() + totalBytesRead), bytesLeft, 0);
            if (received <= 0) break;
            totalBytesRead += received;
        }

        if (totalBytesRead < static_cast<int>(imageSize)) continue; // Fragmented or broken frame

        // 3. Decode the compressed JPEG bytes straight into our current "Write" buffer
        // This takes the heavy CPU compression-lifting out of your main logic thread.
        cv::imdecode(byteBuffer, cv::IMREAD_UNCHANGED, m_writePtr);

        if (!m_writePtr->empty()) {
            // 4. Instantly swap pointers. The new image becomes readable, 
            // and the old image memory space is recycled for the next incoming data.
            std::lock_guard<std::mutex> lock(m_bufferMutex);
            std::swap(m_writePtr, m_readPtr);
        }
    }
}

cv::Mat NetworkCamera::getLatestFrame() {
    std::lock_guard<std::mutex> lock(m_bufferMutex);
    // Returns a fast, shallow-copy header of the matrix. 
    // It is fully isolated and thread-safe because the capture loop will 
    // write to the *other* buffer pointer on its next iteration.
    return *m_readPtr;
}