#include "Context.hpp"

#ifdef _WINDOWS
#define WIN32_LEAN_AND_MEAN
#include <WinSock2.h>
#include <Windows.h>
#include <assert.h>
#include <iphlpapi.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <ws2tcpip.h>

#ifndef SOCKET_TYPE
#define SOCKET_TYPE SOCKET
#endif
#ifndef SOCKET_CLOSE
#define SOCKET_CLOSE(x) closesocket(x)
#endif
#ifndef WSA_LAST_ERROR
#define WSA_LAST_ERROR(x) WSAGetLastError()
#endif
#ifndef socklen_t
#define socklen_t int
#endif

#else /* Linux */

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/types.h>

#ifndef __USE_XOPEN2K
#define __USE_XOPEN2K
#endif
#ifndef __USE_POSIX
#define __USE_POSIX
#endif
#include <arpa/inet.h>
#include <errno.h>
#include <netdb.h>
#include <netinet/in.h>
#include <sys/select.h>

#ifndef SOCKET_TYPE
#define SOCKET_TYPE int
#endif
#ifndef INVALID_SOCKET
#define INVALID_SOCKET -1
#endif
#ifndef SOCKET_CLOSE
#define SOCKET_CLOSE(x) close(x)
#endif
#ifndef WSA_LAST_ERROR
#define WSA_LAST_ERROR(x) ((long)(x))
#endif
#endif

#include <picoquic/picoquic_internal.h>
#include <picoquic/picoquic_utils.h>
//FIXME: Windows compat
//MAYBE: templatise callback switch

using namespace net;

bool stream_ctx::IsUnidirId(uint64_t id) { return !IS_BIDIR_STREAM_ID(id); }
bool stream_ctx::IsServerId(uint64_t id) { return IsUnidirId(id) && !IS_CLIENT_STREAM_ID(id); }
bool stream_ctx::IsClientId(uint64_t id) { return IsUnidirId(id) && IS_CLIENT_STREAM_ID(id); }

Context::Context(uint32_t max_connections, char const *cert, char const *key,
    picoquic_stream_data_cb_fn callback, void* callback_ctx,
    char const *tickets, char const *tokens):
    ticket_store_filename(tickets), token_store_filename(tokens)
{
    if (max_connections > 0) {
        uint64_t current_time = picoquic_current_time();
        quic = picoquic_create(max_connections, cert, key, NULL, ALPN, callback, callback_ctx,
            NULL, NULL, NULL, current_time, NULL, ticket_store_filename, NULL, 0);
        if (quic == NULL) {
            FATAL("Network context creation failed");
        }

        if (token_store_filename != nullptr) {
            if (picoquic_load_retry_tokens(quic, token_store_filename) != 0) {
                LOG_W("No token file present. Will create one as " << token_store_filename);
            }
        }

        picoquic_set_default_congestion_algorithm(quic, picoquic_bbr_algorithm);
        //MAYBE: picoquic_set_cookie_mode(quic, 2);

#if TRACE
        picoquic_set_key_log_file_from_env(quic);
        picoquic_set_log_level(quic, 1);
        //MAYBE: picoquic_set_textlog
#endif
        /*
        FIXME: esni from ct.address
        //MAYBE: multi universe
        if (esni_key_file_name != NULL && esni_rr_file_name != NULL) {
            ret = picoquic_esni_load_key(qserver, esni_key_file_name);
            if (ret == 0) {
                ret = picoquic_esni_server_setup(qserver, esni_rr_file_name);
            }
        }
        */

    } else {
        quic = nullptr;
        LOG_D("Local only context");
    }
}
Context::~Context() {
    for (int i = 0; i < nb_sockets; i++) {
        if (sockets[i] != INVALID_SOCKET) {
            SOCKET_CLOSE(sockets[i]);
            sockets[i] = INVALID_SOCKET;
        }
    }

    if (quic != NULL) {
        if (ticket_store_filename != NULL) {
            if (picoquic_save_session_tickets(quic, ticket_store_filename) != 0) {
                LOG_E("Could not store the saved session tickets.");
            }
        }
        if (token_store_filename != NULL) {
            if (picoquic_save_retry_tokens(quic, token_store_filename) != 0) {
                LOG_E("Could not save tokens to " << token_store_filename);
            }
        }
        picoquic_free(quic);
    }
}

void Context::openSockets(int port, int family) {
    assert(quic != nullptr);
#ifdef _WINDOWS
    WSADATA wsaData = {0};
    (void)WSA_START(MAKEWORD(2, 2), &wsaData);
#endif
    memset(sockets_family, 0, sizeof(sockets_family));

    {
        nb_sockets = (family == AF_UNSPEC) ? 2 : 1;

        /* Compute how many sockets are necessary */
        if (family == AF_UNSPEC) {
            sockets_family[0] = AF_INET;
            sockets_family[1] = AF_INET6;
        } else if (family == AF_INET || family == AF_INET6) {
            sockets_family[0] = family;
        } else {
            FATAL("Could not open socket. Unsupported AF " << family);
        }

        for (int i = 0; i < nb_sockets; i++) {
            int recv_set = 0;
            int send_set = 0;

            if ((sockets[i] = socket(sockets_family[i], SOCK_DGRAM, IPPROTO_UDP)) == INVALID_SOCKET ||
                picoquic_socket_set_ecn_options(sockets[i], sockets_family[i], &recv_set, &send_set) != 0 ||
                picoquic_socket_set_pkt_info(sockets[i], sockets_family[i]) != 0 ||
                (port != 0 && picoquic_bind_to_port(sockets[i], sockets_family[i], port) != 0)) {
                for (int j = 0; j < i; j++) {
                    if (sockets[i] != INVALID_SOCKET) {
                        SOCKET_CLOSE(sockets[i]);
                        sockets[i] = INVALID_SOCKET;
                    }
                }
                FATAL("Cannot setup socket. AF " << sockets_family[i] << ", port " << port);
            }
        }
    }
}

void Context::pull(const uint64_t MAX_DELAY, const uint64_t MAX_LOOPS, const uint64_t MAX_INNER_LOOPS) {
    if (quic == nullptr)
        return;

    const uint64_t start_time = picoquic_get_quic_time(quic);
    uint64_t current_time = start_time;
    uint64_t next_wake_time = current_time;
    uint64_t nb_loops = 0;

    struct sockaddr_storage addr_from;
    struct sockaddr_storage addr_to;

    int if_index_to;
    uint8_t buffer[PICOQUIC_MAX_PACKET_SIZE];
    uint8_t send_buffer[PICOQUIC_MAX_PACKET_SIZE];
    size_t send_length = 0;

    //MAYBE: port migration

    while (isRunning() && nb_loops < MAX_LOOPS && next_wake_time - start_time < MAX_DELAY) {
        /* Wait for packets */
        unsigned char received_ecn;
        int socket_rank = -1;
        if_index_to = 0;
        int64_t wake_delay = next_wake_time - current_time;
        int bytes_recv = picoquic_select_ex(sockets, nb_sockets,
            &addr_from, &addr_to,
            &if_index_to, &received_ecn,
            buffer, sizeof(buffer),
            wake_delay, &socket_rank, &current_time);

        nb_loops++;

        if (bytes_recv < 0) {
            LOG_E("Socket pull error");
            break;
        }

        uint64_t loop_time = current_time;
        uint16_t current_recv_port = socket_port;

        if (bytes_recv > 0) {
            /* track the local port value if not known yet */
            if (socket_port == 0 && nb_sockets == 1) {
                struct sockaddr_storage local_address;
                if (picoquic_get_local_address(sockets[0], &local_address) != 0) {
                    memset(&local_address, 0, sizeof(struct sockaddr_storage));
                    LOG_E("Could not read local address");
                } else if (addr_to.ss_family == AF_INET6) {
                    socket_port = ((struct sockaddr_in6*) & local_address)->sin6_port;
                } else if (addr_to.ss_family == AF_INET) {
                    socket_port = ((struct sockaddr_in*) & local_address)->sin_port;
                }
                current_recv_port = socket_port;
            }
            /* Document incoming port */
            if (addr_to.ss_family == AF_INET6) {
                ((struct sockaddr_in6*) & addr_to)->sin6_port = current_recv_port;
            } else if (addr_to.ss_family == AF_INET) {
                ((struct sockaddr_in*) & addr_to)->sin_port = current_recv_port;
            }
            /* Read packet */
            (void)picoquic_incoming_packet(quic, buffer,
                (size_t)bytes_recv, (struct sockaddr*) & addr_from,
                (struct sockaddr*) & addr_to, if_index_to, received_ecn,
                current_time);
        }

        /* Emit packets */
        uint64_t nb_inner_loops = 0;
        while (nb_inner_loops < MAX_INNER_LOOPS && loop_time - start_time < MAX_DELAY) {
            nb_inner_loops++;

            struct sockaddr_storage peer_addr;
            struct sockaddr_storage local_addr;
            int if_index = 0;
            int sock_err = 0;

            int ret = picoquic_prepare_next_packet(quic, loop_time,
                send_buffer, sizeof(send_buffer), &send_length,
                &peer_addr, &local_addr, &if_index, NULL, NULL);
            if (ret != 0 || send_length <= 0)
                break;

            SOCKET_TYPE send_socket = INVALID_SOCKET;
            for (int i = 0; i < nb_sockets; i++) {
                if (sockets_family[i] == peer_addr.ss_family) {
                    send_socket = sockets[i];
                    break;
                }
            }

            int sock_ret = picoquic_send_through_socket(send_socket,
                (struct sockaddr*) & peer_addr, (struct sockaddr*) & local_addr, if_index,
                (const char*)send_buffer, (int)send_length, &sock_err);
            if (sock_ret <= 0) {
                LOG_E("Cannot send data through socket");
            }
        }
        next_wake_time = picoquic_get_next_wake_time(quic, current_time);
    }
}

Connection::~Connection() {
    if (cnx != nullptr)
        release(0);
}
void Connection::setup(picoquic_quic_t* ctx, sockaddr* addr, const char* sni, picoquic_stream_data_cb_fn cb_fn, void *cb_ctx) {
    auto cnx = picoquic_create_cnx(ctx, picoquic_null_connection_id, picoquic_null_connection_id,
        addr, picoquic_current_time(), 0, sni, ALPN, is_client);
    if (cnx == NULL) {
        FATAL("Network connection creation failed");
    }
    setHandle(cnx);
    setCallback(cb_fn, cb_ctx);
    /* Client connection parameters could be set here, before starting the connection. */
    if (picoquic_start_client_cnx(cnx) < 0) {
        FATAL("Could not activate connection");
    }
}
void Connection::setHandle(picoquic_cnx_t* c) {
    assert(cnx == nullptr);
    cnx = c;
#if TRACE
    picoquic_connection_id_t icid = picoquic_get_initial_cnxid(cnx);
    printf("QUIC connection ID: ");
    for (uint8_t i = 0; i < icid.id_len; i++) {
        printf("%02x", icid.id[i]);
    }
    printf("\n");
#endif
}
void Connection::setCallback(picoquic_stream_data_cb_fn cb_fn, void *ctx) {
    if (cnx != nullptr)
        picoquic_set_callback(cnx, cb_fn, ctx);
}
void Connection::release(uint16_t reason) {
    if (cnx != nullptr) {
        picoquic_close(cnx, reason);
        cnx = nullptr;
    }
}

uint64_t Connection::getRTT() {
    return cnx != nullptr ? picoquic_get_rtt(cnx) : INT64_MAX;
}
uint16_t Connection::getErrorCode(bool is_app) {
    return cnx != nullptr ? (is_app ? cnx->remote_application_error : cnx->remote_error) : 42;
}
std::string Connection::getAddress() {
    auto str = std::string("?:?");
    if (cnx != nullptr) {
        sockaddr* addr = NULL;
        picoquic_get_peer_addr(cnx, &addr);
        str.resize(128, '\0');
        picoquic_addr_text(addr, str.data(), str.size());
    }
    return str;
}

void Connection::sendCopy(const uint8_t* ptr, size_t len, uint8_t queue, size_t queue_size) {
    auto data = (uint8_t*)malloc(len);
    assert(data != nullptr);
    memcpy(data, ptr, len);
    send(data::out_buffer(data, len), queue, queue_size);
}
void Connection::send(const data::out_view &view, uint8_t queue, size_t queue_size) {
    send(data::out_buffer(view), queue, queue_size);
}
void Connection::send(const data::out_buffer& buffer, uint8_t queue_id, size_t queue_size) {
    if (cnx == nullptr)
        return;

    assert(queue_id < outgoing.size());
    auto& queue = outgoing.at(queue_id);
    if (queue_size == 0 || GetSize(queue.streams) < queue_size) {
        auto stream = &queue.streams.emplace_front(nextOutgoingId(), queue_id, buffer);
        auto res = picoquic_mark_active_stream(cnx, stream->stream_id, 1, stream);
        if (res != 0) {
            LOG_E("Cannot initialize stream " << res);
        }
    } else {
        queue.pending = buffer;
    }
}

void Connection::emit(const uint8_t *ptr, size_t size) {
    assert(cnx != nullptr);
    auto res = picoquic_queue_datagram_frame(cnx, size, ptr);
    if (res != 0) {
        LOG_E("Cannot initialize datagram " << res);
    }
}

in_stream_ctx *Connection::receive(uint64_t streamId) {
    auto stream = &incoming.emplace_front(streamId);
    if (cnx != nullptr) {
        picoquic_set_app_stream_ctx(cnx, streamId, stream);
    }
    return stream;
}

void Connection::close(out_stream_ctx *str) {
    auto& queue = outgoing.at(str->queue_id);
    if (queue.pending.handle) {
        send(queue.pending, str->queue_id);
        queue.pending.handle = nullptr;
    }
    queue.streams.remove(*str);
}
void Connection::close(in_stream_ctx *str) {
    reset(str->stream_id);
    incoming.remove(*str);
}

void Connection::reset(uint64_t id) {
    if (cnx != nullptr) {
        picoquic_reset_stream(cnx, id, 0);
    }
}