#include "Connection.hpp"

#include <optional>
#include <algorithm>
#include <boost/asio.hpp>



namespace Connection {
asio::awaitable<Packet> Client::receivePacket() {
    Packet packet;
    // Receive operation
    packet.op = co_await receiveProtobuf<Generic::Operation>();
    // Receive data
    switch (packet.op.code()) {
    case Generic::Operation::Error: {
        packet.data = std::make_unique<Generic::Error>(co_await receiveProtobuf<Generic::Error>());
    } break;
    case Generic::Operation::SimpleAuth: {
        packet.data = std::make_unique<Generic::SimpleAuth>(co_await receiveProtobuf<Generic::SimpleAuth>(30));
    } break;
    case Generic::Operation::Disconnect:
    case Generic::Operation::OK:
        break;
    default: throw FatalError("Invalid opcode: "+std::to_string(packet.op.code()), Generic::Error::InvalidOpcode);
    }
}

asio::awaitable<void> Client::sendProtobuf(const google::protobuf::Message& buffer) {
    // Send size
    uint16_t size = buffer.ByteSizeLong();
    co_await asio::async_write(socket, asio::buffer(&size, sizeof(size)), asio::use_awaitable);
    // Send buffer
    std::vector<char> rawBuffer(size);
    buffer.SerializeToArray(rawBuffer.data(), rawBuffer.size());
    co_await asio::async_write(socket, asio::buffer(rawBuffer.data(), rawBuffer.size()), asio::use_awaitable);
}

asio::awaitable<void> Client::sendPacket(const Packet& packet) {
    // Send operation
    co_await sendProtobuf(packet.op);
    // Send data
    if (packet.data) {
        co_await sendProtobuf(*packet.data);
    }
}

asio::awaitable<void> Client::handlePacket(const Packet& packet) {
    // Check if authentication is required and completed
    if (packet.op.code() >= Generic::Operation::_MinAuth && !isAuthed()) {
        throw FatalError("Authentication required", Generic::Error::AuthenticationRequired);
    }
    // Handle the packet
    switch (packet.op.code()) {
    case Generic::Operation::Disconnect: {
        good = 0;
    } break;
    case Generic::Operation::SimpleAuth: {
        const Generic::SimpleAuth& authData = *static_cast<Generic::SimpleAuth*>(packet.data.get());
        // Check that username is not already taken
        if (std::find_if(clients.begin(), clients.end(), [&, this] (const auto& o) {return o->getUsername() == authData.username();}) != clients.end()) {
            throw Error("Username has already been taken", Generic::Error::UsernameAlreadyTaken);
        }
        // Check that username is "good" (no unprintable characters and stuff like that)
        //TODO...
        // Set username
        username = std::move(authData.username());
        // Send OK
        co_await sendOK();
    } break;
    default: {}
    }
    co_return;
}

asio::awaitable<void> Client::connect() {
    // Send initial "OK"
    try {
        co_await sendOK();
    } catch (...) {
        good = false;
    }
    // Run while connection is good
    while (good) {
        std::optional<Error> error;
        try {
            co_await handlePacket(co_await receivePacket());
        } catch (FatalError& e) {
            error = std::move(e);
            good = false;
        } catch (Error& e) {
            error = std::move(e);
        } catch (std::exception& e) {
            error = Error("Internal server error: "+std::string(e.what()), Generic::Error::Unknown);
            good = false;
        }
        if (error.has_value()) {
            co_await sendPacket(error.value().packet);
        }
    }
    // Try to send disconnect
    try {
        Packet packet;
        packet.op.set_code(Generic::Operation::Disconnect);
        co_await sendPacket(packet);
    } catch (...) {}
    // Close socket
    boost::system::error_code ec;
    socket.close(ec);
}

asio::awaitable<void> Server::listen(int port) {
    try {
        // Create acceptor
        asio::ip::tcp::acceptor acceptor(ioc, asio::ip::tcp::endpoint(asio::ip::tcp::v4(), port));
        // Handle all incoming connections
        while (true) {
            auto client = std::make_shared<Client>(ioc, clients);
            co_await acceptor.async_accept(client->getSocket(), asio::use_awaitable);
            asio::co_spawn(ioc, client->connect(), asio::detached);
            clients.push_back(std::move(client));
        }
    } catch (std::exception& e) {
        std::cerr << "Failed to stay alive: " << e.what() << std::endl;
    }
}
}