// SPDX-License-Identifier: AGPL-3.0-only

//! Synvael application handshake over an established QUIC connection.

use quinn::{Connection, Incoming, RecvStream, SendStream, VarInt};
use shared::protocol::{
    ClientHello, ControlMessage, Disconnect, HandshakeAck, HandshakeReject, PROTOCOL_VERSION,
    RejectReason, StreamLayout,
};
use tracing::{info, warn};

use crate::codec::{MAX_CONTROL_FRAME_LEN, read_frame, write_frame};
use crate::error::HandshakeError;

/// Application close code used when a peer is rejected during the handshake.
const CLOSE_CODE_REJECTED: u32 = 1;

/// Application close code used for an orderly, graceful disconnect.
const CLOSE_CODE_GRACEFUL: u32 = 0;

/// Upper bound on the wait for a rejected client to read the `HandshakeReject` and close, before the server tears the connection down anyway.
const REJECT_DELIVERY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);

/// A completed client-side handshake: an established connection, the retained control stream, and the server's acceptance parameters.
#[derive(Debug)]
pub struct Connected {
    /// The established QUIC connection. Additional streams are opened from it.
    pub connection: Connection,
    /// The control stream (stream 0), retained so it can carry chat, commands, and disconnect. Never finished after the handshake.
    pub control: (SendStream, RecvStream),
    /// The negotiated session parameters returned by the server.
    pub ack: HandshakeAck,
}

/// A completed server-side handshake for one connection: the established connection, the retained control stream, and the client's presented identity.
#[derive(Debug)]
pub struct ServerConnection {
    /// The established QUIC connection. Additional streams are accepted from it.
    pub connection: Connection,
    /// The control stream (stream 0), retained so it can carry chat, commands, and disconnect. Never finished after the handshake.
    pub control: (SendStream, RecvStream),
    /// The `ClientHello` the accepted client presented.
    pub hello: ClientHello,
}

/// Connects to a Synvael server: establishes the QUIC connection, opens the control stream, sends `ClientHello`, and awaits the server's response.
///
/// # Errors
///
/// Returns [`HandshakeError::Connect`] if the connection attempt cannot be initiated, [`HandshakeError::Connection`] if the QUIC connection fails to establish, [`HandshakeError::Codec`] if the `ClientHello` cannot be written or the reply cannot be read, [`HandshakeError::Rejected`] if the server refuses the handshake, and [`HandshakeError::UnexpectedMessage`] if the server replies with a control message other than `HandshakeAck` or `HandshakeReject`.
pub async fn connect(
    endpoint: &quinn::Endpoint,
    server_addr: std::net::SocketAddr,
    server_name: &str,
    hello: ClientHello,
) -> Result<Connected, HandshakeError> {
    let connection = endpoint.connect(server_addr, server_name)?.await?;

    // The client opens the control stream; the server accepts it. The stream first appears on the server once the `ClientHello` bytes are written below.
    let (mut send, mut recv) = connection.open_bi().await?;

    write_frame(&mut send, &ControlMessage::ClientHello(hello)).await?;

    match read_frame::<ControlMessage>(&mut recv, MAX_CONTROL_FRAME_LEN).await? {
        ControlMessage::HandshakeAck(ack) => {
            info!(
                protocol_version = ack.protocol_version,
                server_build = %ack.server_build,
                "handshake accepted by server"
            );
            Ok(Connected {
                connection,
                control: (send, recv),
                ack,
            })
        }
        ControlMessage::HandshakeReject(rej) => {
            warn!(reason = ?rej.reason, detail = %rej.detail, "handshake rejected by server");
            Err(HandshakeError::Rejected(rej))
        }
        _ => Err(HandshakeError::UnexpectedMessage),
    }
}

/// Accepts one incoming connection: completes the QUIC handshake, reads the client's `ClientHello`, validates it, and replies with `HandshakeAck` or `HandshakeReject`.
///
/// On a protocol-version mismatch a [`HandshakeReject`] is sent to the client, the connection is closed with [`CLOSE_CODE_REJECTED`], and [`HandshakeError::VersionMismatch`] is returned locally.
///
/// # Errors
///
/// Returns [`HandshakeError::Connection`] if the QUIC connection fails to establish or the control stream cannot be accepted, [`HandshakeError::Codec`] if the `ClientHello` cannot be read or a reply cannot be written, [`HandshakeError::UnexpectedMessage`] if the first control message is not a `ClientHello`, and [`HandshakeError::VersionMismatch`] if the client's protocol version does not match the server's.
pub async fn accept_connection(
    incoming: Incoming,
    server_build: String,
    tick_rate_hint: u16,
) -> Result<ServerConnection, HandshakeError> {
    let connection = incoming.await?;

    // The client opened the control stream; the server accepts it here.
    let (mut send, mut recv) = connection.accept_bi().await?;

    let ControlMessage::ClientHello(hello) =
        read_frame::<ControlMessage>(&mut recv, MAX_CONTROL_FRAME_LEN).await?
    else {
        return Err(HandshakeError::UnexpectedMessage);
    };

    if hello.protocol_version != PROTOCOL_VERSION {
        let reject = HandshakeReject {
            reason: RejectReason::ProtocolMismatch,
            detail: format!(
                "client protocol version {}, server protocol version {PROTOCOL_VERSION}",
                hello.protocol_version
            ),
            upgrade_url: None,
        };
        warn!(
            client = hello.protocol_version,
            server = PROTOCOL_VERSION,
            "rejecting client on protocol mismatch"
        );
        // Send the rejection, then keep the connection alive until the client has read it and closed. `Connection::close` (and dropping the connection) discards buffered stream data, so an immediate close would race the client's read and lose the reject frame. The wait is bounded so a misbehaving client cannot park the accept task indefinitely.
        write_frame(&mut send, &ControlMessage::HandshakeReject(reject)).await?;
        let _ = send.finish();
        let _ = tokio::time::timeout(REJECT_DELIVERY_TIMEOUT, connection.closed()).await;
        connection.close(
            VarInt::from_u32(CLOSE_CODE_REJECTED),
            b"protocol version mismatch",
        );
        return Err(HandshakeError::VersionMismatch {
            client: hello.protocol_version,
            server: PROTOCOL_VERSION,
        });
    }

    // TODO: validate installed_packs against the world's required packs once the modlist-matching concept lands.
    let ack = HandshakeAck {
        protocol_version: PROTOCOL_VERSION,
        server_build,
        world_packs: vec![],
        missing_packs: vec![],
        stream_layout: StreamLayout::default(),
        tick_rate_hint,
    };
    write_frame(&mut send, &ControlMessage::HandshakeAck(ack)).await?;

    info!(
        display_name = %hello.player_identity.display_name,
        "handshake accepted"
    );

    Ok(ServerConnection {
        connection,
        control: (send, recv),
        hello,
    })
}

/// Sends a `Disconnect` control message and then closes the connection cleanly.
///
/// The control stream's send half is borrowed mutably to write the final frame; the connection is then closed with [`CLOSE_CODE_GRACEFUL`]. A failure to write the `Disconnect` frame is logged rather than propagated, since the connection is closed unconditionally afterwards. Both [`Connected`] and [`ServerConnection`] expose their parts as `connection` and `control`, so either can call this with `(&conn.connection, &mut conn.control.0, reason)`.
pub async fn graceful_disconnect(connection: &Connection, send: &mut SendStream, reason: &str) {
    if let Err(error) = write_frame(
        send,
        &ControlMessage::Disconnect(Disconnect {
            reason: reason.to_owned(),
        }),
    )
    .await
    {
        warn!(%error, "failed to send disconnect frame; closing connection regardless");
    }
    connection.close(VarInt::from_u32(CLOSE_CODE_GRACEFUL), reason.as_bytes());
}