//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 <thread>
#include <mutex>
#include <csignal>

#include "../lib/error.h++"
#include "../lib/cli.h++"

#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::tcp_acceptor tcpAcceptor;
std::mutex tcpSocketMutex;
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();
    // tell sockpp to close TCP socket if open because it blocks when trying
    // to read and there is no data
    if (*tcpSocket) {
        // Intentionally not using the mutex here
        tcpSocket->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();
}

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();
}

int main(int argc, char* argv[]){

    signal(SIGINT, signalHandler);
    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"));

    std::vector<CLI::UnpositionalArgument> unpositionalArguments;
    unpositionalArguments.push_back(CLI::UnpositionalArgument('c', "connect", "HOST", "connect to HOST, listen for incoming connections if omitted"));
    unpositionalArguments.push_back(CLI::UnpositionalArgument('m', "mtu-optimize", "MTU", "Optimize for a specific maximum transfer unit by reading MTU bytes at a time."));

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

    CLI::ArgumentsParser cliParser = CLI::ArgumentsParser(argc, argv, flags, unpositionalArguments, positionalArguments);

    if (cliParser.wrongUsage) {
        //TODO: spit out usage information generated by the parser
        return EXIT_USAGE;
    }
    if (cliParser.getFlag('4').value || cliParser.getFlag('6').value) {
        ipv4 = cliParser.getFlag('4').value;
        ipv6 = cliParser.getFlag('6').value;
    }
    if (cliParser.getFlag('t').value || cliParser.getFlag('u').value) {
        tcp = cliParser.getFlag('t').value;
        udp = cliParser.getFlag('u').value;
    }
    if (cliParser.getUnpositionalArgument('c').errorCode == ErrorCodes::NOT_PRESENT) {
        listenMode = true;
    }
    if (cliParser.getUnpositionalArgument('m').errorCode == ErrorCodes::SUCCESS) {
        mtu = std::stol(cliParser.getUnpositionalArgument('m').value);
    }
    host = cliParser.getUnpositionalArgument('c').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.getPositionalArgument(0).value);

    if (listenMode) {
        if (ipv6) {
            std::cerr << "IPv6 support is not implented yet." << std::endl;
            return EXIT_UNIMPLEMENTED;
        }
        if (udp) {
            std::cerr << "UDP support is not implemented yet." << std::endl;
            return EXIT_UNIMPLEMENTED;
        }
        std::cerr << "Listening on port " << port << "." << std::endl;

        sockpp::socket_initializer socketInitializer;
        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;
        }

        readFromTCPSocket(tcpSocket, mtu);

        delete tcpSocket;
        return EXIT_SUCCESS;

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