#include "NetworkManager.hpp"
#include <iostream>
#include <cstring>

const std::vector<std::string> NetworkManager::AVAILABLE_COLORS = {
    "red", "green", "orange", "purple", "white"
};

NetworkManager::NetworkManager() : serverEndpoint(ip::make_address("127.0.0.1"), 9999) {
    socket.open(udp::v4());
    startReceive();
    ioThread = std::thread([this] { ioContext.run(); });
}

NetworkManager::~NetworkManager() {
    ioContext.stop();
    if (ioThread.joinable()) ioThread.join();
}

void NetworkManager::startReceive() {
    socket.async_receive_from(
        buffer(recvBuffer), serverEndpoint,
        [this](std::error_code ec, std::size_t bytes) {
            if (!ec && bytes > 0) {
                processMessage(recvBuffer.data(), bytes);
            }
            startReceive();
        }
    );
}

void NetworkManager::processMessage(const uint8_t* data, std::size_t size) {
    if (size == 0) return;

    auto msgType = static_cast<MessageType>(data[0]);

    switch (msgType) {
    case MessageType::Spawn:
        handleSpawn(data, size);
        break;
    case MessageType::Update:
        handleUpdate(data, size);
        break;
    case MessageType::PlayerJoined:
        handlePlayerJoined(data, size);
        break;
    case MessageType::PlayerLeft:
        handlePlayerLeft(data, size);
        break;
    case MessageType::PlayerList:
        handlePlayerList(data, size);
        break;
    case MessageType::ColorChanged:
        handleColorChanged(data, size);
        break;
    case MessageType::LoginResponse:
        handleLoginResponse(data, size);
        break;
    default:
        break;
    }
}

void NetworkManager::handleSpawn(const uint8_t* data, std::size_t size) {
    // Message format: [type(1)][id(4)][x(4)][y(4)][z(4)][colorLen(1)][color(colorLen)]
    if (size < 18) return;

    uint32_t id;
    float x, y, z;
    std::memcpy(&id, &data[1], sizeof(id));
    std::memcpy(&x, &data[5], sizeof(x));
    std::memcpy(&y, &data[9], sizeof(y));
    std::memcpy(&z, &data[13], sizeof(z));

    uint8_t colorLen = data[17];
    if (size >= 18 + colorLen) {
        playerColor = std::string(reinterpret_cast<const char*>(&data[18]), colorLen);
    }

    playerID = id;
    {
        std::lock_guard lock(positionMutex);
        serverPosition = {x, y, z};
    }
    connected = true;
    std::cout << "Connected as player " << id << " with color " << playerColor << "\n";
}

void NetworkManager::handleUpdate(const uint8_t* data, std::size_t size) {
    if (size < 17) return;

    uint32_t id;
    float x, y, z;
    std::memcpy(&id, &data[1], sizeof(id));
    std::memcpy(&x, &data[5], sizeof(x));
    std::memcpy(&y, &data[9], sizeof(y));
    std::memcpy(&z, &data[13], sizeof(z));

    if (id == playerID) {
        std::lock_guard lock(positionMutex);
        serverPosition = {x, y, z};
    } else {
        std::lock_guard lock(remotePlayersMutex);
        if (remotePlayers.find(id) != remotePlayers.end()) {
            remotePlayers[id].position = {x, y, z};
            remotePlayers[id].lastUpdate = GetTime();
        }
    }
}

void NetworkManager::handlePlayerJoined(const uint8_t* data, std::size_t size) {
    // Message format: [type(1)][id(4)][x(4)][y(4)][z(4)][colorLen(1)][color(colorLen)]
    if (size < 18) return;

    uint32_t id;
    float x, y, z;
    std::memcpy(&id, &data[1], sizeof(id));
    std::memcpy(&x, &data[5], sizeof(x));
    std::memcpy(&y, &data[9], sizeof(y));
    std::memcpy(&z, &data[13], sizeof(z));

    uint8_t colorLen = data[17];
    std::string color = "red";
    if (size >= 18 + colorLen) {
        color = std::string(reinterpret_cast<const char*>(&data[18]), colorLen);
    }

    if (id != playerID) {
        std::lock_guard lock(remotePlayersMutex);
        remotePlayers[id] = {id, {x, y, z}, color, static_cast<float>(GetTime())};
        std::cout << "Player " << id << " joined with color " << color << "\n";
    }
}

void NetworkManager::handlePlayerLeft(const uint8_t* data, std::size_t size) {
    if (size < 5) return;

    uint32_t id;
    std::memcpy(&id, &data[1], sizeof(id));

    std::lock_guard lock(remotePlayersMutex);
    remotePlayers.erase(id);
    std::cout << "Player " << id << " left\n";
}

void NetworkManager::handlePlayerList(const uint8_t* data, std::size_t size) {
    if (size < 2) return;

    uint8_t count = data[1];
    size_t offset = 2;

    std::lock_guard lock(remotePlayersMutex);
    remotePlayers.clear();

    for (uint8_t i = 0; i < count && offset + 17 < size; i++) {
        uint32_t id;
        float x, y, z;
        std::memcpy(&id, &data[offset], sizeof(id));
        std::memcpy(&x, &data[offset + 4], sizeof(x));
        std::memcpy(&y, &data[offset + 8], sizeof(y));
        std::memcpy(&z, &data[offset + 12], sizeof(z));

        uint8_t colorLen = data[offset + 16];
        std::string color = "red";

        if (offset + 17 + colorLen <= size) {
            color = std::string(reinterpret_cast<const char*>(&data[offset + 17]), colorLen);
        }

        offset += 17 + colorLen;

        if (id != playerID) {
            remotePlayers[id] = {id, {x, y, z}, color, static_cast<float>(GetTime())};
        }
    }

    std::cout << "Received list of " << (int)count << " players\n";
}

void NetworkManager::sendLogin() {
    std::array<uint8_t, 1> msg{static_cast<uint8_t>(MessageType::Login)};
    socket.send_to(buffer(msg), serverEndpoint);
}

void NetworkManager::sendLoginWithUsername(const std::string& username) {
    std::cout << "Attempting login with username: " << username << "\n";
    loginErrorMsg = "";
    currentUsername = username;
    std::vector<uint8_t> msg(2 + username.size());
    msg[0] = static_cast<uint8_t>(MessageType::Login);
    msg[1] = static_cast<uint8_t>(username.size());
    std::memcpy(&msg[2], username.data(), username.size());
    socket.send_to(buffer(msg), serverEndpoint);
    std::cout << "Login packet sent for username: " << username << "\n";
}

void NetworkManager::sendLogout() {
    if (!connected) {
        std::cout << "Warning: sendLogout called but not connected\n";
        return;
    }
    
    std::cout << "Sending logout for player ID " << playerID << " (username: " << currentUsername << ")\n";
    
    std::array<uint8_t, 5> msg{};
    msg[0] = static_cast<uint8_t>(MessageType::Logout);
    uint32_t id = playerID;
    std::memcpy(&msg[1], &id, sizeof(id));
    socket.send_to(buffer(msg), serverEndpoint);
    
    std::cout << "Logout packet sent, resetting client state\n";
    
    // Reset state
    connected = false;
    playerID = 0;
    currentUsername = "";
    loginErrorMsg = "";  // Clear any error messages
    {
        std::lock_guard lock(remotePlayersMutex);
        remotePlayers.clear();
    }
}

void NetworkManager::sendMove(float dx, float dy, float dz) {
    if (!connected) return;

    std::array<uint8_t, 17> msg{};
    msg[0] = static_cast<uint8_t>(MessageType::Move);

    uint32_t id = playerID;
    std::memcpy(&msg[1], &id, sizeof(id));
    std::memcpy(&msg[5], &dx, sizeof(dx));
    std::memcpy(&msg[9], &dy, sizeof(dy));
    std::memcpy(&msg[13], &dz, sizeof(dz));

    socket.send_to(buffer(msg), serverEndpoint);
}

void NetworkManager::sendColorChange(const std::string& newColor) {
    if (!connected) return;

    std::vector<uint8_t> msg(6 + newColor.size());
    msg[0] = static_cast<uint8_t>(MessageType::ChangeColor);

    uint32_t id = playerID;
    std::memcpy(&msg[1], &id, sizeof(id));
    msg[5] = static_cast<uint8_t>(newColor.size());
    std::memcpy(&msg[6], newColor.data(), newColor.size());

    socket.send_to(buffer(msg), serverEndpoint);
}

void NetworkManager::sendHeartbeat() {
    if (!connected) return;
    
    std::array<uint8_t, 5> msg{};
    msg[0] = static_cast<uint8_t>(MessageType::Heartbeat);
    uint32_t id = playerID;
    std::memcpy(&msg[1], &id, sizeof(id));
    socket.send_to(buffer(msg), serverEndpoint);
}

Vector3 NetworkManager::getPosition() {
    std::lock_guard lock(positionMutex);
    return serverPosition;
}

void NetworkManager::handleColorChanged(const uint8_t* data, std::size_t size) {
    // Message format: [type(1)][id(4)][colorLen(1)][color(colorLen)]
    if (size < 6) return;

    uint32_t id;
    std::memcpy(&id, &data[1], sizeof(id));

    uint8_t colorLen = data[5];
    if (size >= 6 + colorLen) {
        std::string newColor(reinterpret_cast<const char*>(&data[6]), colorLen);

        if (id == playerID) {
            playerColor = newColor;
            std::cout << "Your color changed to " << newColor << "\n";
        } else {
            std::lock_guard lock(remotePlayersMutex);
            if (remotePlayers.find(id) != remotePlayers.end()) {
                remotePlayers[id].color = newColor;
                std::cout << "Player " << id << " changed color to " << newColor << "\n";
            }
        }
    }
}

std::unordered_map<uint32_t, RemotePlayer> NetworkManager::getRemotePlayers() {
    std::lock_guard lock(remotePlayersMutex);
    return remotePlayers;
}

void NetworkManager::handleLoginResponse(const uint8_t* data, std::size_t size) {
    // Message format: [type(1)][success(1)][messageLen(1)][message(messageLen)]
    if (size < 3) return;
    
    uint8_t success = data[1];
    uint8_t msgLen = data[2];
    
    if (size >= 3 + msgLen) {
        std::string message(reinterpret_cast<const char*>(&data[3]), msgLen);
        
        if (success == 0) {
            // Login failed
            loginErrorMsg = message;
            connected = false;
            std::cout << "Login failed: " << message << "\n";
        } else {
            // Login succeeded, wait for spawn message
            std::cout << "Login accepted, waiting for spawn...\n";
        }
    }
}