#pragma once

#include <string>
#include <vector>
#include <thread>
#include <atomic>
#include <functional>
#include <mutex>
#include <winsock2.h>

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

    bool Start(int port);
    void Stop();
    
    // Send a file to the connected client
    bool SendFile(const std::string& filePath);
    bool SendFolder(const std::string& folderPath);

    // Set callback for received files
    // Callback signature: (filename, data) -> void
    // Note: For large files, we should probably stream to disk directly, 
    // but for simplicity in this demo, we might just notify path.
    using FileReceivedCallback = std::function<void(const std::string& filename)>;
    void SetFileReceivedCallback(FileReceivedCallback cb);

private:
    void AcceptLoop();
    void ClientHandler(SOCKET clientSocket);
    bool ReceiveBytes(SOCKET sock, void* buffer, int size);
    bool SendBytes(SOCKET sock, const void* buffer, int size);
    
    // File receiving logic
    void HandleFileHeader(SOCKET sock, const std::vector<uint8_t>& payload);
    void HandleFileData(const std::vector<uint8_t>& payload);
    void HandleFileEnd();
    void HandleFolderHeader(const std::vector<uint8_t>& payload);
    void HandleDirEntry(const std::vector<uint8_t>& payload);
    void HandleFileHeaderV2(const std::vector<uint8_t>& payload);

    SOCKET listenSocket_ = INVALID_SOCKET;
    SOCKET clientSocket_ = INVALID_SOCKET; // Only support 1 client for now
    std::thread acceptThread_;
    std::thread clientThread_;
    std::atomic<bool> running_ = false;

    FileReceivedCallback fileReceivedCallback_;
    
    // Current receiving state
    std::string currentFileName_;
    uint64_t currentFileSize_ = 0;
    uint64_t receivedBytes_ = 0;
    std::ofstream* currentFileStream_ = nullptr;
    std::mutex fileMutex_;
    std::string baseFolderRoot_;
};