#include "commands.hpp"
#include "config.hpp"
#include <iostream>
#include <string>
#include <vector>
#include <sstream>
#include <iomanip>
#include <cstdio>
#include <stdexcept>
extern "C" {
#include "monocypher.h"
}

// Utility functions
namespace {
    class HexError : public std::runtime_error {
    public:
        explicit HexError(const std::string& msg) : std::runtime_error(msg) {}
    };

    class FileError : public std::runtime_error {
    public:
        explicit FileError(const std::string& msg) : std::runtime_error(msg) {}
    };

    std::string toHex(const uint8_t* data, size_t len) {
        std::ostringstream oss;
        oss << std::hex;
        for (size_t i = 0; i < len; i++) {
            oss << std::setw(2) << std::setfill('0') << (int)data[i];
        }
        return oss.str();
    }

    std::vector<uint8_t> fromHex(const std::string& hex) {
        if (hex.size() % 2 != 0) {
            throw HexError("Hex string length must be even");
        }
        
        std::vector<uint8_t> data;
        data.reserve(hex.size()/2);
        
        for(size_t i = 0; i < hex.size(); i += 2) {
            try {
                uint8_t val = (uint8_t)std::stoi(hex.substr(i, 2), nullptr, 16);
                data.push_back(val);
            } catch (const std::invalid_argument&) {
                throw HexError("Invalid hex character in string");
            } catch (const std::out_of_range&) {
                throw HexError("Hex value out of range");
            }
        }
        return data;
    }

    void readRandomBytes(uint8_t* buffer, size_t size) {
        FILE* f = fopen("/dev/urandom", "rb");
        if (!f) {
            throw FileError("Failed to open /dev/urandom");
        }
        
        size_t bytesRead = fread(buffer, 1, size, f);
        fclose(f);
        
        if (bytesRead != size) {
            throw FileError("Failed to read enough random bytes");
        }
    }

    void hexStrToKey(const std::string &hex, uint8_t outKey[32]) {
        if (hex.size() != 64) {
            throw HexError("Key must be 64 hex characters (32 bytes)");
        }
        
        std::vector<uint8_t> buf = fromHex(hex);
        std::copy(buf.begin(), buf.end(), outKey);
    }
}

static const size_t NONCE_SIZE = 24;
static const size_t KEY_SIZE = 32;
static const size_t MAC_SIZE = 16;

static void handleNickCommand(const std::string &args, AppConfig &config) {
    std::istringstream iss(args);
    std::string sub;
    iss >> sub;
    if(sub == "set") {
        std::string name;
        std::getline(iss, name);
        if(!name.empty() && name[0] == ' ') name.erase(0, 1);
        if(name.empty()) {
            std::cout << CLR_RED << "[nick] Error: Nickname cannot be empty" << CLR_RESET << "\n";
            return;
        }
        config.nickname = name;
        std::cout << CLR_GREEN << "[nick] Nickname set to: " << config.nickname << CLR_RESET << "\n";
    } else if(sub == "generatekey") {
        try {
            readRandomBytes(config.sharedSecret, KEY_SIZE);
            config.haveSharedSecret = true;
            std::cout << CLR_GREEN << "[nick] 256-bit key generated: " << toHex(config.sharedSecret, 32) << CLR_RESET << "\n";
        } catch (const FileError& e) {
            std::cout << CLR_RED << "[nick] Error: " << e.what() << CLR_RESET << "\n";
        }
    } else {
        std::cout << CLR_RED << "[nick] Error: Unknown subcommand '" << sub << "'" << CLR_RESET << "\n";
    }
}

static void handleWebCommand(const std::string &args, AppConfig &config);
static void handleSoundCommand(const std::string &args, AppConfig &config);

static void handleMakeTea(const std::string& input, AppConfig& config) {
    std::istringstream iss(input);
    std::string plaintext;
    iss >> plaintext;
    if(plaintext.empty()) {
        std::cout << CLR_RED << "[makeTea] Error: Text to encrypt cannot be empty" << CLR_RESET << "\n";
        return;
    }

    std::string keyHex;
    iss >> keyHex;
    uint8_t localKey[32];
    bool useLocal = false;
    
    if(!keyHex.empty()) {
        try {
            hexStrToKey(keyHex, localKey);
            useLocal = true;
        } catch (const HexError& e) {
            std::cout << CLR_RED << "[makeTea] Error: " << e.what() << CLR_RESET << "\n";
            return;
        }
    }

    if(!config.haveSharedSecret && !useLocal) {
        try {
            readRandomBytes(config.sharedSecret, KEY_SIZE);
            config.haveSharedSecret = true;
            std::cout << CLR_YELLOW << "[makeTea] No key found, random generated: " << toHex(config.sharedSecret, 32) << CLR_RESET << "\n";
        } catch (const FileError& e) {
            std::cout << CLR_RED << "[makeTea] Error: " << e.what() << CLR_RESET << "\n";
            return;
        }
    }

    std::vector<uint8_t> nonce(NONCE_SIZE), mac(MAC_SIZE);
    std::vector<uint8_t> ciphertext(plaintext.size());
    
    try {
        readRandomBytes(nonce.data(), NONCE_SIZE);
    } catch (const FileError& e) {
        std::cout << CLR_RED << "[makeTea] Error: " << e.what() << CLR_RESET << "\n";
        return;
    }

    const uint8_t* usedKey = useLocal ? localKey : config.sharedSecret;
    crypto_aead_lock(
        ciphertext.data(),
        mac.data(),
        usedKey,
        nonce.data(),
        nullptr, 0,
        (const uint8_t*)plaintext.data(),
        plaintext.size()
    );

    std::vector<uint8_t> out;
    out.insert(out.end(), nonce.begin(), nonce.end());
    out.insert(out.end(), mac.begin(), mac.end());
    out.insert(out.end(), ciphertext.begin(), ciphertext.end());
    
    std::cout << CLR_GREEN << "[makeTea] keyUsed=" << toHex(usedKey, 32) << CLR_RESET << "\n";
    std::cout << CLR_GREEN << "[makeTea] encrypted: " << toHex(out.data(), out.size()) << CLR_RESET << "\n";
}

static void handleDrinkTea(const std::string& input, AppConfig& config) {
    std::istringstream iss(input);
    std::string hexIn;
    iss >> hexIn;
    if(hexIn.empty()) {
        std::cout << CLR_RED << "[drinkTea] Error: Encrypted data cannot be empty" << CLR_RESET << "\n";
        return;
    }

    std::string keyHex;
    iss >> keyHex;
    uint8_t localKey[32];
    bool useLocal = false;
    
    if(!keyHex.empty()) {
        try {
            hexStrToKey(keyHex, localKey);
            useLocal = true;
        } catch (const HexError& e) {
            std::cout << CLR_RED << "[drinkTea] Error: " << e.what() << CLR_RESET << "\n";
            return;
        }
    }

    if(!config.haveSharedSecret && !useLocal) {
        std::cout << CLR_RED << "[drinkTea] Error: No key available for decryption" << CLR_RESET << "\n";
        return;
    }

    std::vector<uint8_t> data;
    try {
        data = fromHex(hexIn);
    } catch (const HexError& e) {
        std::cout << CLR_RED << "[drinkTea] Error: " << e.what() << CLR_RESET << "\n";
        return;
    }

    if(data.size() < NONCE_SIZE + MAC_SIZE) {
        std::cout << CLR_RED << "[drinkTea] Error: Input data too short" << CLR_RESET << "\n";
        return;
    }

    std::vector<uint8_t> nonce(data.begin(), data.begin() + NONCE_SIZE);
    std::vector<uint8_t> mac(data.begin() + NONCE_SIZE, data.begin() + NONCE_SIZE + MAC_SIZE);
    std::vector<uint8_t> cipher(data.begin() + NONCE_SIZE + MAC_SIZE, data.end());
    std::vector<uint8_t> plain(cipher.size());

    const uint8_t* usedKey = useLocal ? localKey : config.sharedSecret;
    int rc = crypto_aead_unlock(
        plain.data(),
        mac.data(),
        usedKey,
        nonce.data(),
        nullptr, 0,
        cipher.data(),
        cipher.size()
    );

    if(rc != 0) {
        std::cout << CLR_RED << "[drinkTea] Error: MAC verification failed" << CLR_RESET << "\n";
        return;
    }

    std::string s((char*)plain.data(), plain.size());
    std::cout << CLR_GREEN << "[drinkTea] keyUsed=" << toHex(usedKey, 32) << CLR_RESET << "\n";
    std::cout << CLR_GREEN << "[drinkTea] decrypted: " << s << CLR_RESET << "\n";
}

void processCommand(const std::string& input, AppConfig& config) {
    if(input.empty()) {
        return;
    }

    try {
        if(input.rfind("nick ", 0) == 0) {
            handleNickCommand(input.substr(5), config);
        } else if(input.rfind("web ", 0) == 0) {
            handleWebCommand(input.substr(4), config);
        } else if(input.rfind("sound ", 0) == 0) {
            handleSoundCommand(input.substr(6), config);
        } else if(input.rfind("cerber maketea ", 0) == 0) {
            handleMakeTea(input.substr(15), config);
        } else if(input.rfind("cerber drinktea ", 0) == 0) {
            handleDrinkTea(input.substr(16), config);
        } else if(input == "exit") {
            std::cout << CLR_CYAN << "[cli] Exiting..." << CLR_RESET << "\n";
            exit(0);
        } else if(input == "help") {
            std::cout << CLR_CYAN << "Available commands:\n"
                      << "  nick set <usernick> - Set your nickname\n"
                      << "  nick generatekey - Generate a new encryption key\n"
                      << "  web start - Start the web server\n"
                      << "  web connect <type> <ip> - Connect to a web server\n"
                      << "  web stop - Stop the web server\n"
                      << "  sound find - Start listening for sound signals\n"
                      << "  sound lose - Stop listening for sound signals\n"
                      << "  cerber maketea <text> [hexKey] - Encrypt text\n"
                      << "  cerber drinktea <hex> [hexKey] - Decrypt text\n"
                      << "  help - Show this help message\n"
                      << "  exit - Exit the program\n"
                      << CLR_RESET;
        } else {
            std::cout << CLR_RED << "[cli] Error: Unknown command. Type 'help' for available commands." << CLR_RESET << "\n";
        }
    } catch (const std::exception& e) {
        std::cout << CLR_RED << "[cli] Error: " << e.what() << CLR_RESET << "\n";
    } catch (...) {
        std::cout << CLR_RED << "[cli] Error: Unknown error occurred" << CLR_RESET << "\n";
    }
}

static void handleWebCommand(const std::string &args, AppConfig &config) {
    std::istringstream iss(args);
    std::string cmd;
    iss >> cmd;
    
    if(cmd == "start") {
        if(config.webServerRunning) {
            std::cout << CLR_YELLOW << "[web] Warning: Web server is already running" << CLR_RESET << "\n";
            return;
        }
        extern void webServerStart(AppConfig&);
        try {
            webServerStart(config);
            config.webServerRunning = true;
            std::cout << CLR_GREEN << "[web] Server started successfully" << CLR_RESET << "\n";
        } catch (const std::exception& e) {
            std::cout << CLR_RED << "[web] Error: " << e.what() << CLR_RESET << "\n";
        }
    } else if(cmd == "connect") {
        std::string type, ip;
        iss >> type >> ip;
        
        if(type.empty() || ip.empty()) {
            std::cout << CLR_RED << "[web] Error: Type and IP address are required" << CLR_RESET << "\n";
            return;
        }
        
        extern void webServerConnect(AppConfig&, const std::string&, const std::string&);
        try {
            webServerConnect(config, type, ip);
            std::cout << CLR_GREEN << "[web] Connected to " << ip << " as " << type << CLR_RESET << "\n";
        } catch (const std::exception& e) {
            std::cout << CLR_RED << "[web] Error: " << e.what() << CLR_RESET << "\n";
        }
    } else if(cmd == "stop") {
        if(!config.webServerRunning) {
            std::cout << CLR_YELLOW << "[web] Warning: Web server is not running" << CLR_RESET << "\n";
            return;
        }
        extern void webServerStop(AppConfig&);
        try {
            webServerStop(config);
            config.webServerRunning = false;
            std::cout << CLR_GREEN << "[web] Server stopped successfully" << CLR_RESET << "\n";
        } catch (const std::exception& e) {
            std::cout << CLR_RED << "[web] Error: " << e.what() << CLR_RESET << "\n";
        }
    } else {
        std::cout << CLR_RED << "[web] Error: Unknown command '" << cmd << "'" << CLR_RESET << "\n";
    }
}

static void handleSoundCommand(const std::string &args, AppConfig &config) {
    std::istringstream iss(args);
    std::string cmd;
    iss >> cmd;
    
    if(cmd == "find") {
        if(config.soundExchangeActive) {
            std::cout << CLR_YELLOW << "[sound] Warning: Sound exchange is already active" << CLR_RESET << "\n";
            return;
        }
        extern void soundFind(AppConfig&);
        try {
            soundFind(config);
            config.soundExchangeActive = true;
            std::cout << CLR_GREEN << "[sound] Started listening for sound signals" << CLR_RESET << "\n";
        } catch (const std::exception& e) {
            std::cout << CLR_RED << "[sound] Error: " << e.what() << CLR_RESET << "\n";
        }
    } else if(cmd == "lose") {
        if(!config.soundExchangeActive) {
            std::cout << CLR_YELLOW << "[sound] Warning: Sound exchange is not active" << CLR_RESET << "\n";
            return;
        }
        extern void soundLose(AppConfig&);
        try {
            soundLose(config);
            config.soundExchangeActive = false;
            std::cout << CLR_GREEN << "[sound] Stopped listening for sound signals" << CLR_RESET << "\n";
        } catch (const std::exception& e) {
            std::cout << CLR_RED << "[sound] Error: " << e.what() << CLR_RESET << "\n";
        }
    } else {
        std::cout << CLR_RED << "[sound] Error: Unknown command '" << cmd << "'" << CLR_RESET << "\n";
    }
}