use crate::config::{ServerConfig, load_server_config};
use crate::crypto;
use crate::native::{
    self, ForwardingKind, ForwardingRequest, NativeAuthOk, NativeServerHello,
    derive_native_session_key, generate_native_ephemeral, host_public_key, load_or_create_host_key,
    sign_server_hello, verify_native_user_auth_from_config,
};
use crate::protocol::{
    self, AttachReject, CLIENT_TO_SERVER, NativeAuthOkBody, NativeClientHelloBody,
    NativeServerHelloBody, NativeUserAuthBody, PacketKind, SERVER_TO_CLIENT,
};
use crate::transport::{
    DoshTransport, SessionEvent, SessionRole, SessionTransportConfig, TransportConfig,
    service_name_from_target,
};
use crate::udp::is_transient_udp_send_error;
use anyhow::{Context, Result, anyhow, bail};
use ed25519_dalek::SigningKey;
use std::collections::{HashMap, HashSet};
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::net::UdpSocket;

#[derive(Debug, Clone)]
pub struct DoshServerConfig {
    pub server: ServerConfig,
    pub bind_addr: Option<SocketAddr>,
    pub services: HashSet<String>,
    pub transport: TransportConfig,
    pub require_current_user: bool,
    pub auth_timeout: Duration,
}

impl DoshServerConfig {
    pub fn new(server: ServerConfig) -> Self {
        Self {
            server,
            bind_addr: None,
            services: HashSet::new(),
            transport: TransportConfig::default(),
            require_current_user: true,
            auth_timeout: Duration::from_secs(30),
        }
    }

    pub fn bind_addr(mut self, addr: SocketAddr) -> Self {
        self.bind_addr = Some(addr);
        self
    }

    pub fn service(mut self, name: impl Into<String>) -> Result<Self> {
        self.services.insert(validate_service_name(name.into())?);
        Ok(self)
    }

    pub fn services(mut self, names: impl IntoIterator<Item = impl Into<String>>) -> Result<Self> {
        for name in names {
            self.services.insert(validate_service_name(name.into())?);
        }
        Ok(self)
    }

    pub fn require_current_user(mut self, value: bool) -> Self {
        self.require_current_user = value;
        self
    }

    pub fn transport(mut self, transport: TransportConfig) -> Self {
        self.transport = transport;
        self
    }
}

impl Default for DoshServerConfig {
    fn default() -> Self {
        Self::new(ServerConfig::default())
    }
}

#[derive(Debug, Clone)]
pub struct DoshAccepted {
    pub conn_id: [u8; 16],
    pub user: String,
    pub session: String,
    pub services: Vec<String>,
    pub peer_addr: SocketAddr,
}

#[derive(Debug, Clone)]
pub enum DoshServerEvent {
    Accepted(DoshAccepted),
    Session {
        conn_id: [u8; 16],
        event: SessionEvent,
    },
    Ignored,
}

#[derive(Debug, Clone)]
struct PendingServerAuth {
    client: native::NativeClientHello,
    server: NativeServerHello,
    session_key: [u8; 32],
    peer: SocketAddr,
    created_at: Instant,
}

pub struct DoshServer {
    socket: Arc<UdpSocket>,
    config: DoshServerConfig,
    host_signing: SigningKey,
    pending: HashMap<[u8; 16], PendingServerAuth>,
    transports: HashMap<[u8; 16], DoshTransport>,
    accepted: HashMap<[u8; 16], DoshAccepted>,
}

impl DoshServer {
    pub async fn load() -> Result<Self> {
        Self::bind(DoshServerConfig::new(load_server_config(None)?)).await
    }

    pub async fn load_from_path(path: Option<PathBuf>) -> Result<Self> {
        Self::bind(DoshServerConfig::new(load_server_config(path)?)).await
    }

    pub async fn bind(config: DoshServerConfig) -> Result<Self> {
        let bind_addr = match config.bind_addr {
            Some(addr) => addr,
            None => format!("{}:{}", config.server.bind, config.server.port)
                .parse()
                .with_context(|| {
                    format!(
                        "parse Dosh bind address {}:{}",
                        config.server.bind, config.server.port
                    )
                })?,
        };
        let host_signing = load_or_create_host_key(&config.server)?;
        let socket = Arc::new(UdpSocket::bind(bind_addr).await?);
        Ok(Self {
            socket,
            config,
            host_signing,
            pending: HashMap::new(),
            transports: HashMap::new(),
            accepted: HashMap::new(),
        })
    }

    pub fn local_addr(&self) -> Result<SocketAddr> {
        Ok(self.socket.local_addr()?)
    }

    pub fn connection(&self, conn_id: &[u8; 16]) -> Option<&DoshAccepted> {
        self.accepted.get(conn_id)
    }

    pub fn transport(&self, conn_id: &[u8; 16]) -> Option<&DoshTransport> {
        self.transports.get(conn_id)
    }

    pub fn transport_mut(&mut self, conn_id: &[u8; 16]) -> Option<&mut DoshTransport> {
        self.transports.get_mut(conn_id)
    }

    pub async fn recv(&mut self) -> Result<DoshServerEvent> {
        let mut buf = vec![0u8; 65535];
        loop {
            self.expire_pending();
            let (n, peer) = self.socket.recv_from(&mut buf).await?;
            let packet = match protocol::decode(&buf[..n]) {
                Ok(packet) => packet,
                Err(_) => continue,
            };
            if packet.header.kind == PacketKind::NativeClientHello {
                self.handle_client_hello(peer, packet.body).await?;
                return Ok(DoshServerEvent::Ignored);
            }
            if packet.header.kind == PacketKind::NativeUserAuth {
                return self.handle_user_auth(peer, &packet).await;
            }
            if let Some(transport) = self.transports.get_mut(&packet.header.conn_id) {
                let event = transport.handle_datagram(&buf[..n], peer).await?;
                return Ok(DoshServerEvent::Session {
                    conn_id: packet.header.conn_id,
                    event,
                });
            }
            self.send_reject(peer, packet.header.conn_id, "unknown Dosh connection")
                .await?;
            return Ok(DoshServerEvent::Ignored);
        }
    }

    pub async fn accept_stream(&mut self, conn_id: [u8; 16], stream_id: u64) -> Result<()> {
        self.transport_mut(&conn_id)
            .ok_or_else(|| anyhow!("unknown Dosh connection"))?
            .accept_stream(stream_id)
            .await
    }

    pub async fn reject_stream(
        &mut self,
        conn_id: [u8; 16],
        stream_id: u64,
        reason: impl Into<String>,
    ) -> Result<()> {
        self.transport_mut(&conn_id)
            .ok_or_else(|| anyhow!("unknown Dosh connection"))?
            .reject_stream(stream_id, reason)
            .await
    }

    pub async fn send(
        &mut self,
        conn_id: [u8; 16],
        stream_id: u64,
        bytes: impl Into<Vec<u8>>,
    ) -> Result<()> {
        self.transport_mut(&conn_id)
            .ok_or_else(|| anyhow!("unknown Dosh connection"))?
            .send(stream_id, bytes)
            .await
    }

    pub async fn close(&mut self, conn_id: [u8; 16], stream_id: u64) -> Result<()> {
        self.transport_mut(&conn_id)
            .ok_or_else(|| anyhow!("unknown Dosh connection"))?
            .close(stream_id)
            .await
    }

    pub async fn maintenance(&mut self) -> Result<()> {
        for transport in self.transports.values_mut() {
            transport.maintenance().await?;
        }
        Ok(())
    }

    async fn handle_client_hello(&mut self, peer: SocketAddr, body: Vec<u8>) -> Result<()> {
        let req: NativeClientHelloBody = protocol::from_body(&body)?;
        if let Err(err) =
            native::check_native_protocol_version(req.hello.protocol_version, "client")
        {
            self.send_reject(peer, [0u8; 16], &err.to_string()).await?;
            return Ok(());
        }
        let result = self.build_server_hello(req.hello, peer);
        let (pending_id, hello) = match result {
            Ok(value) => value,
            Err(err) => {
                self.send_reject(peer, [0u8; 16], &err.to_string()).await?;
                return Ok(());
            }
        };
        let body = protocol::to_body(&NativeServerHelloBody { hello })?;
        let out = protocol::encode_plain(PacketKind::NativeServerHello, pending_id, 1, 0, &body)?;
        let _ = send_udp(&self.socket, &out, peer).await?;
        Ok(())
    }

    fn build_server_hello(
        &mut self,
        client: native::NativeClientHello,
        peer: SocketAddr,
    ) -> Result<([u8; 16], NativeServerHello)> {
        if !self.config.server.native_auth {
            bail!("native auth disabled");
        }
        if !client
            .supported_aead
            .iter()
            .any(|algorithm| algorithm == "chacha20poly1305")
        {
            bail!("native auth requires chacha20poly1305");
        }
        if !client
            .supported_user_key_algorithms
            .iter()
            .any(|algorithm| native::is_supported_user_signature_algorithm(algorithm))
        {
            bail!("native auth requires a supported user key algorithm");
        }
        if self.config.require_current_user {
            let current_user = std::env::var("USER").unwrap_or_else(|_| "unknown".to_string());
            if client.requested_user != current_user {
                bail!("native auth user mismatch");
            }
        }

        let (server_secret, server_public) = generate_native_ephemeral();
        let mut server = NativeServerHello {
            protocol_version: native::NATIVE_PROTOCOL_VERSION,
            server_random: crypto::random_32(),
            server_ephemeral_public: server_public,
            host_key: host_public_key(&self.host_signing),
            chosen_aead: "chacha20poly1305".to_string(),
            server_key_epoch: 1,
            auth_challenge: crypto::random_32(),
            rate_limit_remaining: None,
            host_signature: Vec::new(),
        };
        sign_server_hello(&self.host_signing, &client, &mut server)?;
        let session_key = derive_native_session_key(
            &server_secret,
            client.client_ephemeral_public,
            &client,
            &server,
        )?;
        let mut pending_id = [0u8; 16];
        pending_id.copy_from_slice(&server.auth_challenge[..16]);
        self.pending.insert(
            pending_id,
            PendingServerAuth {
                client,
                server: server.clone(),
                session_key,
                peer,
                created_at: Instant::now(),
            },
        );
        Ok((pending_id, server))
    }

    async fn handle_user_auth(
        &mut self,
        peer: SocketAddr,
        packet: &protocol::Packet,
    ) -> Result<DoshServerEvent> {
        let Some(pending) = self.pending.get(&packet.header.conn_id).cloned() else {
            self.send_reject(peer, packet.header.conn_id, "unknown native auth")
                .await?;
            return Ok(DoshServerEvent::Ignored);
        };
        if pending.peer != peer {
            self.send_reject(peer, packet.header.conn_id, "native auth peer changed")
                .await?;
            return Ok(DoshServerEvent::Ignored);
        }
        if pending.created_at.elapsed() > self.config.auth_timeout {
            self.pending.remove(&packet.header.conn_id);
            self.send_reject(peer, packet.header.conn_id, "native auth expired")
                .await?;
            return Ok(DoshServerEvent::Ignored);
        }

        let body = match protocol::decrypt_body(packet, &pending.session_key, CLIENT_TO_SERVER) {
            Ok(body) => body,
            Err(_) => {
                self.send_reject(peer, packet.header.conn_id, "native auth decrypt failed")
                    .await?;
                return Ok(DoshServerEvent::Ignored);
            }
        };
        let req: NativeUserAuthBody = match protocol::from_body(&body) {
            Ok(req) => req,
            Err(_) => {
                self.send_reject(peer, packet.header.conn_id, "invalid native auth body")
                    .await?;
                return Ok(DoshServerEvent::Ignored);
            }
        };
        let services = match self.verify_auth_and_services(&pending, &req.auth) {
            Ok(services) => services,
            Err(err) => {
                self.send_reject(peer, packet.header.conn_id, &format!("{err:#}"))
                    .await?;
                return Ok(DoshServerEvent::Ignored);
            }
        };
        self.pending.remove(&packet.header.conn_id);

        let conn_id = crypto::random_16();
        let key_id = protocol::session_key_id(&pending.session_key);
        let ok = NativeAuthOk {
            client_id: conn_id,
            session: pending.client.requested_session.clone(),
            mode: pending.client.requested_mode.clone(),
            session_key: pending.session_key,
            session_key_id: key_id,
            attach_ticket: Vec::new(),
            attach_ticket_psk: crypto::random_32(),
            initial_seq: 1,
            snapshot: Vec::new(),
            policy_flags: services
                .iter()
                .map(|service| format!("service:{service}"))
                .collect(),
        };
        let body = protocol::to_body(&NativeAuthOkBody { ok })?;
        let out = protocol::encode_encrypted(
            PacketKind::NativeAuthOk,
            conn_id,
            1,
            packet.header.seq,
            &pending.session_key,
            SERVER_TO_CLIENT,
            &body,
        )?;
        let _ = send_udp(&self.socket, &out, peer).await?;

        let transport = DoshTransport::new(
            Arc::clone(&self.socket),
            SessionTransportConfig {
                role: SessionRole::Server,
                conn_id,
                session_key: pending.session_key,
                peer_addr: peer,
                initial_send_seq: 2,
                initial_ack: packet.header.seq,
                stream: self.config.transport.clone(),
            },
        );
        let accepted = DoshAccepted {
            conn_id,
            user: pending.client.requested_user,
            session: pending.client.requested_session,
            services,
            peer_addr: peer,
        };
        self.transports.insert(conn_id, transport);
        self.accepted.insert(conn_id, accepted.clone());
        Ok(DoshServerEvent::Accepted(accepted))
    }

    fn verify_auth_and_services(
        &self,
        pending: &PendingServerAuth,
        auth: &native::NativeUserAuth,
    ) -> Result<Vec<String>> {
        verify_native_user_auth_from_config(
            &self.config.server,
            &pending.client,
            &pending.server,
            auth,
            Some(pending.peer.ip()),
        )
        .context("verify native user auth")?;
        validate_requested_services(&self.config.services, &auth.requested_forwardings)
    }

    async fn send_reject(&self, peer: SocketAddr, conn_id: [u8; 16], reason: &str) -> Result<()> {
        let body = protocol::to_body(&AttachReject {
            reason: reason.to_string(),
        })?;
        let out = protocol::encode_plain(PacketKind::AttachReject, conn_id, 0, 0, &body)?;
        let _ = send_udp(&self.socket, &out, peer).await?;
        Ok(())
    }

    fn expire_pending(&mut self) {
        let timeout = self.config.auth_timeout;
        self.pending
            .retain(|_, pending| pending.created_at.elapsed() <= timeout);
    }
}

async fn send_udp(socket: &UdpSocket, packet: &[u8], peer: SocketAddr) -> Result<bool> {
    match socket.send_to(packet, peer).await {
        Ok(_) => Ok(true),
        Err(err) if is_transient_udp_send_error(&err) => Ok(false),
        Err(err) => Err(err).with_context(|| format!("send UDP packet to {peer}")),
    }
}

fn validate_requested_services(
    allowed_services: &HashSet<String>,
    requests: &[ForwardingRequest],
) -> Result<Vec<String>> {
    let mut services = Vec::new();
    for request in requests {
        if request.kind != ForwardingKind::Local {
            bail!("embedded Dosh services only accept local service requests");
        }
        if request.listen_port != 0 || request.target_port != Some(0) {
            bail!("embedded Dosh service requests must use port 0");
        }
        let target = request
            .target_host
            .as_deref()
            .ok_or_else(|| anyhow!("embedded Dosh service request is missing target host"))?;
        let service = service_name_from_target(target)
            .ok_or_else(|| anyhow!("target {target:?} is not a Dosh service"))?;
        if !allowed_services.contains(service) {
            bail!("Dosh service {service:?} is not registered");
        }
        if !services.iter().any(|existing| existing == service) {
            services.push(service.to_string());
        }
    }
    Ok(services)
}

fn validate_service_name(name: String) -> Result<String> {
    crate::transport::service_target(&name)?;
    Ok(name)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::client::DoshClient;
    use crate::config::{ClientConfig, HostsConfig};
    use crate::protocol::{CLIENT_TO_SERVER, NativeUserAuthBody};
    use crate::transport::TransportEvent;
    use ed25519_dalek::SigningKey;

    #[tokio::test]
    async fn sdk_client_and_server_exchange_service_stream() {
        let dir = tempfile::tempdir().unwrap();
        let host_key = dir.path().join("host_key");
        let authorized_keys = dir.path().join("authorized_keys");
        let identity = dir.path().join("id_ed25519");
        let known_hosts = dir.path().join("known_hosts");

        let signing = SigningKey::from_bytes(&[91u8; 32]);
        let keypair = ssh_key::private::Ed25519Keypair::from(&signing);
        let private =
            ssh_key::PrivateKey::new(ssh_key::private::KeypairData::from(keypair), "").unwrap();
        private
            .write_openssh_file(&identity, ssh_key::LineEnding::LF)
            .unwrap();
        std::fs::write(
            &authorized_keys,
            format!("{}\n", private.public_key().to_openssh().unwrap()),
        )
        .unwrap();

        let server_config = ServerConfig {
            host_key: host_key.to_string_lossy().to_string(),
            authorized_keys: vec![authorized_keys.to_string_lossy().to_string()],
            ..ServerConfig::default()
        };
        let server_config = DoshServerConfig::new(server_config)
            .bind_addr("127.0.0.1:0".parse().unwrap())
            .service("echo")
            .unwrap()
            .require_current_user(false);
        let mut server = DoshServer::bind(server_config).await.unwrap();
        let server_port = server.local_addr().unwrap().port();

        let client_config = ClientConfig {
            dosh_port: server_port,
            trust_on_first_use: true,
            known_hosts: known_hosts.to_string_lossy().to_string(),
            identity_files: vec![identity.to_string_lossy().to_string()],
            use_ssh_agent: false,
            ..ClientConfig::default()
        };
        let client = DoshClient::with_config(client_config, HostsConfig::default());
        let connect = client
            .connect("127.0.0.1")
            .user("sdk-user")
            .service("echo")
            .connect();
        let accept = async {
            loop {
                if let DoshServerEvent::Accepted(accepted) = server.recv().await.unwrap() {
                    break accepted;
                }
            }
        };
        let (connected, accepted) = tokio::join!(connect, accept);
        let mut client_transport = connected.unwrap().into_transport();
        let conn_id = accepted.conn_id;
        assert_eq!(accepted.services, vec!["echo".to_string()]);

        let stream_id = client_transport.open_service("echo").await.unwrap();
        match server.recv().await.unwrap() {
            DoshServerEvent::Session {
                event: SessionEvent::Stream(TransportEvent::Open(open)),
                ..
            } => {
                assert_eq!(open.stream_id, stream_id);
                server.accept_stream(conn_id, open.stream_id).await.unwrap();
            }
            other => panic!("unexpected event {other:?}"),
        }
        assert!(matches!(
            client_transport.recv().await.unwrap(),
            SessionEvent::Stream(TransportEvent::OpenOk { .. })
        ));

        client_transport
            .send(stream_id, b"ping".to_vec())
            .await
            .unwrap();
        loop {
            match server.recv().await.unwrap() {
                DoshServerEvent::Session {
                    event: SessionEvent::Stream(TransportEvent::Data(data)),
                    ..
                } => {
                    assert_eq!(data.chunks, vec![b"ping".to_vec()]);
                    server
                        .send(conn_id, data.stream_id, b"pong".to_vec())
                        .await
                        .unwrap();
                    break;
                }
                DoshServerEvent::Session { .. } | DoshServerEvent::Ignored => {}
                other => panic!("unexpected event {other:?}"),
            }
        }
        loop {
            match client_transport.recv().await.unwrap() {
                SessionEvent::Stream(TransportEvent::Data(data)) => {
                    assert_eq!(data.chunks, vec![b"pong".to_vec()]);
                    break;
                }
                SessionEvent::Stream(_) | SessionEvent::Ignored => {}
                other => panic!("unexpected event {other:?}"),
            }
        }
    }

    #[tokio::test]
    async fn bad_native_auth_does_not_consume_pending_challenge() {
        let dir = tempfile::tempdir().unwrap();
        let host_key = dir.path().join("host_key");
        let authorized_keys = dir.path().join("authorized_keys");
        let signing = SigningKey::from_bytes(&[93u8; 32]);
        let keypair = ssh_key::private::Ed25519Keypair::from(&signing);
        let private =
            ssh_key::PrivateKey::new(ssh_key::private::KeypairData::from(keypair), "").unwrap();
        std::fs::write(
            &authorized_keys,
            format!("{}\n", private.public_key().to_openssh().unwrap()),
        )
        .unwrap();

        let server_config = ServerConfig {
            host_key: host_key.to_string_lossy().to_string(),
            authorized_keys: vec![authorized_keys.to_string_lossy().to_string()],
            ..ServerConfig::default()
        };
        let server_config = DoshServerConfig::new(server_config)
            .bind_addr("127.0.0.1:0".parse().unwrap())
            .require_current_user(false);
        let mut server = DoshServer::bind(server_config).await.unwrap();
        let peer: SocketAddr = "127.0.0.1:9".parse().unwrap();
        let (client_secret, client_public) = native::generate_native_ephemeral();
        let hello = native::NativeClientHello {
            protocol_version: native::NATIVE_PROTOCOL_VERSION,
            client_random: crypto::random_32(),
            client_ephemeral_public: client_public,
            requested_host: "127.0.0.1".to_string(),
            requested_user: "sdk-user".to_string(),
            requested_session: "test".to_string(),
            requested_mode: "forward-only".to_string(),
            terminal_size: (80, 24),
            supported_aead: vec!["chacha20poly1305".to_string()],
            supported_user_key_algorithms: vec!["ssh-ed25519".to_string()],
            cached_host_key_fingerprint: None,
            attach_ticket_envelope: None,
            requested_env: Vec::new(),
        };
        let (pending_id, server_hello) = server.build_server_hello(hello.clone(), peer).unwrap();
        let session_key = native::derive_native_session_key(
            &client_secret,
            server_hello.server_ephemeral_public,
            &hello,
            &server_hello,
        )
        .unwrap();
        let auth = native::sign_user_auth(&signing, &hello, &server_hello, Vec::new()).unwrap();
        let body = protocol::to_body(&NativeUserAuthBody { auth }).unwrap();
        let bad = protocol::encode_encrypted(
            PacketKind::NativeUserAuth,
            pending_id,
            2,
            1,
            &[7u8; 32],
            CLIENT_TO_SERVER,
            &body,
        )
        .unwrap();
        let bad = protocol::decode(&bad).unwrap();

        assert!(matches!(
            server.handle_user_auth(peer, &bad).await.unwrap(),
            DoshServerEvent::Ignored
        ));
        assert!(server.pending.contains_key(&pending_id));

        let valid = protocol::encode_encrypted(
            PacketKind::NativeUserAuth,
            pending_id,
            2,
            1,
            &session_key,
            CLIENT_TO_SERVER,
            &body,
        )
        .unwrap();
        let valid = protocol::decode(&valid).unwrap();
        assert!(matches!(
            server.handle_user_auth(peer, &valid).await.unwrap(),
            DoshServerEvent::Accepted(_)
        ));
        assert!(!server.pending.contains_key(&pending_id));
    }
}