use std::collections::HashMap;
use std::future::Future;
use std::net::SocketAddr;
use std::path::Path;
use std::pin::Pin;
use std::sync::{Arc, atomic::AtomicUsize};
use std::time::Duration;

use anyhow::{Context, Result, anyhow};
use russh::server::{self, Msg, Session};
use russh::{Channel, ChannelId, Disconnect};
use socket2::{SockRef, TcpKeepalive};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{OwnedSemaphorePermit, Semaphore, broadcast, watch};
use tokio_rustls::TlsAcceptor;
use tracing::{debug, info, warn};

use rustls::ServerConfig;

use crate::rtr::cache::SharedRtrCache;
use crate::rtr::server::config::RtrServiceConfig;
use crate::rtr::server::connection::{
    ConnectionGuard, handle_tcp_connection, handle_tls_connection, is_expected_disconnect,
};
use crate::rtr::server::ssh::RtrSshRuntimeConfig;
use crate::rtr::server::tls::load_rustls_server_config_with_options;
use crate::rtr::session::RtrSession;

type TransportFuture = Pin<Box<dyn Future<Output = Result<()>> + Send>>;

pub trait TransportAcceptor: Clone + Send + Sync + 'static {
    fn name(&self) -> &'static str;
    fn requires_handshake_limit(&self) -> bool {
        false
    }

    fn handle_connection(
        &self,
        cache: SharedRtrCache,
        stream: TcpStream,
        peer_addr: SocketAddr,
        notify_tx: broadcast::Sender<()>,
        shutdown_tx: watch::Sender<bool>,
        handshake_permit: Option<OwnedSemaphorePermit>,
    ) -> TransportFuture;
}

#[derive(Clone)]
struct TcpTransport;

impl TransportAcceptor for TcpTransport {
    fn name(&self) -> &'static str {
        "TCP"
    }

    fn handle_connection(
        &self,
        cache: SharedRtrCache,
        stream: TcpStream,
        peer_addr: SocketAddr,
        notify_tx: broadcast::Sender<()>,
        shutdown_tx: watch::Sender<bool>,
        _handshake_permit: Option<OwnedSemaphorePermit>,
    ) -> TransportFuture {
        Box::pin(async move {
            handle_tcp_connection(
                cache,
                stream,
                peer_addr,
                notify_tx.subscribe(),
                shutdown_tx.subscribe(),
            )
            .await
        })
    }
}

#[derive(Clone)]
struct TlsTransport {
    acceptor: TlsAcceptor,
    enforce_client_san_ip_match: bool,
}

impl TransportAcceptor for TlsTransport {
    fn name(&self) -> &'static str {
        "TLS"
    }
    fn requires_handshake_limit(&self) -> bool {
        true
    }

    fn handle_connection(
        &self,
        cache: SharedRtrCache,
        stream: TcpStream,
        peer_addr: SocketAddr,
        notify_tx: broadcast::Sender<()>,
        shutdown_tx: watch::Sender<bool>,
        handshake_permit: Option<OwnedSemaphorePermit>,
    ) -> TransportFuture {
        let acceptor = self.acceptor.clone();
        let enforce_client_san_ip_match = self.enforce_client_san_ip_match;
        Box::pin(async move {
            handle_tls_connection(
                cache,
                stream,
                peer_addr,
                acceptor,
                enforce_client_san_ip_match,
                notify_tx.subscribe(),
                shutdown_tx.subscribe(),
                handshake_permit,
            )
            .await
        })
    }
}

#[derive(Clone)]
struct SshTransport {
    runtime: Arc<RtrSshRuntimeConfig>,
}

impl TransportAcceptor for SshTransport {
    fn name(&self) -> &'static str {
        "SSH"
    }

    fn handle_connection(
        &self,
        cache: SharedRtrCache,
        stream: TcpStream,
        peer_addr: SocketAddr,
        notify_tx: broadcast::Sender<()>,
        shutdown_tx: watch::Sender<bool>,
        _handshake_permit: Option<OwnedSemaphorePermit>,
    ) -> TransportFuture {
        let runtime = self.runtime.clone();
        Box::pin(async move {
            let handler = RtrSshHandler::new(
                cache,
                notify_tx.subscribe(),
                shutdown_tx.subscribe(),
                peer_addr,
                runtime.authorized_keys.clone(),
                runtime.username.clone(),
                runtime.subsystem_name.clone(),
                runtime.password.clone(),
            );

            let running = server::run_stream(runtime.server_config.clone(), stream, handler)
                .await
                .with_context(|| format!("failed to start SSH session for {}", peer_addr))?;

            let handle = running.handle();
            let mut connection_shutdown_rx = shutdown_tx.subscribe();

            tokio::select! {
                session_res = running => {
                    session_res.map_err(|err| anyhow!(err))
                }
                changed = connection_shutdown_rx.changed() => {
                    match changed {
                        Ok(()) if *connection_shutdown_rx.borrow() => {
                            let _ = handle
                                .disconnect(
                                    Disconnect::ByApplication,
                                    "service shutdown".to_string(),
                                    "".to_string(),
                                )
                                .await;
                        }
                        Ok(()) | Err(_) => {}
                    }
                    Ok(())
                }
            }
        })
    }
}

pub struct RtrServer {
    bind_addr: SocketAddr,
    cache: SharedRtrCache,
    notify_tx: broadcast::Sender<()>,
    shutdown_tx: watch::Sender<bool>,
    connection_limiter: Arc<Semaphore>,
    handshake_limiter: Arc<Semaphore>,
    active_connections: Arc<AtomicUsize>,
    config: RtrServiceConfig,
}

impl RtrServer {
    pub fn new(
        bind_addr: SocketAddr,
        cache: SharedRtrCache,
        notify_tx: broadcast::Sender<()>,
        shutdown_tx: watch::Sender<bool>,
        connection_limiter: Arc<Semaphore>,
        handshake_limiter: Arc<Semaphore>,
        active_connections: Arc<AtomicUsize>,
        config: RtrServiceConfig,
    ) -> Self {
        Self {
            bind_addr,
            cache,
            notify_tx,
            shutdown_tx,
            connection_limiter,
            handshake_limiter,
            active_connections,
            config,
        }
    }

    pub fn bind_addr(&self) -> SocketAddr {
        self.bind_addr
    }

    pub fn cache(&self) -> SharedRtrCache {
        self.cache.clone()
    }

    pub fn active_connections(&self) -> usize {
        self.active_connections
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    pub async fn run_tcp(self) -> Result<()> {
        self.run_with_transport(TcpTransport).await
    }

    pub async fn run_tls_from_pem(
        self,
        cert_path: impl AsRef<Path>,
        key_path: impl AsRef<Path>,
        client_ca_path: impl AsRef<Path>,
    ) -> Result<()> {
        let tls_config = Arc::new(load_rustls_server_config_with_options(
            cert_path,
            key_path,
            client_ca_path,
            self.config.require_tls_server_dns_name_san,
        )?);
        self.run_tls(tls_config).await
    }

    pub async fn run_tls(self, tls_config: Arc<ServerConfig>) -> Result<()> {
        let transport = TlsTransport {
            acceptor: TlsAcceptor::from(tls_config),
            enforce_client_san_ip_match: self.config.enforce_tls_client_san_ip_match,
        };
        self.run_with_transport(transport).await
    }

    pub async fn run_ssh(self, runtime_config: Arc<RtrSshRuntimeConfig>) -> Result<()> {
        let transport = SshTransport {
            runtime: runtime_config,
        };
        self.run_with_transport(transport).await
    }

    async fn run_with_transport<T>(self, transport: T) -> Result<()>
    where
        T: TransportAcceptor,
    {
        let listener = TcpListener::bind(self.bind_addr).await.with_context(|| {
            format!(
                "failed to bind {} RTR server on {}",
                transport.name(),
                self.bind_addr
            )
        })?;
        let mut shutdown_rx = self.shutdown_tx.subscribe();

        info!(
            "RTR {} server listening on {}",
            transport.name(),
            self.bind_addr
        );

        loop {
            tokio::select! {
                changed = shutdown_rx.changed() => {
                    match changed {
                        Ok(()) => {
                            if *shutdown_rx.borrow() {
                                info!(
                                    "RTR {} listener {} shutting down",
                                    transport.name(),
                                    self.bind_addr
                                );
                                return Ok(());
                            }
                        }
                        Err(_) => {
                            info!(
                                "RTR {} listener {} shutdown channel closed",
                                transport.name(),
                                self.bind_addr
                            );
                            return Ok(());
                        }
                    }
                }

                accept_res = listener.accept() => {
                    let (stream, peer_addr) = match accept_res {
                        Ok(v) => v,
                        Err(err) => {
                            warn!("RTR {} accept failed: {}", transport.name(), err);
                            continue;
                        }
                    };

                    if let Err(err) = apply_keepalive(&stream, self.config.tcp_keepalive) {
                        warn!(
                            "failed to configure TCP keepalive for {} peer {}: {}",
                            transport.name(),
                            peer_addr,
                            err
                        );
                    }

                    let permit = match self.connection_limiter.clone().try_acquire_owned() {
                        Ok(permit) => permit,
                        Err(_) => {
                            warn!(
                                "RTR {} connection rejected for {}: max connections reached ({})",
                                transport.name(),
                                peer_addr,
                                self.config.max_connections
                            );
                            drop(stream);
                            continue;
                        }
                    };
                    let handshake_permit = if transport.requires_handshake_limit() {
                        match self.handshake_limiter.clone().try_acquire_owned() {
                            Ok(permit) => Some(permit),
                            Err(_) => {
                                warn!(
                                    "RTR {} connection rejected for {}: max concurrent handshakes reached ({})",
                                    transport.name(),
                                    peer_addr,
                                    self.config.max_concurrent_handshakes
                                );
                                drop(stream);
                                drop(permit);
                                continue;
                            }
                        }
                    } else {
                        None
                    };

                    let cache = self.cache.clone();
                    let notify_tx = self.notify_tx.clone();
                    let shutdown_tx = self.shutdown_tx.clone();
                    let active_connections = self.active_connections.clone();
                    let transport_instance = transport.clone();

                    debug!(
                        "RTR {} client connected: peer_addr={}, active_connections(before_spawn)={}",
                        transport_instance.name(),
                        peer_addr,
                        self.active_connections()
                    );

                    tokio::spawn(async move {
                        let guard = ConnectionGuard::new(active_connections, permit);
                        info!(
                            "RTR {} connection established: peer_addr={}, active_connections={}",
                            transport_instance.name(),
                            peer_addr,
                            guard.active_count()
                        );

                        if let Err(err) = transport_instance
                            .handle_connection(
                                cache,
                                stream,
                                peer_addr,
                                notify_tx,
                                shutdown_tx,
                                handshake_permit,
                            )
                            .await
                        {
                            let active_after_close = guard.active_count().saturating_sub(1);
                            if is_expected_disconnect(&err) {
                                info!(
                                    "RTR {} session closed by peer: peer_addr={}, active_connections={}, err={}",
                                    transport_instance.name(),
                                    peer_addr,
                                    active_after_close,
                                    err
                                );
                            } else {
                                warn!(
                                    "RTR {} session closed with error: peer_addr={}, active_connections={}, err={}",
                                    transport_instance.name(),
                                    peer_addr,
                                    active_after_close,
                                    err
                                );
                            }
                        } else {
                            let active_after_close = guard.active_count().saturating_sub(1);
                            info!(
                                "RTR {} session closed cleanly: peer_addr={}, active_connections={}",
                                transport_instance.name(),
                                peer_addr,
                                active_after_close
                            );
                        }
                    });
                }
            }
        }
    }
}

struct RtrSshHandler {
    cache: SharedRtrCache,
    notify_rx: broadcast::Receiver<()>,
    shutdown_rx: watch::Receiver<bool>,
    peer_addr: SocketAddr,
    authorized_keys: Arc<Vec<russh::keys::ssh_key::PublicKey>>,
    username: Arc<str>,
    subsystem_name: Arc<str>,
    password: Option<Arc<str>>,
    channels: HashMap<ChannelId, Channel<Msg>>,
    subsystem_started: bool,
}

impl RtrSshHandler {
    fn new(
        cache: SharedRtrCache,
        notify_rx: broadcast::Receiver<()>,
        shutdown_rx: watch::Receiver<bool>,
        peer_addr: SocketAddr,
        authorized_keys: Arc<Vec<russh::keys::ssh_key::PublicKey>>,
        username: Arc<str>,
        subsystem_name: Arc<str>,
        password: Option<Arc<str>>,
    ) -> Self {
        Self {
            cache,
            notify_rx,
            shutdown_rx,
            peer_addr,
            authorized_keys,
            username,
            subsystem_name,
            password,
            channels: HashMap::new(),
            subsystem_started: false,
        }
    }

    fn is_authorized_key(&self, key: &russh::keys::ssh_key::PublicKey) -> bool {
        self.authorized_keys
            .iter()
            .any(|allowed| allowed.key_data() == key.key_data())
    }

    fn is_expected_user(&self, user: &str) -> bool {
        user == self.username.as_ref()
    }
}

impl server::Handler for RtrSshHandler {
    type Error = anyhow::Error;

    async fn auth_none(&mut self, _user: &str) -> Result<server::Auth, Self::Error> {
        Ok(server::Auth::reject())
    }

    async fn auth_password(
        &mut self,
        user: &str,
        password: &str,
    ) -> Result<server::Auth, Self::Error> {
        let accepted = self.is_expected_user(user)
            && self
                .password
                .as_deref()
                .map(|expected| expected == password)
                .unwrap_or(false);
        if accepted {
            info!(
                "RTR SSH password auth accepted: peer_addr={}, user={}",
                self.peer_addr, user
            );
            Ok(server::Auth::Accept)
        } else {
            warn!(
                "RTR SSH password auth rejected: peer_addr={}, user={}",
                self.peer_addr, user
            );
            Ok(server::Auth::reject())
        }
    }

    async fn auth_publickey_offered(
        &mut self,
        user: &str,
        public_key: &russh::keys::ssh_key::PublicKey,
    ) -> Result<server::Auth, Self::Error> {
        if self.is_expected_user(user) && self.is_authorized_key(public_key) {
            Ok(server::Auth::Accept)
        } else {
            Ok(server::Auth::reject())
        }
    }

    async fn auth_publickey(
        &mut self,
        user: &str,
        public_key: &russh::keys::ssh_key::PublicKey,
    ) -> Result<server::Auth, Self::Error> {
        if self.is_expected_user(user) && self.is_authorized_key(public_key) {
            info!(
                "RTR SSH publickey auth accepted: peer_addr={}, user={}",
                self.peer_addr, user
            );
            Ok(server::Auth::Accept)
        } else {
            warn!(
                "RTR SSH publickey auth rejected: peer_addr={}, user={}",
                self.peer_addr, user
            );
            Ok(server::Auth::reject())
        }
    }

    async fn channel_open_session(
        &mut self,
        channel: Channel<Msg>,
        _session: &mut Session,
    ) -> Result<bool, Self::Error> {
        if self.subsystem_started {
            return Ok(false);
        }
        self.channels.insert(channel.id(), channel);
        Ok(true)
    }

    async fn subsystem_request(
        &mut self,
        channel: ChannelId,
        name: &str,
        session: &mut Session,
    ) -> Result<(), Self::Error> {
        if name != self.subsystem_name.as_ref() {
            let _ = session.channel_failure(channel);
            warn!(
                "RTR SSH subsystem rejected: peer_addr={}, requested={}, expected={}",
                self.peer_addr, name, self.subsystem_name
            );
            return Ok(());
        }

        let Some(channel) = self.channels.remove(&channel) else {
            let _ = session.channel_failure(channel);
            return Ok(());
        };

        session.channel_success(channel.id())?;
        self.subsystem_started = true;

        let cache = self.cache.clone();
        let notify_rx = self.notify_rx.resubscribe();
        let shutdown_rx = self.shutdown_rx.clone();
        let peer_addr = self.peer_addr;
        let subsystem = self.subsystem_name.clone();

        tokio::spawn(async move {
            let stream = channel.into_stream();
            let session = RtrSession::new(cache, stream, notify_rx, shutdown_rx);
            if let Err(err) = session.run().await {
                warn!(
                    "RTR SSH subsystem session closed with error: peer_addr={}, subsystem={}, err={}",
                    peer_addr, subsystem, err
                );
            } else {
                info!(
                    "RTR SSH subsystem session completed: peer_addr={}, subsystem={}",
                    peer_addr, subsystem
                );
            }
        });

        Ok(())
    }

    async fn shell_request(
        &mut self,
        channel: ChannelId,
        session: &mut Session,
    ) -> Result<(), Self::Error> {
        let _ = session.channel_failure(channel);
        Ok(())
    }

    async fn exec_request(
        &mut self,
        channel: ChannelId,
        _data: &[u8],
        session: &mut Session,
    ) -> Result<(), Self::Error> {
        let _ = session.channel_failure(channel);
        Ok(())
    }

    async fn pty_request(
        &mut self,
        channel: ChannelId,
        _term: &str,
        _col_width: u32,
        _row_height: u32,
        _pix_width: u32,
        _pix_height: u32,
        _modes: &[(russh::Pty, u32)],
        session: &mut Session,
    ) -> Result<(), Self::Error> {
        let _ = session.channel_failure(channel);
        Ok(())
    }
}

fn apply_keepalive(stream: &tokio::net::TcpStream, keepalive: Option<Duration>) -> Result<()> {
    let Some(keepalive) = keepalive else {
        return Ok(());
    };

    let socket = SockRef::from(stream);
    let keepalive = TcpKeepalive::new().with_time(keepalive);
    socket
        .set_tcp_keepalive(&keepalive)
        .context("unable to apply TCP keepalive settings")?;
    Ok(())
}
