#include "net/ipv4.hpp"

#include <algorithm>
#include <memory>
#include <utility>

using namespace net;

IPHandler::IPHandler(IPProvider &backend, std::uint8_t protocol)
    : backend(backend), protocol(protocol)
{
    backend.handlers[protocol] = this;
}

IPHandler::~IPHandler()
{
    if (backend.handlers[protocol] == this) {
        backend.handlers[protocol] = nullptr;
    }
}

void
IPHandler::send(NetOrder<std::uint32_t> dstIP, span<const std::uint8_t> msg)
{
    backend.send(dstIP, protocol, msg);
}

IPProvider::IPProvider(EtherFrameProvider &backend, ARP &arp,
                       NetOrder<std::uint32_t> gatewayIP,
                       NetOrder<std::uint32_t> subnetMask)
    : EtherFrameHandler(backend, 0x800),
      handlers(), arp(arp), gatewayIP(gatewayIP), subnetMask(subnetMask),
      nextIdent(0U)
{
}

bool
IPProvider::onEtherFrameReceived(span<std::uint8_t> msg)
{
    if (msg.size() < IPv4Msg::size()) {
        return false;
    }

    IPv4Msg ipHdr(msg.data());
    const int headerLength = 4*(ipHdr.verAndHdrLen.BE & 0x0f);
    bool sendBack = false;

    if (ipHdr.dstIP == backend.getIPAddress()) {
        int length = std::min<std::ptrdiff_t>(ntoh(ipHdr.totalLength.BE),
                                              msg.size());

        if (IPHandler *handler = handlers[ipHdr.protocol.BE]) {
            span<std::uint8_t> payload = msg.subspan(headerLength,
                                                     length - headerLength);
            sendBack = handler->onIPReceived(ipHdr.srcIP, ipHdr.dstIP, payload);
        }
    }

    if (sendBack) {
        NetOrder<std::uint32_t> temp = ipHdr.dstIP;
        ipHdr.dstIP = ipHdr.srcIP;
        ipHdr.srcIP = temp;

        ipHdr.timeToLive = netOrder<std::uint8_t>(0x40);
        ipHdr.checksum = checksum(msg.first(headerLength));
    }

    return sendBack;
}

void
IPProvider::send(NetOrder<std::uint32_t> dstIP, std::uint8_t protocol,
                 span<const std::uint8_t> msg)
{
    std::unique_ptr<std::uint8_t[]> buffer {
        new std::uint8_t[IPv4Msg::size() + msg.size()]
    };

    IPv4Msg ipHdr(buffer.get());

    ipHdr.verAndHdrLen = netOrder<std::uint8_t>(0x40 | (IPv4Msg::size()/4));
    ipHdr.tos = netOrder<std::uint8_t>(0);
    ipHdr.totalLength = toNetOrder<std::uint16_t>(msg.size() + IPv4Msg::size());
    // TODO: `ident` suppossed to be random.
    ipHdr.ident = toNetOrder<std::uint16_t>(nextIdent++);
    ipHdr.flagsAndOffset = toNetOrder<std::uint16_t>(0x4000);
    ipHdr.timeToLive = netOrder<std::uint8_t>(0x40);
    ipHdr.protocol = netOrder<std::uint8_t>(protocol);

    ipHdr.dstIP = dstIP;
    ipHdr.srcIP = backend.getIPAddress();

    ipHdr.checksum = netOrder<std::uint16_t>(0);
    ipHdr.checksum = checksum({ buffer.get(), IPv4Msg::size() });

    std::uint8_t* databuffer = buffer.get() + IPv4Msg::size();
    for (int i = 0; i < msg.size(); ++i) {
        databuffer[i] = msg[i];
    }

    NetOrder<std::uint32_t> route = dstIP;
    if ((dstIP.BE & subnetMask.BE) != (ipHdr.srcIP.BE & subnetMask.BE)) {
        route = gatewayIP;
    }

    backend.send(arp.resolve(route), etherType,
                 { buffer.get(), IPv4Msg::size() + msg.size() });
}

NetOrder<std::uint16_t>
IPProvider::checksum(span<const std::uint8_t> data)
{
    std::uint32_t temp = 0U;

    for (int i = 0U; i + 1 < data.size(); i += 2) {
        temp += (static_cast<std::uint16_t>(data[i]) << 8) | data[i + 1];
    }

    if (data.size() % 2) {
        temp += static_cast<uint16_t>(data[data.size() - 1U]) << 8;
    }

    while (temp & 0xffff0000) {
        temp = (temp & 0xffff) + (temp >> 16);
    }

    return toNetOrder<std::uint16_t>(~temp);
}