#include "network_interface.hh"

#include "arp_message.hh"
#include "ethernet_frame.hh"

#include <iostream>

// Dummy implementation of a network interface
// Translates from {IP datagram, next hop address} to link-layer frame, and from link-layer frame to IP datagram

// For Lab 5, please replace with a real implementation that passes the
// automated checks run by `make check_lab5`.

// You will need to add private members to the class declaration in `network_interface.hh`

template <typename... Targs>
void DUMMY_CODE(Targs &&.../* unused */) {}

using namespace std;

//! \param[in] ethernet_address Ethernet (what ARP calls "hardware") address of the interface
//! \param[in] ip_address IP (what ARP calls "protocol") address of the interface
NetworkInterface::NetworkInterface(const EthernetAddress &ethernet_address, const Address &ip_address)
    : _ethernet_address(ethernet_address), _ip_address(ip_address) {
    cerr << "DEBUG: Network interface has Ethernet address " << to_string(_ethernet_address) << " and IP address "
         << ip_address.ip() << "\n";
}

//! \param[in] dgram the IPv4 datagram to be sent
//! \param[in] next_hop the IP address of the interface to send it to (typically a router or default gateway, but may also be another host if directly connected to the same network as the destination)
//! (Note: the Address type can be converted to a uint32_t (raw 32-bit IP address) with the Address::ipv4_numeric() method.)
void NetworkInterface::send_datagram(const InternetDatagram &dgram, const Address &next_hop) {
    // convert IP address of next hop to raw 32-bit representation (used in ARP header)
    const uint32_t next_hop_ip = next_hop.ipv4_numeric();
    auto const &itr_next_hop_eth = _arp_cache.find(next_hop_ip);
    if (itr_next_hop_eth != _arp_cache.end()) {
        // found
        auto const &next_hop_eth = itr_next_hop_eth->second;
        if (next_hop_eth.first.has_value()) {
            // a valid arp cache entry
            EthernetFrame frame;
            EthernetHeader hdr;
            hdr.dst = next_hop_eth.first.value();
            hdr.src = _ethernet_address;
            hdr.type = EthernetHeader::TYPE_IPv4;
            frame.header() = hdr;
            frame.payload() = dgram.serialize();
            _frames_out.push(frame);
        } else {
            // an arp has been sent for it within 5 seconds
            _suspended_datagrams.push_back(std::make_pair(dgram, next_hop));
        }
    } else {
        // not found and no arp has been sent for the ip
        _suspended_datagrams.push_back(std::make_pair(dgram, next_hop));
        ARPMessage arpmsg;
        arpmsg.opcode = ARPMessage::OPCODE_REQUEST;
        arpmsg.sender_ethernet_address = _ethernet_address;
        arpmsg.sender_ip_address = _ip_address.ipv4_numeric();
        arpmsg.target_ip_address = next_hop_ip;
        EthernetFrame frame;
        EthernetHeader hdr;
        hdr.dst = ETHERNET_BROADCAST;
        hdr.src = _ethernet_address;
        hdr.type = EthernetHeader::TYPE_ARP;
        frame.header() = hdr;
        frame.payload() = arpmsg.serialize();
        _frames_out.push(frame);
        _arp_cache[next_hop_ip] = std::make_pair(std::optional<EthernetAddress>{}, 5 * 1000);
        // allocate an entry for the ARP Req we sent
    }
}

//! \param[in] frame the incoming Ethernet frame
optional<InternetDatagram> NetworkInterface::recv_frame(const EthernetFrame &frame) {
    if (frame.header().dst != _ethernet_address && frame.header().dst != ETHERNET_BROADCAST)
        return {};
    if (frame.header().type == EthernetHeader::TYPE_IPv4) {
        auto payload = frame.payload();
        InternetDatagram dgram;
        auto parse_res = dgram.parse(payload);
        if (parse_res == ParseResult::NoError) {
            return dgram;
        }
    } else if (frame.header().type == EthernetHeader::TYPE_ARP) {
        if (frame.header().src == _ethernet_address) {
            cerr << "receive loopback arp req\n";
            return {};
        }
        ARPMessage arpmsg;
        auto parse_res = arpmsg.parse(frame.payload());
        if (parse_res != ParseResult::NoError) {
            return {};
        }
        if (arpmsg.opcode != ARPMessage::OPCODE_REPLY && arpmsg.opcode != ARPMessage::OPCODE_REQUEST) {
            cerr << "unknown ARP opcode\n";
            return {};
        }
        if (arpmsg.opcode == ARPMessage::OPCODE_REQUEST && arpmsg.target_ip_address == _ip_address.ipv4_numeric()) {
            ARPMessage arpreply;
            arpreply.opcode = ARPMessage::OPCODE_REPLY;
            arpreply.sender_ethernet_address = _ethernet_address;
            arpreply.sender_ip_address = _ip_address.ipv4_numeric();
            arpreply.target_ip_address = arpmsg.sender_ip_address;
            arpreply.target_ethernet_address = arpmsg.sender_ethernet_address;
            // i think this should be directly the requester's EA
            EthernetFrame framereply;
            EthernetHeader hdr;
            hdr.dst = arpmsg.sender_ethernet_address;
            hdr.src = _ethernet_address;
            hdr.type = EthernetHeader::TYPE_ARP;
            framereply.header() = hdr;
            framereply.payload() = arpreply.serialize();
            _frames_out.push(framereply);
        }
        // no matter what type of arp it is, learn from it
        _arp_cache[arpmsg.sender_ip_address] =
            std::make_pair(std::optional<EthernetAddress>{arpmsg.sender_ethernet_address}, 30 * 1000);
        auto itr_dgram = _suspended_datagrams.begin();

        while (itr_dgram != _suspended_datagrams.end()) {
            if (itr_dgram->second.ipv4_numeric() == arpmsg.sender_ip_address) {
                auto dgram = itr_dgram->first;
                auto next_hop = itr_dgram->second;
                itr_dgram = _suspended_datagrams.erase(itr_dgram);
                this->send_datagram(dgram, next_hop);
            } else {
                itr_dgram++;
            }
        }
    }
    return {};
}

//! \param[in] ms_since_last_tick the number of milliseconds since the last call to this method
void NetworkInterface::tick(const size_t ms_since_last_tick) {
    for (auto &arp_entry : _arp_cache) {
        if (arp_entry.second.second >= ms_since_last_tick)
            arp_entry.second.second -= ms_since_last_tick;
        else
            arp_entry.second.second = 0;
    }
    auto arp_entry = _arp_cache.begin();
    while (arp_entry != _arp_cache.end()) {
        if (arp_entry->second.second == 0) {
            arp_entry = _arp_cache.erase(arp_entry);
        } else {
            arp_entry++;
        }
    }
    // TODO: maybe drop packets that couldn't get ARP reply, though not required in this lab
}