#include "UdpVpn.hpp"

#include <chrono>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <string.h>
#include <poll.h>
#include <errno.h>
#include <math.h>

#include "ip_header.hpp"

static const size_t VPN_MTU = 1460; // TODO determine this -- issue #3

UdpVpn::UdpVpn()
    : _stopped(false), _dump_requested(false), _vpn_mtu(VPN_MTU),
    _tun_dev("cvpn%d"), _peer(nullptr)
{
    _last_tick = _last_control_sent =
        std::chrono::steady_clock::now() - std::chrono::seconds(1);
    _tun_dev.set_mtu(VpnPacket::get_tunnelled_mtu(_vpn_mtu));
    _socket = socket(AF_INET6, SOCK_DGRAM, 0);
    if(_socket < 0)
        throw UdpVpn::InitializationError("Cannot create socket", errno, true);
}

UdpVpn::~UdpVpn() {
    close(_socket);
}

void UdpVpn::run() {
    int rc;
    int start_at_fd = 0;  // read from polled fds in round-robin fashion
    int cur_fd;
    int nfds = 2;
    struct pollfd poll_fds[2];

    // poll_fds[0]: tun device
    poll_fds[0].fd = _tun_dev.get_fd();
    poll_fds[0].events = POLLIN;

    // poll_fds[1]: UDP socket device
    poll_fds[1].fd = _socket;
    poll_fds[1].events = POLLIN;

    while(!_stopped) {
        if(_dump_requested) {
            dump_state();
            _dump_requested = false;
        }

        rc = poll(poll_fds, nfds, 10); // timeout every 10ms

        if(rc < 0) {
            if(errno == EINTR)  // Interrupt.
                continue;
            throw UdpVpn::NetError(
                    "Error polling from interface", errno, true);
        }

        // ## Check periodic actions
        if(_peer) {
            if(std::chrono::steady_clock::now() - _last_control_sent
                > std::chrono::milliseconds(100))
            {
                if(_peer->send_control_packet())
                    _last_control_sent = std::chrono::steady_clock::now();
            }
        }
        if(std::chrono::steady_clock::now() - _last_tick
                > std::chrono::milliseconds(50))
        {
            _last_tick = std::chrono::steady_clock::now();
            if(_peer)
                _peer->tick();
        }

        if(rc == 0)  // Nothing to read -- timeout
            continue;

        cur_fd = start_at_fd;
        do {
            if(poll_fds[cur_fd].revents & POLLIN) {
                if(cur_fd == 0)
                    receive_from_tun();
                else if(cur_fd == 1)
                    receive_from_udp();
                break;
            }

            cur_fd = (cur_fd + 1) % nfds;
        } while(cur_fd != start_at_fd);

        start_at_fd = (start_at_fd + 1) % nfds;
    }
}

size_t UdpVpn::read_from_tun(char* buffer, size_t len) {
    // We know that there is data available -- use `read()`
    return _tun_dev.read(buffer, len);
}

size_t UdpVpn::read_from_tun(VpnDataPacket& packet) {
    size_t payload_space = packet.get_payload_space();
    size_t nread = read_from_tun(packet.get_payload(), payload_space);
    packet.set_payload_size(nread);
    if(!packet.parse_as_ipv6()) {
        debugf("Ignoring packet with invalid header\n");
        return 0;
    }
    if(nread != packet.get_ipv6_header().packet_length()) {
        debugf("Ignoring packet with bad size (expected %d, got %d, buffer %d)\n",
               packet.get_ipv6_header().packet_length(), nread, payload_space);
        return 0;
    }
    return nread;
}

size_t UdpVpn::read_from_udp(char* buffer, size_t len,
        sockaddr_in6& peer_addr)
{
    ssize_t nread;
    socklen_t peer_addr_len = sizeof(peer_addr);
    nread = recvfrom(_socket, buffer, len, 0,
            (struct sockaddr*) &peer_addr, &peer_addr_len);

    if(nread < 0)
        throw UdpVpn::NetError("Cannot receive datagram", errno, true);
    if(nread == 0)
        return 0;

    if(peer_addr.sin6_family != AF_INET6) {
        debugf("WARNING: Received non-ipv6 family datagram %d. Ignoring.\n",
                peer_addr.sin6_family);
        return 0;
    }
    if(peer_addr_len != sizeof(peer_addr)) {
        debugf("WARNING: received unexpected source address length %u."
                "Ignoring.\n",
                peer_addr_len);
        return 0;
    }

    return nread;
}

size_t UdpVpn::read_from_udp(VpnPacket& packet, sockaddr_in6& peer_addr) {
    packet.upon_reception(); // The packet is not read yet, but it has arrived
    size_t nread =
        read_from_udp(packet.get_data(), packet.get_data_space(), peer_addr);
    packet.set_data_size(nread);
    return nread;
}

size_t UdpVpn::transmit_to_peer(VpnPacket& packet) {
    if(!_peer) {
        debugf("Dropping packet: no peer yet.\n");
        return 0;
    }
    return _peer->write(packet);
}

void UdpVpn::receive_from_tun() {
    VpnDataPacket packet(_vpn_mtu, false);
    size_t nread = read_from_tun(packet);
    if(nread == 0)
        return;

    if(!_peer) {
        debugf("Dropping packet: no peer yet.\n");
        return;
    }
    packet.set_peer(_peer.get());

    kdebugf("Transmitting %s -> %s, size %d\n",
            format_address(packet.get_ipv6_header().source.s6_addr),
            format_address(packet.get_ipv6_header().dest.s6_addr),
            nread);

    packet.prepare_for_sending();
    transmit_to_peer(packet);
}

void UdpVpn::receive_from_udp() {
    VpnPacket packet(_vpn_mtu, true);
    sockaddr_in6 peer_ext_addr;
    size_t nread = read_from_udp(packet, peer_ext_addr);
    if(nread == 0)
        return;

    // If we don't have a peer yet -- we're just setting the peer to nullptr.
    packet.set_peer(_peer.get());

    if(packet.is_control()) {
        VpnControlPacket ctrl_packet(std::move(packet));

        for(VpnPacketTLV tlv=ctrl_packet.first_tlv();
                !tlv.past_the_end();
                tlv.seek_next_tlv())
        {
            switch(tlv.get_type()) {
                case VpnPacketTLV::PAYLOAD_TYPE_LOSS_REPORT:
                    if(_peer)
                        _peer->log_loss_report(VpnTlvLossReport(tlv));
                    break;
                case VpnPacketTLV::PAYLOAD_TYPE_RTTQ:
                    if(_peer)
                        _peer->make_rtta_for(VpnTlvRTTQ(tlv));
                    break;
                case VpnPacketTLV::PAYLOAD_TYPE_RTTA:
                    if(_peer)
                        _peer->log_rtta(VpnTlvRTTA(tlv));
                    break;
                case VpnPacketTLV::PAYLOAD_TYPE_UNDEF:
                default:
                    debugf("#%d+%lu: ignoring TLV with bad type %d.\n",
                            ctrl_packet.get_seqno(), tlv.get_offset(),
                            tlv.get_type());
                    break;
            }
        }
    } else {
        VpnDataPacket data_packet(std::move(packet));
        acquire_peer(data_packet, peer_ext_addr);
        receive_tunnelled_tlv(data_packet);
    }
}

void UdpVpn::receive_tunnelled_tlv(VpnDataPacket& packet) {
    if(!packet.parse_as_ipv6()) {
        debugf("Reinjection: dropping packet with bad IPv6 header.\n");
        return;
    }

    // Reinject into tun
    kdebugf("Reinjecting tunnelled packet of size %d [#%ld, TS=%ld μs]\n",
            packet.get_payload_size(),
            packet.get_seqno(),
            packet.get_sending_timestamp());
    _tun_dev.write(packet.get_payload(), packet.get_payload_size());
}

void UdpVpn::dump_state() const {
    if(!_peer) {
        printf("===== Cannot dump state, no peer yet =====\n");
        return;
    }

    const CongestionController& congest_ctrl =
        _peer->get_congestion_controller();

    printf("====== State dump ======\n");
    printf("Packet loss rate (inbound): %.0lf%%\n",
            round(_peer->get_loss_logger().get_loss_rate() * 100));
    printf("Packet loss rate (outbound): %.0lf%%\n",
            round(_peer->get_loss_reports().loss_rate() * 100));
    printf("RTT: %.02lf ms avg, %.02lf ms last [last updated: %lu ms ago]\n",
            (double)_peer->get_rtt().avg_rtt() / 1e3,
            (double)_peer->get_rtt().cur_rtt() / 1e3,
            std::chrono::duration_cast<std::chrono::milliseconds>(
                std::chrono::steady_clock::now()
                - _peer->get_rtt().get_last_update()).count()
          );
    printf("Total bytes transmitted: %s\n",
            human_readable_unit(_peer->get_tot_bytes_sent(), "B"));
    printf("Current outbound byte rate: %s\n",
            human_readable_unit(_peer->get_outbound_byte_rate(), "Bps"));
    printf("Available bandwidth:\n\t%s [loss based controller]\n",
            human_readable_unit(congest_ctrl.get_lossbased_bandwidth(), "B")
          );
    printf("==== End state dump ====\n");
}