#include <signal.h>
#include <chrono>
#include <condition_variable>
#include <cstring>
#include <fstream>
#include <iostream>
#include <mutex>
#include <sstream>
#include <thread>
#include <unordered_set>

#include "common.h"
#include "tundev.h"

using h804::debug;
using h804::error;
using h804::info;
using h804::notice;
using h804::warn;

using namespace std::literals::chrono_literals;

std::mutex main_mut;
std::condition_variable main_waiter;
bool do_exit = false;

// If the first byte of a frame is this, then it specifies a special control message
constexpr uint8_t H804_PROT_MARKER = 0x84;
// On receiving this packet, a machine should add the sender's MAC to its peer list
constexpr uint8_t H804_PROT_PEER_REQ = 0x01;
// Request peer list from a machine
constexpr uint8_t H804_PROT_SHARE_REQ = 0x03;
// Peer list response
constexpr uint8_t H804_PROT_SHARE_RSP = 0x04;

std::unordered_set<h804::tun::macaddr_t> peers;
std::mutex peers_mut;
#define PEERS_LOCK() std::lock_guard<std::mutex> guard(peers_mut)

struct __attribute__((packed)) h804_raw_mac_list {
    size_t num_macs;
    uint8_t macs[][6];
};

struct __attribute__((packed)) h804_prot {
    uint8_t marker;
    uint8_t type;
    union {
        h804_raw_mac_list share_list;
        uint8_t other;
    };
};

void send_peer_req(h804::tun::EthSocket& source, h804::tun::macaddr_t to) {
    std::array<uint8_t, h804::tun::TUN_MTU> buf{};
    h804_prot* prot = reinterpret_cast<h804_prot*>(buf.data());
    prot->marker = H804_PROT_MARKER;
    prot->type = H804_PROT_PEER_REQ;
    source.write(buf, 2, to);
}

void send_share_req(h804::tun::EthSocket& source, h804::tun::macaddr_t to) {
    info << "Sending share req to " + to.to_string();
    std::array<uint8_t, h804::tun::TUN_MTU> buf{};
    h804_prot* prot = reinterpret_cast<h804_prot*>(buf.data());
    prot->marker = H804_PROT_MARKER;
    prot->type = H804_PROT_SHARE_REQ;
    source.write(buf, 2, to);
}

void send_share_rsp(h804::tun::EthSocket& source, h804::tun::macaddr_t to) {
    PEERS_LOCK();
    std::array<uint8_t, h804::tun::TUN_MTU> buf{};
    h804_prot* prot = reinterpret_cast<h804_prot*>(buf.data());
    prot->marker = H804_PROT_MARKER;
    prot->type = H804_PROT_SHARE_RSP;
    prot->share_list.num_macs = peers.size();
    size_t i = 0;
    for (const auto& it : peers) {
        for (size_t j = 0; j < 6; j++) {
            prot->share_list.macs[i][j] = it[j];
        }
        i++;
    }
    source.write(buf, h804::tun::TUN_MTU, to);
}

void handle_share_rsp(const h804::tun::EthSocket& source, const h804_prot* prot, size_t size,
                      h804::tun::macaddr_t from) {
    PEERS_LOCK();
    size_t num_macs = prot->share_list.num_macs;
    if (num_macs * 6 + sizeof(size_t) + 2 * sizeof(uint8_t) > size) {
        warn << "Out of bounds num_macs in share rsp";
        return;
    }
    peers.insert(from);
    for (size_t i = 0; i < num_macs; i++) {
        peers.emplace(prot->share_list.macs[i]);
    }
    peers.erase(peers.find(source.get_source_mac()));
}

void handle_peer_req(h804::tun::EthSocket& source, h804::tun::macaddr_t from) {
    {
        PEERS_LOCK();
        peers.insert(from);
    }
    send_share_rsp(source, from);
}

bool handle_ll(const std::array<uint8_t, h804::tun::TUN_MTU>& buf, size_t size,
               h804::tun::macaddr_t from, h804::tun::EthSocket& source) {
    if (size < 2) {
        return false;
    }
    const h804_prot* prot = reinterpret_cast<const h804_prot*>(buf.data());
    if (prot->marker == H804_PROT_MARKER) {
        switch (prot->type) {
            case H804_PROT_PEER_REQ:
                handle_peer_req(source, from);
                break;
            case H804_PROT_SHARE_RSP:
                handle_share_rsp(source, prot, size, from);
                break;
            case H804_PROT_SHARE_REQ:
                send_share_rsp(source, from);
                break;
            default:
                break;
        }

        return true;
    }
    return false;
}

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

    try {
        std::ifstream infile("h804tun.conf");
        std::string interface;
        std::string ip;
        if (!std::getline(infile, interface)) {
            error << "Need interface in conf";
        }
        if (!std::getline(infile, ip)) {
            error << "Need local IP in conf";
        }

        info << "Starting tunnel on " + interface + " " + ip;
        h804::tun::TunnelDevice device{ip, "255.255.255.0"};
        h804::tun::EthSocket llsock{interface.c_str()};
        info << "Started";
        std::string cmd1("ip link show ");
        cmd1 += device.get_name();
        system(cmd1.c_str());
        cmd1 = "ip addr show ";
        cmd1 += device.get_name();
        system(cmd1.c_str());
        system("ip route");

        struct sigaction handler;
        handler.sa_handler = [](int) {
            // std::unique_lock<std::mutex> lk(main_mut);
            do_exit = true;
            std::cout << std::endl;
            // main_waiter.notify_all();
        };
        sigemptyset(&handler.sa_mask);
        handler.sa_flags = 0;
        // sigaction(SIGINT, &handler, NULL);

        auto tun_thread = [&device, &llsock]() {
            std::array<uint8_t, h804::tun::TUN_MTU> buf;
            while (!do_exit) {
                size_t len = device.read(buf);
                if (len > 0) {
                    PEERS_LOCK();
                    for (const auto& addr : peers) {
                        llsock.write(buf, len, addr);
                    }
                }
            }
        };

        auto eth_thread = [&device, &llsock]() {
            std::array<uint8_t, h804::tun::TUN_MTU> buf;
            h804::tun::macaddr_t addr;
            while (!do_exit) {
                size_t len = llsock.read(buf, addr);
                if (len > 0) {
                    if (handle_ll(buf, len, addr, llsock)) {
                        continue;
                    }

                    device.write(buf, len);
                }
            }
        };

        auto share_thread = [&llsock] {
            while (!do_exit) {
                std::this_thread::sleep_for(300s);
                {
                    PEERS_LOCK();
                    for (const auto& it : peers) {
                        send_share_req(llsock, it);
                    }
                }
            }
        };

        info << "Sending peer requests";
        h804::tun::macaddr_t bcast{"ff:ff:ff:ff:ff:ff"};
        for (size_t i = 0; i < 5; i++) {
            send_peer_req(llsock, bcast);
        }

        std::thread tun_thread_obj(tun_thread);
        std::thread eth_thread_obj(eth_thread);
        std::thread share_thread_obj(share_thread);

        tun_thread_obj.join();
        eth_thread_obj.join();
        share_thread_obj.join();

        // {
        // std::unique_lock<std::mutex> lk(main_mut);
        // main_waiter.wait(lk, [] { return do_exit; });
        // }

        info << "Exiting";
    } catch (std::exception& e) {
        error << "Caught exception:";
        error << e.what();
    }
    return 0;
}