//! Platform-neutral remote client relay session.
//!
//! The Windows binary will own TAP discovery, routing, and adapter I/O. This
//! crate owns the shared relay-facing state machine: connect to the relay,
//! announce the client's virtual MAC, and exchange Ethernet frames as QUIC
//! datagrams after the control-plane welcome.

use std::{
    fs,
    io::ErrorKind,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    path::{Path, PathBuf},
    sync::Arc,
};

use anyhow::{Context, Result, bail};
use bytes::Bytes;
use lanparty_ctrl::{
    CONTROL_LENGTH_PREFIX_LEN, ControlMessage, EndpointHello, MAX_CONTROL_MESSAGE_LEN, RELAY_ALPN,
    RoomCode, ServerWelcome, decode_control_frame, encode_control_message,
};
use lanparty_proto::{EthernetFrame, FrameType, MacAddr, decode_datagram, encode_datagram};
use quinn::{ClientConfig, Endpoint, crypto::rustls::QuicClientConfig};
use rustls::pki_types::CertificateDer;

const MAX_CONTROL_FRAME_LEN: usize = CONTROL_LENGTH_PREFIX_LEN + MAX_CONTROL_MESSAGE_LEN;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ClientIdentity {
    virtual_mac: MacAddr,
}

impl ClientIdentity {
    pub fn new(virtual_mac: MacAddr) -> Result<Self> {
        if !virtual_mac.is_valid_client_identity() {
            bail!("client virtual MAC must be locally administered unicast");
        }

        Ok(Self { virtual_mac })
    }

    pub fn generate() -> Result<Self> {
        let mut octets = [0_u8; 6];
        getrandom::fill(&mut octets).context("failed to generate client virtual MAC")?;
        octets[0] = (octets[0] | 0x02) & 0xfe;

        Self::new(MacAddr::new(octets))
    }

    #[must_use]
    pub const fn virtual_mac(self) -> MacAddr {
        self.virtual_mac
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientIdentityStore {
    path: PathBuf,
}

impl ClientIdentityStore {
    pub fn new(path: impl Into<PathBuf>) -> Result<Self> {
        let path = path.into();
        if path.as_os_str().is_empty() {
            bail!("client identity path cannot be empty");
        }

        Ok(Self { path })
    }

    #[must_use]
    pub fn path(&self) -> &Path {
        &self.path
    }

    pub fn load_or_create(&self) -> Result<ClientIdentity> {
        match fs::read(&self.path) {
            Ok(bytes) => read_identity(&bytes)
                .with_context(|| format!("failed to read client identity {}", self.path.display())),
            Err(error) if error.kind() == ErrorKind::NotFound => {
                let identity = ClientIdentity::generate()?;
                write_identity(&self.path, identity)?;
                Ok(identity)
            }
            Err(error) => Err(error)
                .with_context(|| format!("failed to read client identity {}", self.path.display())),
        }
    }
}

#[derive(Debug, serde::Deserialize, serde::Serialize)]
struct ClientIdentityFile {
    virtual_mac: String,
}

fn read_identity(bytes: &[u8]) -> Result<ClientIdentity> {
    let file: ClientIdentityFile =
        serde_json::from_slice(bytes).context("failed to parse client identity JSON")?;
    let virtual_mac = file
        .virtual_mac
        .parse()
        .context("failed to parse client identity virtual MAC")?;

    ClientIdentity::new(virtual_mac)
}

fn write_identity(path: &Path, identity: ClientIdentity) -> Result<()> {
    let file_name = path
        .file_name()
        .context("client identity path must include a file name")?;
    if let Some(parent) = path
        .parent()
        .filter(|parent| !parent.as_os_str().is_empty())
    {
        fs::create_dir_all(parent)
            .with_context(|| format!("failed to create identity directory {}", parent.display()))?;
    }

    let mut temp_file_name = file_name.to_os_string();
    temp_file_name.push(".tmp");
    let temp_path = path.with_file_name(temp_file_name);
    let contents = serde_json::to_vec_pretty(&ClientIdentityFile {
        virtual_mac: identity.virtual_mac().to_string(),
    })
    .context("failed to encode client identity JSON")?;

    fs::write(&temp_path, contents)
        .with_context(|| format!("failed to write client identity {}", temp_path.display()))?;
    fs::rename(&temp_path, path)
        .with_context(|| format!("failed to persist client identity {}", path.display()))?;

    Ok(())
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientSessionConfig {
    relay_addr: SocketAddr,
    server_name: String,
    relay_ca_cert_der: Vec<u8>,
    room: RoomCode,
    virtual_mac: MacAddr,
    max_datagram_size: u16,
}

impl ClientSessionConfig {
    pub fn new(
        relay_addr: SocketAddr,
        server_name: impl Into<String>,
        relay_ca_cert_der: Vec<u8>,
        room: RoomCode,
        virtual_mac: MacAddr,
        max_datagram_size: u16,
    ) -> Result<Self> {
        let server_name = server_name.into();
        if server_name.trim().is_empty() {
            bail!("relay server name cannot be empty");
        }

        if relay_ca_cert_der.is_empty() {
            bail!("relay CA certificate cannot be empty");
        }

        EndpointHello::client(room.clone(), virtual_mac, max_datagram_size)
            .context("invalid client identity or datagram budget")?;

        Ok(Self {
            relay_addr,
            server_name,
            relay_ca_cert_der,
            room,
            virtual_mac,
            max_datagram_size,
        })
    }

    #[must_use]
    pub const fn relay_addr(&self) -> SocketAddr {
        self.relay_addr
    }

    #[must_use]
    pub fn server_name(&self) -> &str {
        &self.server_name
    }

    #[must_use]
    pub fn relay_ca_cert_der(&self) -> &[u8] {
        &self.relay_ca_cert_der
    }

    #[must_use]
    pub const fn room(&self) -> &RoomCode {
        &self.room
    }

    #[must_use]
    pub const fn virtual_mac(&self) -> MacAddr {
        self.virtual_mac
    }

    #[must_use]
    pub const fn max_datagram_size(&self) -> u16 {
        self.max_datagram_size
    }
}

#[derive(Debug)]
pub struct ClientSession {
    endpoint: Endpoint,
    connection: quinn::Connection,
    config: ClientSessionConfig,
    welcome: ServerWelcome,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ReceivedEthernetFrame {
    source_peer_id: u32,
    payload: Bytes,
}

impl ReceivedEthernetFrame {
    #[must_use]
    pub const fn source_peer_id(&self) -> u32 {
        self.source_peer_id
    }

    #[must_use]
    pub fn payload(&self) -> &[u8] {
        &self.payload
    }
}

impl ClientSession {
    #[must_use]
    pub const fn config(&self) -> &ClientSessionConfig {
        &self.config
    }

    #[must_use]
    pub const fn welcome(&self) -> &ServerWelcome {
        &self.welcome
    }

    #[must_use]
    pub fn relay_io(&self) -> ClientRelayIo {
        ClientRelayIo::new(self.connection.clone(), self.welcome.clone())
    }

    pub fn send_ethernet(&self, frame: &[u8]) -> Result<()> {
        self.relay_io().send_ethernet(frame)
    }

    pub async fn recv_ethernet(&self) -> Result<ReceivedEthernetFrame> {
        self.relay_io().recv_ethernet().await
    }

    pub async fn shutdown(self, reason: &str) {
        self.connection.close(0_u32.into(), reason.as_bytes());
        self.endpoint.wait_idle().await;
    }
}

#[derive(Debug, Clone)]
pub struct ClientRelayIo {
    connection: quinn::Connection,
    welcome: ServerWelcome,
}

impl ClientRelayIo {
    #[must_use]
    fn new(connection: quinn::Connection, welcome: ServerWelcome) -> Self {
        Self {
            connection,
            welcome,
        }
    }

    #[must_use]
    pub const fn welcome(&self) -> &ServerWelcome {
        &self.welcome
    }

    pub fn send_ethernet(&self, frame: &[u8]) -> Result<()> {
        EthernetFrame::parse(frame).context("client Ethernet frame is malformed")?;
        let datagram = encode_datagram(
            FrameType::Ethernet,
            self.welcome.room_id(),
            self.welcome.peer_id(),
            0,
            frame,
        )
        .context("failed to encode client Ethernet datagram")?;

        self.connection
            .send_datagram(Bytes::from(datagram))
            .context("failed to send client Ethernet datagram")?;

        Ok(())
    }

    pub async fn recv_ethernet(&self) -> Result<ReceivedEthernetFrame> {
        loop {
            let datagram = self.connection.read_datagram().await?;
            let Ok(packet) = decode_datagram(&datagram) else {
                continue;
            };
            let header = packet.header();
            if header.frame_type() != FrameType::Ethernet
                || header.room_id() != self.welcome.room_id()
                || header.peer_id() == self.welcome.peer_id()
            {
                continue;
            }
            if EthernetFrame::parse(packet.payload()).is_err() {
                continue;
            }

            return Ok(ReceivedEthernetFrame {
                source_peer_id: header.peer_id(),
                payload: Bytes::copy_from_slice(packet.payload()),
            });
        }
    }
}

pub async fn connect_client(config: ClientSessionConfig) -> Result<ClientSession> {
    let client_config = relay_client_config(config.relay_ca_cert_der())?;
    let mut endpoint = Endpoint::client(client_bind_addr(config.relay_addr()))
        .context("failed to bind client QUIC endpoint")?;
    endpoint.set_default_client_config(client_config);

    let connection = endpoint
        .connect(config.relay_addr(), config.server_name())?
        .await
        .with_context(|| format!("failed to connect to relay {}", config.relay_addr()))?;
    let peer_datagram_size = connection
        .max_datagram_size()
        .context("relay did not negotiate QUIC DATAGRAM support")?;
    let hello_datagram_size = usize::from(config.max_datagram_size())
        .min(peer_datagram_size)
        .min(usize::from(u16::MAX)) as u16;
    let hello = EndpointHello::client(
        config.room().clone(),
        config.virtual_mac(),
        hello_datagram_size,
    )
    .context("failed to build client hello")?;
    let response = request_control_message(&connection, ControlMessage::Hello(hello)).await?;

    match response {
        ControlMessage::Welcome(welcome) => Ok(ClientSession {
            endpoint,
            connection,
            config,
            welcome,
        }),
        ControlMessage::Reject(reject) => bail!(
            "relay rejected client hello: {:?}: {}",
            reject.reason(),
            reject.message()
        ),
        other => bail!("relay sent unexpected client handshake response: {other:?}"),
    }
}

fn relay_client_config(relay_ca_cert_der: &[u8]) -> Result<ClientConfig> {
    let mut roots = rustls::RootCertStore::empty();
    roots
        .add(CertificateDer::from(relay_ca_cert_der.to_vec()))
        .context("failed to trust relay CA certificate")?;

    let mut client_crypto = rustls::ClientConfig::builder()
        .with_root_certificates(roots)
        .with_no_client_auth();
    client_crypto.alpn_protocols = vec![RELAY_ALPN.to_vec()];

    Ok(ClientConfig::new(Arc::new(
        QuicClientConfig::try_from(client_crypto).context("failed to build QUIC client config")?,
    )))
}

fn client_bind_addr(relay_addr: SocketAddr) -> SocketAddr {
    match relay_addr.ip() {
        IpAddr::V4(_) => SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0),
        IpAddr::V6(_) => SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 0),
    }
}

async fn request_control_message(
    connection: &quinn::Connection,
    message: ControlMessage,
) -> Result<ControlMessage> {
    let (mut send, mut recv) = connection.open_bi().await?;
    let request = encode_control_message(&message)?;
    send.write_all(&request).await?;
    send.finish()?;

    let response = recv.read_to_end(MAX_CONTROL_FRAME_LEN).await?;

    Ok(decode_control_frame(&response)?)
}

#[cfg(test)]
mod tests {
    use std::time::{Duration, SystemTime, UNIX_EPOCH};

    use bytes::Bytes;
    use lanparty_ctrl::Role;
    use quinn::{ServerConfig, TransportConfig, crypto::rustls::QuicServerConfig};
    use rustls::pki_types::{PrivateKeyDer, PrivatePkcs8KeyDer};

    use super::*;

    #[test]
    fn validates_client_config() {
        let room = RoomCode::new("ROOM1").unwrap();
        let cert = vec![1, 2, 3];
        let mac = MacAddr::new([0x02, 0, 0, 0, 0, 1]);

        assert!(
            ClientSessionConfig::new(
                "127.0.0.1:443".parse().unwrap(),
                "relay.local",
                cert.clone(),
                room.clone(),
                mac,
                1400,
            )
            .is_ok()
        );
        assert!(
            ClientSessionConfig::new(
                "127.0.0.1:443".parse().unwrap(),
                "",
                cert.clone(),
                room.clone(),
                mac,
                1400,
            )
            .is_err()
        );
        assert!(
            ClientSessionConfig::new(
                "127.0.0.1:443".parse().unwrap(),
                "relay.local",
                Vec::new(),
                room.clone(),
                mac,
                1400,
            )
            .is_err()
        );
        assert!(
            ClientSessionConfig::new(
                "127.0.0.1:443".parse().unwrap(),
                "relay.local",
                cert,
                room,
                MacAddr::BROADCAST,
                1400,
            )
            .is_err()
        );
    }

    #[test]
    fn generates_valid_client_identity() {
        let identity = ClientIdentity::generate().unwrap();

        assert!(identity.virtual_mac().is_valid_client_identity());
    }

    #[test]
    fn identity_store_creates_and_reuses_mac() {
        let path = unique_temp_identity_path();
        let store = ClientIdentityStore::new(&path).unwrap();

        let first = store.load_or_create().unwrap();
        let second = store.load_or_create().unwrap();
        let stored = std::fs::read_to_string(&path).unwrap();

        assert!(path.exists());
        assert!(first.virtual_mac().is_valid_client_identity());
        assert_eq!(first, second);
        assert!(stored.contains(&first.virtual_mac().to_string()));

        std::fs::remove_file(path).unwrap();
    }

    #[test]
    fn identity_store_rejects_invalid_saved_mac() {
        let path = unique_temp_identity_path();
        std::fs::write(
            &path,
            r#"{
  "virtual_mac": "ff:ff:ff:ff:ff:ff"
}"#,
        )
        .unwrap();
        let store = ClientIdentityStore::new(&path).unwrap();

        assert!(store.load_or_create().is_err());

        std::fs::remove_file(path).unwrap();
    }

    #[tokio::test]
    async fn connects_to_relay_control_stream_as_client() {
        let (server_config, certificate) = test_server_config();
        let endpoint = Endpoint::server(server_config, "127.0.0.1:0".parse().unwrap()).unwrap();
        let server_addr = endpoint.local_addr().unwrap();
        let server_task = tokio::spawn(async move {
            let incoming = endpoint.accept().await.unwrap();
            let connection = incoming.await.unwrap();
            let (mut send, mut recv) = connection.accept_bi().await.unwrap();
            let request = recv.read_to_end(MAX_CONTROL_FRAME_LEN).await.unwrap();
            let message = decode_control_frame(&request).unwrap();
            let ControlMessage::Hello(hello) = message else {
                panic!("expected client hello");
            };

            assert_eq!(hello.role(), Role::Client);
            assert_eq!(hello.room().as_str(), "ROOM1");
            assert_eq!(
                hello.announced_mac(),
                Some(MacAddr::new([0x02, 0, 0, 0, 0, 1]))
            );

            let response = encode_control_message(&ControlMessage::Welcome(
                ServerWelcome::new(7, 2, 1200).unwrap(),
            ))
            .unwrap();
            send.write_all(&response).await.unwrap();
            send.finish().unwrap();

            let datagram = connection.read_datagram().await.unwrap();
            let packet = decode_datagram(&datagram).unwrap();
            let header = packet.header();
            assert_eq!(header.frame_type(), FrameType::Ethernet);
            assert_eq!(header.room_id(), 7);
            assert_eq!(header.peer_id(), 2);
            assert_eq!(packet.payload(), ethernet_frame(b"to relay").as_slice());

            let response =
                encode_datagram(FrameType::Ethernet, 7, 1, 0, &ethernet_frame(b"from relay"))
                    .unwrap();
            connection.send_datagram(Bytes::from(response)).unwrap();

            connection.closed().await;
            endpoint.close(0_u32.into(), b"test complete");
            endpoint.wait_idle().await;
        });
        let config = ClientSessionConfig::new(
            server_addr,
            "lanparty-relay.local",
            certificate.as_ref().to_vec(),
            RoomCode::new("ROOM1").unwrap(),
            MacAddr::new([0x02, 0, 0, 0, 0, 1]),
            1400,
        )
        .unwrap();

        let client = connect_client(config).await.unwrap();

        assert_eq!(
            client.config().virtual_mac(),
            MacAddr::new([0x02, 0, 0, 0, 0, 1])
        );
        assert_eq!(client.welcome().room_id(), 7);
        assert_eq!(client.welcome().peer_id(), 2);
        let relay_io = client.relay_io();
        assert_eq!(relay_io.welcome().peer_id(), 2);

        relay_io
            .send_ethernet(&ethernet_frame(b"to relay"))
            .unwrap();
        let received = tokio::time::timeout(Duration::from_secs(5), relay_io.recv_ethernet())
            .await
            .unwrap()
            .unwrap();
        assert_eq!(received.source_peer_id(), 1);
        assert_eq!(received.payload(), ethernet_frame(b"from relay").as_slice());

        client.shutdown("test complete").await;
        tokio::time::timeout(Duration::from_secs(5), server_task)
            .await
            .unwrap()
            .unwrap();
    }

    fn test_server_config() -> (ServerConfig, CertificateDer<'static>) {
        let certified_key =
            rcgen::generate_simple_self_signed(vec!["lanparty-relay.local".into()]).unwrap();
        let certificate = certified_key.cert.der().clone();
        let cert_chain = vec![certificate.clone()];
        let private_key = PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(
            certified_key.signing_key.serialize_der(),
        ));
        let mut tls_config = rustls::ServerConfig::builder()
            .with_no_client_auth()
            .with_single_cert(cert_chain, private_key)
            .unwrap();
        tls_config.alpn_protocols = vec![RELAY_ALPN.to_vec()];

        let mut server_config =
            ServerConfig::with_crypto(Arc::new(QuicServerConfig::try_from(tls_config).unwrap()));
        let mut transport = TransportConfig::default();
        transport.datagram_receive_buffer_size(Some(4 * 1024 * 1024));
        transport.datagram_send_buffer_size(4 * 1024 * 1024);
        server_config.transport_config(Arc::new(transport));

        (server_config, certificate)
    }

    fn ethernet_frame(payload: &[u8]) -> Vec<u8> {
        let mut frame = Vec::new();
        frame.extend_from_slice(&[0x02, 0, 0, 0, 0, 2]);
        frame.extend_from_slice(&[0x02, 0, 0, 0, 0, 1]);
        frame.extend_from_slice(&0x0800_u16.to_be_bytes());
        frame.extend_from_slice(payload);
        frame
    }

    fn unique_temp_identity_path() -> std::path::PathBuf {
        let nanos = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();

        std::env::temp_dir().join(format!(
            "lanparty-client-identity-{}-{nanos}.json",
            std::process::id()
        ))
    }
}