// SPDX-License-Identifier: GPL-3.0-or-later
//! Wire protocol messages for the etr QUIC transport.
//!
//! All messages are protobuf-encoded and transmitted on QUIC bidirectional
//! streams opened by the client.  Each stream begins with a 1-byte tag
//! (see [`crate::quic`]) that identifies the stream type, followed by
//! length-prefixed [`Envelope`] messages on control streams, raw PTY chunks
//! on the PTY stream, and optional framed data on forward streams.
//!
//! ## Stream multiplexing
//!
//! | Tag  | Stream   | Client → Server                  | Server → Client               |
//! |------|----------|----------------------------------|-------------------------------|
//! | 0x01 | Control  | `SessionOpen`, then `Envelope`*  | `SessionAccept`, then `Envelope`* |
//! | 0x02 | PTY      | stdin chunks `[seq][len][data]`  | PTY output chunks `[seq][len][data]` |
//! | 0x03 | Forward  | `StreamOpen` header + raw bytes  | raw bytes                     |
//!
//! `Envelope`* = any combination of `Heartbeat`, `TerminalResize`, `Disconnect`.

/// Stream type carried in [`StreamOpen`].
#[derive(Clone, Copy, Debug, PartialEq, Eq, prost::Enumeration)]
#[repr(i32)]
pub enum StreamType {
    Terminal = 0,
    PortForward = 1,
}

/// Transport protocol for a port-forward stream.
#[derive(Clone, Copy, Debug, PartialEq, Eq, prost::Enumeration)]
#[repr(i32)]
pub enum ForwardProto {
    /// TCP forwarding: one QUIC stream per TCP connection.
    Tcp = 0,
    /// UDP forwarding: one shared QUIC stream per `-L` spec.
    Udp = 1,
}

/// Sent by the client on the control stream to resume or open a session.
#[derive(Clone, PartialEq, prost::Message)]
pub struct SessionOpen {
    /// 16-byte session ID generated by the client.
    #[prost(bytes = "vec", tag = "1")]
    pub session_id: Vec<u8>,
    /// Shared secret bootstrapped via SSH stdin.
    #[prost(string, tag = "2")]
    pub passkey: String,
    /// Highest PTY-output seq the client received from the server, per stream.
    /// The server replays any history above this watermark.
    #[prost(map = "uint32, uint64", tag = "3")]
    pub last_received_seq: std::collections::HashMap<u32, u64>,
    /// Reverse port forwarding specifications (e.g. "remote_port:local_host:local_port[/tcp|/udp]").
    #[prost(string, repeated, tag = "4")]
    pub reverse_forwards: Vec<String>,
    /// If true, reverse-forward listeners on the server bind to all interfaces ([::] dual-stack)
    /// instead of loopback only.  Corresponds to the client's -g / --gateway-ports flag.
    #[prost(bool, tag = "5")]
    pub gateway_ports: bool,
}

/// Sent by the server in reply to [`SessionOpen`] once the session is verified.
#[derive(Clone, PartialEq, prost::Message)]
pub struct SessionAccept {
    /// Highest stdin seq the server received from the client, per stream.
    /// The client replays any history above this watermark.
    #[prost(map = "uint32, uint64", tag = "1")]
    pub last_received_seq: std::collections::HashMap<u32, u64>,
}

/// Header sent by the client at the start of a forward (0x03) stream,
/// describing the remote target to connect to.
#[derive(Clone, PartialEq, prost::Message)]
pub struct StreamOpen {
    #[prost(uint32, tag = "1")]
    pub stream_id: u32,
    #[prost(enumeration = "StreamType", tag = "2")]
    pub stream_type: i32,
    /// Remote host the server should connect to.
    #[prost(string, tag = "3")]
    pub remote_host: String,
    /// Remote port.
    #[prost(uint32, tag = "4")]
    pub remote_port: u32,
    /// TCP or UDP (default: TCP).
    #[prost(enumeration = "ForwardProto", tag = "5")]
    pub forward_proto: i32,
}

/// Close a forward stream cleanly or signal an error.
#[derive(Clone, PartialEq, prost::Message)]
pub struct StreamClose {
    #[prost(uint32, tag = "1")]
    pub stream_id: u32,
    /// 0 = clean close; non-zero = error code.
    #[prost(uint32, tag = "2")]
    pub error_code: u32,
}

/// Terminal window resize event sent on the control stream.
#[derive(Clone, PartialEq, prost::Message)]
pub struct TerminalResize {
    #[prost(uint32, tag = "1")]
    pub rows: u32,
    #[prost(uint32, tag = "2")]
    pub cols: u32,
}

/// Keepalive message exchanged on the control stream to detect idle timeouts.
///
/// `last_received_seq` piggybacks the sender's receive watermark so the peer
/// can continuously trim its send history without waiting for a reconnect.
#[derive(Clone, PartialEq, prost::Message)]
pub struct Heartbeat {
    /// Highest seq received per stream (stream_id → seq).  Empty map = pure keepalive.
    #[prost(map = "uint32, uint64", tag = "1")]
    pub last_received_seq: std::collections::HashMap<u32, u64>,
}

/// Clean-disconnect signal on the control stream.
/// The receiver should not attempt to reconnect.
#[derive(Clone, PartialEq, prost::Message)]
pub struct Disconnect {}

/// Framing for one UDP datagram on a shared UDP forward (0x03) stream.
/// The peer address is embedded so the far side can route replies correctly.
#[derive(Clone, PartialEq, prost::Message)]
pub struct UdpDatagram {
    #[prost(string, tag = "1")]
    pub peer_addr: String,
    #[prost(uint32, tag = "2")]
    pub peer_port: u32,
    #[prost(bytes = "vec", tag = "3")]
    pub data: Vec<u8>,
}

/// Top-level protobuf envelope for control-stream messages.
#[derive(Clone, PartialEq, prost::Message)]
pub struct Envelope {
    #[prost(oneof = "Payload", tags = "3, 4, 7, 8, 9, 10, 11, 12")]
    pub payload: Option<Payload>,
}

#[derive(Clone, PartialEq, prost::Oneof)]
pub enum Payload {
    #[prost(message, tag = "3")]
    StreamOpen(StreamOpen),
    #[prost(message, tag = "4")]
    StreamClose(StreamClose),
    #[prost(message, tag = "7")]
    TerminalResize(TerminalResize),
    #[prost(message, tag = "8")]
    Heartbeat(Heartbeat),
    #[prost(message, tag = "9")]
    Disconnect(Disconnect),
    #[prost(message, tag = "10")]
    SessionOpen(SessionOpen),
    #[prost(message, tag = "11")]
    SessionAccept(SessionAccept),
    #[prost(message, tag = "12")]
    UdpDatagram(UdpDatagram),
}

#[cfg(test)]
mod tests {
    use super::*;
    use prost::Message;

    /// Verify that a [`SessionOpen`] with all fields set (including `reverse_forwards` and
    /// `gateway_ports`) survives a protobuf encode → decode round trip inside an [`Envelope`].
    /// This catches field-tag collisions and ensures the new bool field is not silently dropped.
    #[test]
    fn test_session_open_round_trip() {
        let env = Envelope {
            payload: Some(Payload::SessionOpen(SessionOpen {
                session_id: vec![1u8; 16],
                passkey: "secret".to_string(),
                last_received_seq: [(0u32, 7u64)].into(),
                reverse_forwards: vec!["9090:localhost:90".to_string()],
                gateway_ports: true,
            })),
        };
        let decoded = Envelope::decode(env.encode_to_vec().as_slice()).unwrap();
        assert_eq!(decoded, env);
        // Verify gateway_ports survives the round trip.
        if let Some(Payload::SessionOpen(so)) = decoded.payload {
            assert!(so.gateway_ports);
            assert_eq!(so.reverse_forwards, vec!["9090:localhost:90"]);
        } else {
            panic!("expected SessionOpen payload");
        }
    }

    #[test]
    fn test_session_accept_round_trip() {
        let env = Envelope {
            payload: Some(Payload::SessionAccept(SessionAccept {
                last_received_seq: [(0u32, 3u64)].into(),
            })),
        };
        let decoded = Envelope::decode(env.encode_to_vec().as_slice()).unwrap();
        assert_eq!(decoded, env);
    }

    #[test]
    fn test_stream_open_close_round_trip() {
        let open = Envelope {
            payload: Some(Payload::StreamOpen(StreamOpen {
                stream_id: 1,
                stream_type: StreamType::PortForward as i32,
                remote_host: "localhost".to_string(),
                remote_port: 5432,
                forward_proto: ForwardProto::Tcp as i32,
            })),
        };
        let close = Envelope {
            payload: Some(Payload::StreamClose(StreamClose {
                stream_id: 1,
                error_code: 0,
            })),
        };
        assert_eq!(
            Envelope::decode(open.encode_to_vec().as_slice()).unwrap(),
            open
        );
        assert_eq!(
            Envelope::decode(close.encode_to_vec().as_slice()).unwrap(),
            close
        );
    }

    #[test]
    fn test_heartbeat_disconnect_round_trip() {
        for payload in [
            Payload::Heartbeat(Heartbeat {
                last_received_seq: [(0u32, 42u64)].into(),
            }),
            Payload::Heartbeat(Heartbeat::default()),
            Payload::Disconnect(Disconnect {}),
        ] {
            let env = Envelope {
                payload: Some(payload),
            };
            assert_eq!(
                Envelope::decode(env.encode_to_vec().as_slice()).unwrap(),
                env
            );
        }
    }

    #[test]
    fn test_udp_datagram_round_trip() {
        let env = Envelope {
            payload: Some(Payload::UdpDatagram(UdpDatagram {
                peer_addr: "127.0.0.1".to_string(),
                peer_port: 12345,
                data: b"dns query".to_vec(),
            })),
        };
        let decoded = Envelope::decode(env.encode_to_vec().as_slice()).unwrap();
        assert_eq!(decoded, env);
    }
}
