//Copyright 2022, FOSS-VG Developers and Contributers
//
//This program is free software: you can redistribute it and/or modify it
//under the terms of the GNU Affero General Public License as published
//by the Free Software Foundation, version 3.
//
//This program is distributed in the hope that it will be useful,
//but WITHOUT ANY WARRANTY; without even the implied
//warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
//See the GNU Affero General Public License for more details.
//
//You should have received a copy of the GNU Affero General Public License
//version 3 along with this program.
//If not, see https://www.gnu.org/licenses/agpl-3.0.en.html

#include <iostream>
#include <iomanip>
#include <string>
#include <cstdint>
#include <cctype>
#include <sockpp/tcp_acceptor.h>
#include <sockpp/tcp6_acceptor.h>
#include <sockpp/udp_socket.h>
#include <sockpp/udp6_socket.h>
#include <thread>
#include <mutex>
#include <csignal>

#include "../lib/error.hpp"
#include "../lib/cli.hpp"

#define EXIT_SUCCESS 0
#define EXIT_RUNTIME 1
#define EXIT_USAGE 2
#define EXIT_UNIMPLEMENTED 3

bool ipv4 = true;
bool ipv6 = true;
bool tcp = true;
bool udp = true;
bool listenMode = false;
int64_t mtu = 1500;
std::string host;
in_port_t port;
sockpp::tcp_socket* tcpSocket;
sockpp::tcp6_socket* tcp6Socket;
sockpp::udp_socket* udpSocket;
sockpp::tcp_acceptor tcpAcceptor;
sockpp::tcp6_acceptor tcp6Acceptor;
std::mutex tcpSocketMutex;
std::mutex tcp6SocketMutex;
std::mutex udpSocketMutex;
std::mutex udp6SocketMutex;
std::mutex consoleMutex;
// used for coordinated graceful exit across threads
bool exitProgram = false;

void signalHandler(int signal) {
    exitProgram = true;
    // if still waiting for incoming connection, stop waiting
    tcpAcceptor.shutdown();
    tcp6Acceptor.shutdown();
    // tell sockpp to close TCP socket if open because it blocks when trying
    // to read and there is no data

    if (tcpSocket != nullptr && *tcpSocket) {
        // Intentionally not using the mutex here
        std::cout << "test\n";
        tcpSocket->shutdown(SHUT_RD);
    }

    if (tcp6Socket != nullptr && *tcp6Socket) {
        // Intentionally not using the mutex here
        tcp6Socket->shutdown(SHUT_RD);
    }
    //TODO: figure out if - and how - this applies to UDP

    // Priority is to finish up all unfinished business that can be finished up.
    // If something has the console mutex locked, that should not prevent
    // other threads from winding down. This is why logging happens last.
    consoleMutex.lock();
    std::cerr << "Received signal " << signal << ", shutting down." << std::endl;
    consoleMutex.unlock();
    std::exit(signal);
}

void readFromTCPSocket(sockpp::tcp_socket* socket, int64_t mtu) {
    ssize_t numBytes;
    uint8_t buffer[mtu];
    tcpSocketMutex.lock();
    while (!exitProgram && (numBytes = socket->read(buffer, sizeof(buffer))) > 0) {
        tcpSocketMutex.unlock();
        consoleMutex.lock();
        for (ssize_t i=0; i<numBytes; i++) {
            std::cout << std::hex << std::setfill('0') << std::setw(2) << (short) buffer[i];
        }
        std::cout.flush();
        consoleMutex.unlock();
        tcpSocketMutex.lock();
    }
    tcpSocketMutex.unlock();
    consoleMutex.lock();
    std::cerr << std::endl << "Connection closed." << std::endl;
    consoleMutex.unlock();
}

void readFromTCP6Socket(sockpp::tcp6_socket* socket, int64_t mtu){
    ssize_t numBytes;
    uint8_t buffer[mtu];
    tcp6SocketMutex.lock();
    while (!exitProgram && (numBytes = socket->read(buffer, sizeof(buffer))) > 0) {
        tcp6SocketMutex.unlock();
        consoleMutex.lock();
        for (ssize_t i=0; i<numBytes; i++) {
            std::cout << std::hex << std::setfill('0') << std::setw(2) << (short) buffer[i];
        }
        std::cout.flush();
        consoleMutex.unlock();
        tcp6SocketMutex.lock();
    }
    tcp6SocketMutex.unlock();
    consoleMutex.lock();
    std::cerr << std::endl << "Connection closed." << std::endl;
    consoleMutex.unlock();
}

void readFromUDPSocket(sockpp::udp_socket* socket, int64_t mtu){
    ssize_t numBytes;
    uint8_t buffer[mtu];
    typename sockpp::udp_socket::addr_t srcAddr;
    udpSocketMutex.lock();
    while(!exitProgram && (numBytes = socket->recv_from(buffer, sizeof(buffer), &srcAddr) > 0)){
        udpSocketMutex.unlock();
        consoleMutex.lock();
        for(ssize_t i=0; i<numBytes; i++){
            std::cout << std::hex << std::setfill('0') << std::setw(2) << (short) buffer[i];
        }
        std::cout.flush();
        consoleMutex.unlock();
        udpSocketMutex.lock();
    }
    udpSocketMutex.unlock();
    consoleMutex.lock();
    std::cerr << std::endl << "Connection closed." << std::endl;
    consoleMutex.unlock();
}  

void writeToTCPSocket(sockpp::tcp_socket* socket){
    while (!exitProgram) {
       //TODO: Implement locking/unlocking/threading shenanigans
    }
    consoleMutex.lock();
    std::cerr << std::endl << "Connection closed." << std::endl;
    consoleMutex.unlock();
}

void writeToTCP6Socket(sockpp::tcp6_socket* socket){
    while (!exitProgram) {
       //TODO: Implement locking/unlocking/threading shenanigans
    }
    consoleMutex.lock();
    std::cerr << std::endl << "Connection closed." << std::endl;
    consoleMutex.unlock();
}

int main(int argc, char* argv[]){
    sockpp::socket_initializer socketInitializer;

    std::signal(SIGINT, signalHandler);
    std::signal(SIGTERM, signalHandler);

    std::vector<CLI::Flag> flags;
    flags.push_back(CLI::Flag('4', "ipv4", "use IPv4, defaults to both when -4 and -6 are omitted, otherwise uses what is specified"));
    flags.push_back(CLI::Flag('6', "ipv6", "use IPv6, defaults to both when -4 and -6 are omitted, otherwise uses what is specified"));
    flags.push_back(CLI::Flag('t', "tcp", "use TCP, defaults to both when -t and -u are omitted, otherwise uses what is specified"));
    flags.push_back(CLI::Flag('u', "udp", "use UDP, defaults to both when -t and -u are omitted, otherwise uses what is specified"));
    flags.push_back(CLI::Flag('n', "no-color", "disable coloring the output (intended for terminals that don't work well with color escape sequences)"));
    flags.push_back(CLI::Flag('e', "echo-back", "echo input back to stdout"));
    flags.push_back(CLI::Flag('h', "help", "print this information and exit"));
    flags.push_back(CLI::Flag('l', "license", "print license information and exit"));

    std::vector<CLI::Option> options;
    options.push_back(CLI::Option('c', "connect", "HOST", "connect to HOST, listen for incoming connections if omitted"));
    options.push_back(CLI::Option('m', "mtu-optimize", "MTU", "Optimize for a specific maximum transfer unit by reading MTU bytes at a time."));
    options.push_back(CLI::Option(
        'p', "print-prefixes", "TCPin:UDPin:TCPout:UDPout",
        "override default prefixes for output (defaults to spaces + coloring the output or \"t:u:T:U\" in no-color mode)"
    ));
    options.push_back(CLI::Option(
        'i', "input-prefixes", "TCP:UDP",
        "override default prefixes for input (defaults to \"t:u\")"
    ));

    std::vector<CLI::Argument> arguments;
    arguments.push_back(CLI::Argument("PORT", "the port to use"));

    CLI::ArgumentsParser cliParser = CLI::ArgumentsParser(argc, argv, flags, options, arguments, "Arbitrary tcp/udp connections in hex format.");

    if (cliParser.getFlag("help").value){
        std::cout << cliParser.getUsage() << std::endl;
        return EXIT_SUCCESS;
    }
    if (cliParser.getFlag("license").value){
        std::cout
            << "Copyright 2022, FOSS-VG Developers and Contributers\n"
            << "\n"
            << "Hexnet is part of the FOSS-VG development tool suite.\n"
            << "\n"
            << "This program is free software: you can redistribute it and/or modify it\n"
            << "under the terms of the GNU Affero General Public License as published\n"
            << "by the Free Software Foundation, version 3.\n"
            << "\n"
            << "This program is distributed in the hope that it will be useful,\n"
            << "but WITHOUT ANY WARRANTY; without even the implied\n"
            << "warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.\n"
            << "See the GNU Affero General Public License for more details.\n"
            << "\n"
            << "You should have received a copy of the GNU Affero General Public License\n"
            << "version 3 along with this program.\n"
            << "If not, see https://www.gnu.org/licenses/agpl-3.0.en.html"
            << std::endl;
        return EXIT_SUCCESS;
    }

    if (cliParser.wrongUsage) {
        std::cout << cliParser.getUsage() << std::endl;
        return EXIT_USAGE;
    }
    if (cliParser.getFlag("ipv4").value || cliParser.getFlag("ipv6").value) {
        ipv4 = cliParser.getFlag("ipv4").value;
        ipv6 = cliParser.getFlag("ipv6").value;
    }
    if (cliParser.getFlag("tcp").value || cliParser.getFlag("udp").value) {
        tcp = cliParser.getFlag("tcp").value;
        udp = cliParser.getFlag("udp").value;
    }
    if (cliParser.getOption('c').errorCode == ErrorCodes::NOT_PRESENT) {
        listenMode = true;
    }
    if (cliParser.getOption("mtu-optimize").errorCode == ErrorCodes::SUCCESS) {
        mtu = std::stol(cliParser.getOption("mtu-optimize").value);
    }
    host = cliParser.getOption("connect").value;
    //FIXME: use a function that returns a fixed-width data type instead,
    //       ensure that the given value is a valid port
    port = (in_port_t) std::stoi(cliParser.getArgument(0).value);

    if (listenMode) {
        if (udp && ipv4) {
            std::cerr << "Listening on port " << port << "." << std::endl;
            
            if(!udpSocket->bind(sockpp::inet_address("localhost", port))){
                std::cerr << "Error while binding UDP socket: " << udpSocket->last_error_str() << std::endl;
                return EXIT_RUNTIME;
            }

            std::thread threadReadFromUDP = std::thread(readFromUDPSocket, udpSocket, mtu);
            threadReadFromUDP.join();

            delete udpSocket;
            return EXIT_SUCCESS;
        }

        if (ipv6) {
            std::cerr << "Listening on port " << port << "." << std::endl;
            
            tcp6Acceptor = sockpp::tcp6_acceptor(port);

            if(!tcp6Acceptor){
                std::cerr << "Error while creating TCP6 acceptor: " << tcp6Acceptor.last_error_str() << std::endl;
                return EXIT_RUNTIME;
            }

            sockpp::inet6_address peer;
            tcp6Socket = new sockpp::tcp6_socket();
            *tcp6Socket = tcp6Acceptor.accept(&peer);

            std::cerr << "Incoming connection from " << peer << std::endl;

            if (!(*tcp6Socket)) {
                std::cerr << "Error on incoming connection: " << tcp6Acceptor.last_error_str() << std::endl;
                delete tcp6Socket;
                return EXIT_RUNTIME;
            }

            std::thread threadReadFromTCP6 = std::thread(readFromTCP6Socket, tcp6Socket, mtu);
            threadReadFromTCP6.join();

            delete tcp6Socket;
            return EXIT_SUCCESS;
        }

        if(ipv4){
            std::cerr << "Listening on port " << port << "." << std::endl;

            
            tcpAcceptor = sockpp::tcp_acceptor(port);

            if (!tcpAcceptor) {
                std::cerr << "Error while creating TCP acceptor: " << tcpAcceptor.last_error_str() << std::endl;
                return EXIT_RUNTIME;
            }

            sockpp::inet_address peer;
            tcpSocket = new sockpp::tcp_socket();
            *tcpSocket = tcpAcceptor.accept(&peer);

            std::cerr << "Incoming connection from " << peer << std::endl;

            if (!(*tcpSocket)) {
                std::cerr << "Error on incoming connection: " << tcpAcceptor.last_error_str() << std::endl;
                delete tcpSocket;
                return EXIT_RUNTIME;
            }

            std::thread threadReadFromTCP = std::thread(readFromTCPSocket, tcpSocket, mtu);
            threadReadFromTCP.join();

            delete tcpSocket;
            return EXIT_SUCCESS;
        }

    } else {
        std::cerr << "Client mode is not implemented yet." << std::endl;
        return EXIT_UNIMPLEMENTED;
    }
}