#include "Server.hpp"

using namespace net::server;

namespace net::server {
    int connection_callback(picoquic_cnx_t* cnx,
        uint64_t stream_id, uint8_t* bytes, size_t length,
        picoquic_call_back_event_t fin_or_event, void* callback_ctx, void* v_stream_ctx) {

        auto server = (Server*)picoquic_get_default_callback_context(picoquic_get_quic_ctx(cnx));
        auto peer = (Peer*)callback_ctx;
        if (fin_or_event != picoquic_callback_almost_ready) {
            if (peer == nullptr || peer == (void*)server) {
                assert(fin_or_event == picoquic_callback_ready);
                peer = server->connect(cnx);
                picoquic_set_callback(cnx, connection_callback, peer);
            }
            assert(peer == nullptr || peer->contains(cnx) || fin_or_event == picoquic_callback_close);
        } else if (peer == (void*)server) {
            peer = nullptr;
        }
        assert(v_stream_ctx == nullptr || ((net::stream_ctx*)v_stream_ctx)->stream_id == stream_id);

        return server->connectionCallback(peer, stream_id, bytes, length, fin_or_event, v_stream_ctx);
    }
}

Server::Server(const net::exposure& ct,
    std::function<std::optional<uint16_t>(Peer*)> onConnect,
    std::function<bool(Peer*, bool, uint16_t)> onDisconnect,
    std::function<bool(Peer*, const data::out_view&, net::PacketFlags)> onPacket):
    Context(ct.max_connections, ct.cert.c_str(), ct.key.c_str(), connection_callback, this),
    onConnect(onConnect), onDisconnect(onDisconnect), onPacket(onPacket), max_connections(ct.max_connections)
{
    if (ct.max_connections > 0) {
        openSockets(ct.port, 0);
        LOG_I("Listening on port " << ct.port);
    }
}
Server::~Server() {
    for(auto& peer: peers) {
        constexpr auto reason = (uint16_t)disconnect_reason::CLOSE;
        if (onDisconnect(&peer, true, reason))
            peer.release(reason);
    }
    peers.clear();
}

void Server::emitBroadcast(const uint8_t* ptr, size_t size) {
    for(auto& peer: peers) {
        peer.emit(ptr, size);
    }
}

int Server::connectionCallback(Peer* peer, uint64_t stream_id, uint8_t* bytes, size_t length,
    picoquic_call_back_event_t fin_or_event, void* v_stream_ctx) {

    switch (fin_or_event) {
        case picoquic_callback_stream_data:
        case picoquic_callback_stream_fin: {
            assert(stream_ctx::IsClientId(stream_id));
            auto stream_ctx = (in_stream_ctx *)v_stream_ctx;
            const auto is_fin = fin_or_event == picoquic_callback_stream_fin;

            if (stream_ctx == NULL) {
                if (is_fin) { // Single frame packet
                    if (length > 0) {
                        onPacket(peer, data::out_view(bytes, length), PacketFlags::TINY);
                    }
                    peer->reset(stream_id);
                    break;
                }
                // New long stream from server
                stream_ctx = peer->receive(stream_id);
            }

            if (length > 0) {
                stream_ctx->buffer.write(bytes, length);
            }

            if (is_fin) {
                if (onPacket(peer, data::out_view(stream_ctx->buffer.data.data(), stream_ctx->buffer.data.size()), PacketFlags::NONE)) {
                    peer->close(stream_ctx);
                } else {
                    return -1;
                }
            }
            break;
        }

        case picoquic_callback_datagram:
            if(!onPacket(peer, data::out_view(bytes, length), PacketFlags::DATAGRAM)) {
                return -1;
            }
            break;

        case picoquic_callback_stop_sending: /* Should not happen, treated as reset */
            /* Mark stream as abandoned, close the file, etc. */
            peer->reset(stream_id);
            /* Fall through */
        case picoquic_callback_stream_reset: /* Server reset stream #x */ {
            assert(stream_ctx::IsUnidirId(stream_id));

            // auto remote_error = picoquic_get_remote_stream_error(cnx, stream_id);
            //FIXME: if remote_error callback onError(remote_error)
            if (stream_ctx::IsClientId(stream_id)) {
                auto stream_ctx = (in_stream_ctx*)v_stream_ctx;
                if (stream_ctx == NULL) {
                    peer->close(stream_ctx);
                }
            } else {
                auto stream_ctx = (out_stream_ctx*)v_stream_ctx;
                if (stream_ctx == NULL) {
                    peer->close(stream_ctx);
                }
            }
            break;
        }

        case picoquic_callback_close:
        case picoquic_callback_stateless_reset:
        case picoquic_callback_application_close: /* Received application close */ {
            if (peer == nullptr)
                break;

            const auto is_app = fin_or_event == picoquic_callback_application_close;
            if (onDisconnect(peer, is_app, peer->getErrorCode(is_app))) {
                peer->setCallback(NULL, NULL);
                peers.remove(*peer);
            }
            break;
        }
        case picoquic_callback_version_negotiation:
            /* The server should never receive a version negotiation response */
            break;
        case picoquic_callback_stream_gap:
            /* This callback is never used. */
            break;
        case picoquic_callback_prepare_to_send: {
            /* Active sending API */
            assert(stream_ctx::IsServerId(stream_id));

            auto stream_ctx = (out_stream_ctx*)v_stream_ctx;
            if (stream_ctx == NULL) {
                return -1;
            }

            if (!stream_ctx->buffer.isDone()) {
                size_t available = stream_ctx->buffer.remaining();
                int is_fin = 1;

                /* The length parameter marks the space available in the packet */
                if (available > length) {
                    available = length;
                    is_fin = 0;
                }
                uint8_t* buffer = picoquic_provide_stream_data_buffer(bytes, available, is_fin, !is_fin);
                if (buffer == NULL) {
                    return -1;
                }

                // MAYBE: try zero copy
                stream_ctx->buffer.read(buffer, available);
            } else {
                LOG_W("Out stream reset");
                peer->reset(stream_id);
            }
            if (stream_ctx->buffer.isDone()) {
                //MAYBE: reuse
                peer->close(stream_ctx);
            }
            break;
        }
        case picoquic_callback_almost_ready:
            assert(peer == nullptr);
            // MAYBE: use pre-connect for server status
            if (Connection::GetSize(peers) >= max_connections) {
                LOG_W("Server is full");
                return -1;
            }
            break;
        case picoquic_callback_ready:
            /* TODO: Check that the transport parameters are what the sample expects */
            if (auto reason = onConnect(peer)) {
                peer->release(reason.value());
                peers.remove(*peer);
                return -1;
            }
            break;
        default:
            //MAYBE: picoquic_callback_request_alpn_list
            //MAYBE: picoquic_callback_set_alpn
            //MAYBE: picoquic_callback_pacing_changed
            /* unexpected -- just ignore. */
            break;
    }

    return 0;
}

Peer *Server::connect(picoquic_cnx_t * cnx) {
    return &peers.emplace_front(cnx);
}
