#include "arp_message.hh"
#include "router.hh"
#include "util.hh"

#include <iostream>
#include <list>
#include <unordered_map>

using namespace std;

auto rd = get_random_generator();

EthernetAddress random_host_ethernet_address() {
    EthernetAddress addr;
    for (auto &byte : addr) {
        byte = rd();  // use a random local Ethernet address
    }
    addr.at(0) |= 0x02;  // "10" in last two binary digits marks a private Ethernet address
    addr.at(0) &= 0xfe;

    return addr;
}

EthernetAddress random_router_ethernet_address() {
    EthernetAddress addr;
    for (auto &byte : addr) {
        byte = rd();  // use a random local Ethernet address
    }
    addr.at(0) = 0x02;  // "10" in last two binary digits marks a private Ethernet address
    addr.at(1) = 0;
    addr.at(2) = 0;

    return addr;
}

uint32_t ip(const string &str) { return Address{str}.ipv4_numeric(); }

template <typename T>
void clear(T &queue1, T &queue2) {
    while (not queue1.empty()) {
        queue1.pop();
        queue2.pop();
    }
}

string summary(const EthernetFrame &frame) {
    string ret;
    ret += frame.header().to_string();
    switch (frame.header().type) {
        case EthernetHeader::TYPE_IPv4: {
            InternetDatagram dgram;
            if (dgram.parse(frame.payload()) == ParseResult::NoError) {
                ret += " " + dgram.header().summary();
                ret += " payload=\"" + string(dgram.payload().concatenate()) + "\"";
            } else {
                ret += " (bad IPv4)";
            }
        } break;
        case EthernetHeader::TYPE_ARP: {
            ARPMessage arp;
            if (arp.parse(frame.payload()) == ParseResult::NoError) {
                ret += " " + arp.to_string();
            } else {
                ret += " (bad ARP)";
            }
        }
        default:
            break;
    }
    return ret;
}

class Host {
    string _name;
    Address _my_address;
    AsyncNetworkInterface _interface;
    Address _next_hop;

    std::list<InternetDatagram> _expecting_to_receive{};

    bool expecting(const InternetDatagram &expected) const {
        for (const auto &x : _expecting_to_receive) {
            if (x.serialize().concatenate() == expected.serialize().concatenate()) {
                return true;
            }
        }
        return false;
    }

    void remove_expectation(const InternetDatagram &expected) {
        for (auto it = _expecting_to_receive.begin(); it != _expecting_to_receive.end(); ++it) {
            if (it->serialize().concatenate() == expected.serialize().concatenate()) {
                _expecting_to_receive.erase(it);
                return;
            }
        }
    }

  public:
    Host(const string &name, const Address &my_address, const Address &next_hop)
        : _name(name)
        , _my_address(my_address)
        , _interface(random_host_ethernet_address(), _my_address)
        , _next_hop(next_hop) {}

    InternetDatagram send_to(const Address &destination, const uint8_t ttl = 64) {
        InternetDatagram dgram;
        dgram.header().src = _my_address.ipv4_numeric();
        dgram.header().dst = destination.ipv4_numeric();
        dgram.payload() = "random payload: {" + to_string(rd()) + "}";
        dgram.header().len = dgram.header().hlen * 4 + dgram.payload().size();
        dgram.header().ttl = ttl;

        _interface.send_datagram(dgram, _next_hop);

        cerr << "Host " << _name << " trying to send datagram (with next hop = " << _next_hop.ip()
             << "): " << dgram.header().summary() << " payload=\"" << dgram.payload().concatenate() << "\"\n";

        return dgram;
    }

    const Address &address() { return _my_address; }

    AsyncNetworkInterface &interface() { return _interface; }

    void expect(const InternetDatagram &expected) { _expecting_to_receive.push_back(expected); }

    const string &name() { return _name; }

    void check() {
        while (not _interface.datagrams_out().empty()) {
            const auto &dgram_received = _interface.datagrams_out().front();
            if (not expecting(dgram_received)) {
                throw runtime_error("Host " + _name +
                                    " received unexpected Internet datagram: " + dgram_received.header().summary() +
                                    " payload=\"" + dgram_received.payload().concatenate() + "\"");
            }
            remove_expectation(dgram_received);
            _interface.datagrams_out().pop();
        }

        if (not _expecting_to_receive.empty()) {
            auto &expected = _expecting_to_receive.front();
            throw runtime_error("Host " + _name + " did NOT receive an expected Internet datagram: " +
                                expected.header().summary() + " payload=\"" + expected.payload().concatenate() + "\"");
        }
    }
};

class Network {
  private:
    Router _router{};

    size_t default_id, eth0_id, eth1_id, eth2_id, uun3_id, hs4_id, mit5_id;

    std::unordered_map<string, Host> _hosts{};

    void exchange_frames(const string &x_name,
                         AsyncNetworkInterface &x,
                         const string &y_name,
                         AsyncNetworkInterface &y) {
        auto x_frames = x.frames_out(), y_frames = y.frames_out();

        deliver(x_name, x_frames, y_name, y);
        deliver(y_name, y_frames, x_name, x);

        clear(x_frames, x.frames_out());
        clear(y_frames, y.frames_out());
    }

    void exchange_frames(const string &x_name,
                         AsyncNetworkInterface &x,
                         const string &y_name,
                         AsyncNetworkInterface &y,
                         const string &z_name,
                         AsyncNetworkInterface &z) {
        auto x_frames = x.frames_out(), y_frames = y.frames_out(), z_frames = z.frames_out();

        deliver(x_name, x_frames, y_name, y);
        deliver(x_name, x_frames, z_name, z);

        deliver(y_name, y_frames, x_name, x);
        deliver(y_name, y_frames, z_name, z);

        deliver(z_name, z_frames, x_name, x);
        deliver(z_name, z_frames, y_name, y);

        clear(x_frames, x.frames_out());
        clear(y_frames, y.frames_out());
        clear(z_frames, z.frames_out());
    }

    void deliver(const string &src_name,
                 const queue<EthernetFrame> &src,
                 const string &dst_name,
                 AsyncNetworkInterface &dst) {
        queue<EthernetFrame> to_send = src;
        while (not to_send.empty()) {
            to_send.front().payload() = to_send.front().payload().concatenate();
            cerr << "Transferring frame from " << src_name << " to " << dst_name << ": " << summary(to_send.front())
                 << "\n";
            dst.recv_frame(move(to_send.front()));
            to_send.pop();
        }
    }

  public:
    Network()
        : default_id(_router.add_interface({random_router_ethernet_address(), {"171.67.76.46"}}))
        , eth0_id(_router.add_interface({random_router_ethernet_address(), {"10.0.0.1"}}))
        , eth1_id(_router.add_interface({random_router_ethernet_address(), {"172.16.0.1"}}))
        , eth2_id(_router.add_interface({random_router_ethernet_address(), {"192.168.0.1"}}))
        , uun3_id(_router.add_interface({random_router_ethernet_address(), {"198.178.229.1"}}))
        , hs4_id(_router.add_interface({random_router_ethernet_address(), {"143.195.0.2"}}))
        , mit5_id(_router.add_interface({random_router_ethernet_address(), {"128.30.76.255"}})) {
        _hosts.insert({"applesauce", {"applesauce", {"10.0.0.2"}, {"10.0.0.1"}}});
        _hosts.insert({"default_router", {"default_router", {"171.67.76.1"}, {"0"}}});
        ;
        _hosts.insert({"cherrypie", {"cherrypie", {"192.168.0.2"}, {"192.168.0.1"}}});
        _hosts.insert({"hs_router", {"hs_router", {"143.195.0.1"}, {"0"}}});
        _hosts.insert({"dm42", {"dm42", {"198.178.229.42"}, {"198.178.229.1"}}});
        _hosts.insert({"dm43", {"dm43", {"198.178.229.43"}, {"198.178.229.1"}}});

        _router.add_route(ip("0.0.0.0"), 0, host("default_router").address(), default_id);
        _router.add_route(ip("10.0.0.0"), 8, {}, eth0_id);
        _router.add_route(ip("172.16.0.0"), 16, {}, eth1_id);
        _router.add_route(ip("192.168.0.0"), 24, {}, eth2_id);
        _router.add_route(ip("198.178.229.0"), 24, {}, uun3_id);
        _router.add_route(ip("143.195.0.0"), 17, host("hs_router").address(), hs4_id);
        _router.add_route(ip("143.195.128.0"), 18, host("hs_router").address(), hs4_id);
        _router.add_route(ip("143.195.192.0"), 19, host("hs_router").address(), hs4_id);
        _router.add_route(ip("128.30.76.255"), 16, Address{"128.30.0.1"}, mit5_id);
    }

    void simulate_physical_connections() {
        exchange_frames(
            "router.default", _router.interface(default_id), "default_router", host("default_router").interface());
        exchange_frames("router.eth0", _router.interface(eth0_id), "applesauce", host("applesauce").interface());
        exchange_frames("router.eth2", _router.interface(eth2_id), "cherrypie", host("cherrypie").interface());
        exchange_frames("router.hs4", _router.interface(hs4_id), "hs_router", host("hs_router").interface());
        exchange_frames("router.uun3",
                        _router.interface(uun3_id),
                        "dm42",
                        host("dm42").interface(),
                        "dm43",
                        host("dm43").interface());
    }

    void simulate() {
        for (unsigned int i = 0; i < 256; i++) {
            _router.route();
            simulate_physical_connections();
        }

        for (auto &host : _hosts) {
            host.second.check();
        }
    }

    Host &host(const string &name) {
        auto it = _hosts.find(name);
        if (it == _hosts.end()) {
            throw runtime_error("unknown host: " + name);
        }
        if (it->second.name() != name) {
            throw runtime_error("invalid host: " + name);
        }
        return it->second;
    }
};

void network_simulator() {
    const string green = "\033[32;1m", normal = "\033[m";

    cerr << green << "Constructing network." << normal << "\n";

    Network network;

    cout << green << "\n\nTesting traffic between two ordinary hosts (applesauce to cherrypie)..." << normal << "\n\n";
    {
        auto dgram_sent = network.host("applesauce").send_to(network.host("cherrypie").address());
        dgram_sent.header().ttl--;
        network.host("cherrypie").expect(dgram_sent);
        network.simulate();
    }

    cout << green << "\n\nTesting traffic between two ordinary hosts (cherrypie to applesauce)..." << normal << "\n\n";
    {
        auto dgram_sent = network.host("cherrypie").send_to(network.host("applesauce").address());
        dgram_sent.header().ttl--;
        network.host("applesauce").expect(dgram_sent);
        network.simulate();
    }

    cout << green << "\n\nSuccess! Testing applesauce sending to the Internet." << normal << "\n\n";
    {
        auto dgram_sent = network.host("applesauce").send_to({"1.2.3.4"});
        dgram_sent.header().ttl--;
        network.host("default_router").expect(dgram_sent);
        network.simulate();
    }

    cout << green << "\n\nSuccess! Testing sending to the HS network and Internet." << normal << "\n\n";
    {
        auto dgram_sent = network.host("applesauce").send_to({"143.195.131.17"});
        dgram_sent.header().ttl--;
        network.host("hs_router").expect(dgram_sent);
        network.simulate();

        dgram_sent = network.host("cherrypie").send_to({"143.195.193.52"});
        dgram_sent.header().ttl--;
        network.host("hs_router").expect(dgram_sent);
        network.simulate();

        dgram_sent = network.host("cherrypie").send_to({"143.195.223.255"});
        dgram_sent.header().ttl--;
        network.host("hs_router").expect(dgram_sent);
        network.simulate();

        dgram_sent = network.host("cherrypie").send_to({"143.195.224.0"});
        dgram_sent.header().ttl--;
        network.host("default_router").expect(dgram_sent);
        network.simulate();
    }

    cout << green << "\n\nSuccess! Testing two hosts on the same network (dm42 to dm43)..." << normal << "\n\n";
    {
        auto dgram_sent = network.host("dm42").send_to(network.host("dm43").address());
        dgram_sent.header().ttl--;
        network.host("dm43").expect(dgram_sent);
        network.simulate();
    }

    cout << green << "\n\nSuccess! Testing TTL expiration..." << normal << "\n\n";
    {
        auto dgram_sent = network.host("applesauce").send_to({"1.2.3.4"}, 1);
        network.simulate();

        dgram_sent = network.host("applesauce").send_to({"1.2.3.4"}, 0);
        network.simulate();
    }

    cout << "\n\n\033[32;1mCongratulations! All datagrams were routed successfully.\033[m\n";
}

int main() {
    try {
        network_simulator();
    } catch (const exception &e) {
        cerr << "\n\n\n";
        cerr << "\033[31;1mError: " << e.what() << "\033[m\n";
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}