use std::net::{IpAddr, SocketAddr};
use std::sync::{
    Arc,
    atomic::{AtomicUsize, Ordering},
};

use anyhow::{Context, Result, anyhow};
use tokio::net::TcpStream;
use tokio::sync::{OwnedSemaphorePermit, broadcast, watch};
use tracing::{info, warn};
use x509_parser::extensions::GeneralName;
use x509_parser::prelude::{FromDer, X509Certificate};

use tokio_rustls::TlsAcceptor;

use crate::rtr::cache::SharedRtrCache;
use crate::rtr::session::RtrSession;

pub struct ConnectionGuard {
    active_connections: Arc<AtomicUsize>,
    _permit: OwnedSemaphorePermit,
}

impl ConnectionGuard {
    pub fn new(active_connections: Arc<AtomicUsize>, permit: OwnedSemaphorePermit) -> Self {
        active_connections.fetch_add(1, Ordering::Relaxed);
        Self {
            active_connections,
            _permit: permit,
        }
    }

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

impl Drop for ConnectionGuard {
    fn drop(&mut self) {
        self.active_connections.fetch_sub(1, Ordering::Relaxed);
    }
}

pub async fn handle_tcp_connection(
    cache: SharedRtrCache,
    stream: TcpStream,
    peer_addr: SocketAddr,
    notify_rx: broadcast::Receiver<()>,
    shutdown_rx: watch::Receiver<bool>,
) -> Result<()> {
    let session = RtrSession::new(cache, stream, notify_rx, shutdown_rx);
    session.run().await?;

    info!("RTR TCP session completed normally for {}", peer_addr);
    Ok(())
}

pub async fn handle_tls_connection(
    cache: SharedRtrCache,
    stream: TcpStream,
    peer_addr: SocketAddr,
    acceptor: TlsAcceptor,
    enforce_client_san_ip_match: bool,
    notify_rx: broadcast::Receiver<()>,
    shutdown_rx: watch::Receiver<bool>,
    handshake_permit: Option<OwnedSemaphorePermit>,
) -> Result<()> {
    info!("RTR TLS handshake started for {}", peer_addr);
    let tls_stream = acceptor
        .accept(stream)
        .await
        .with_context(|| format!("TLS handshake failed for {}", peer_addr))?;
    info!("RTR TLS handshake completed for {}", peer_addr);
    match verify_peer_certificate_ip(&tls_stream, peer_addr.ip()) {
        Ok(()) => info!("RTR TLS client certificate SAN IP validated for {}", peer_addr),
        Err(err) => {
            if enforce_client_san_ip_match {
                return Err(err).with_context(|| {
                    format!(
                        "TLS client certificate SAN IP validation failed for {}",
                        peer_addr
                    )
                });
            }
            warn!(
                "RTR TLS client certificate SAN IP validation failed but allowed by configuration for {}: {}",
                peer_addr, err
            );
        }
    }
    drop(handshake_permit);

    let session = RtrSession::new(cache, tls_stream, notify_rx, shutdown_rx);
    session.run().await?;

    info!("RTR TLS session completed normally for {}", peer_addr);
    Ok(())
}

pub fn is_expected_disconnect(err: &anyhow::Error) -> bool {
    err.chain().any(|cause| {
        if let Some(io_err) = cause.downcast_ref::<std::io::Error>() {
            return matches!(
                io_err.kind(),
                std::io::ErrorKind::BrokenPipe
                    | std::io::ErrorKind::UnexpectedEof
                    | std::io::ErrorKind::ConnectionAborted
                    | std::io::ErrorKind::ConnectionReset
                    | std::io::ErrorKind::NotConnected
            ) || io_err.raw_os_error() == Some(10054);
        }
        false
    })
}

fn verify_peer_certificate_ip(
    tls_stream: &tokio_rustls::server::TlsStream<TcpStream>,
    peer_ip: IpAddr,
) -> Result<()> {
    let (_, server_connection) = tls_stream.get_ref();
    let peer_certs = server_connection
        .peer_certificates()
        .ok_or_else(|| anyhow!("missing peer certificate after TLS client authentication"))?;
    let end_entity = peer_certs
        .first()
        .ok_or_else(|| anyhow!("peer did not present an end-entity certificate"))?;

    let (_, cert) = X509Certificate::from_der(end_entity.as_ref())
        .map_err(|err| anyhow!("failed to parse peer certificate: {:?}", err))?;
    let san = cert
        .subject_alternative_name()
        .map_err(|err| anyhow!("failed to parse peer certificate SAN: {:?}", err))?
        .ok_or_else(|| anyhow!("peer certificate is missing subjectAltName"))?;

    let matched = san.value.general_names.iter().any(|name| match name {
        GeneralName::IPAddress(bytes) => {
            let bytes = *bytes;
            match (peer_ip, bytes.len()) {
                (IpAddr::V4(ip), 4) => <[u8; 4]>::try_from(bytes)
                    .map(IpAddr::from)
                    .map(|cert_ip| cert_ip == IpAddr::V4(ip))
                    .unwrap_or(false),
                (IpAddr::V6(ip), 16) => <[u8; 16]>::try_from(bytes)
                    .map(IpAddr::from)
                    .map(|cert_ip| cert_ip == IpAddr::V6(ip))
                    .unwrap_or(false),
                _ => false,
            }
        }
        _ => false,
    });

    if matched {
        Ok(())
    } else {
        warn!(
            "RTR TLS client certificate SAN IP mismatch for peer_ip={}",
            peer_ip
        );
        Err(anyhow!(
            "peer certificate subjectAltName iPAddress does not match {}",
            peer_ip
        ))
    }
}
