pub mod connection;
pub mod session;

use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;

use futures::SinkExt;
use tokio::net::TcpListener;
use tokio::sync::Semaphore;
use tokio_util::codec::Framed;
use tracing::{self, Instrument};

use crate::auth::Authenticator;
use crate::codec::frame::ErrorCode;
use crate::codec::{Frame, MessageCodec};
use crate::config::{DeployConfig, SharedConfig, TcpConfig};
use crate::error::Result;
use crate::namespace::NamespaceManager;
use crate::server::state::{ServerLifecycle, ServerState};
use crate::tls::TlsProvider;
use crate::types::ClientId;

use connection::{ConnectionConfig, run_connection};

/// TCP server that accepts client connections and dispatches them to
/// per-connection handlers.
pub struct TcpServer<M: NamespaceManager> {
    config: TcpConfig,
    deploy_config: DeployConfig,
    shared_config: SharedConfig,
    manager: Arc<M>,
    auth: Arc<dyn Authenticator>,
    client_id_counter: Arc<AtomicU64>,
    message_id_counter: Arc<AtomicU64>,
    lifecycle: Arc<ServerLifecycle>,
    tls_provider: Option<Arc<TlsProvider>>,
}

impl<M: NamespaceManager + 'static> TcpServer<M> {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        config: TcpConfig,
        deploy_config: DeployConfig,
        shared_config: SharedConfig,
        manager: Arc<M>,
        auth: Arc<dyn Authenticator>,
        message_id_counter: Arc<AtomicU64>,
        lifecycle: Arc<ServerLifecycle>,
        tls_provider: Option<Arc<TlsProvider>>,
    ) -> Self {
        Self {
            config,
            deploy_config,
            shared_config,
            manager,
            auth,
            client_id_counter: Arc::new(AtomicU64::new(1)),
            message_id_counter,
            lifecycle,
            tls_provider,
        }
    }

    /// Run the accept loop with a pre-bound listener until shutdown.
    ///
    /// Accepts a `TcpListener` directly so that callers (especially tests) can
    /// bind to port 0 and pass the listener without a race window.
    pub async fn run_with_listener(&self, listener: TcpListener) -> Result<()> {
        tracing::info!(
            address = %listener.local_addr().map(|a| a.to_string()).unwrap_or_default(),
            tls = self.tls_provider.is_some(),
            "TCP server listening"
        );
        self.accept_loop(listener).await
    }

    /// Bind to the configured address/port and run the accept loop.
    pub async fn run(&self) -> Result<()> {
        let addr = format!("{}:{}", self.config.bind_address, self.config.port);
        let listener = TcpListener::bind(&addr).await?;
        tracing::info!(address = %addr, "TCP server listening");
        if self.tls_provider.is_some() {
            tracing::info!("TLS enabled for TCP listener");
        } else {
            tracing::warn!(
                "TLS is disabled — connections are unencrypted. \
                 Set [tls] enabled=true with cert_file/key_file for production use."
            );
        }
        self.accept_loop(listener).await
    }

    async fn accept_loop(&self, listener: TcpListener) -> Result<()> {
        let mut lifecycle_rx = self.lifecycle.subscribe();

        // Wait for initialization before accepting connections.
        while !self.lifecycle.is_initialized() {
            tokio::select! {
                _ = lifecycle_rx.changed() => {
                    if self.lifecycle.is_draining() {
                        tracing::info!("TCP server shutting down before initialization");
                        return Ok(());
                    }
                }
                _ = tokio::time::sleep(Duration::from_millis(50)) => {}
            }
        }

        // Connection limit semaphore. max_connections == 0 means unlimited.
        // Read from SharedConfig for hot-reload support.
        let max_connections = self.shared_config.load().max_connections;
        let conn_semaphore = if max_connections > 0 {
            Some(Arc::new(Semaphore::new(max_connections)))
        } else {
            None
        };

        // Flapping detection: track connections per peer IP in a rolling window.
        // Reset every 60s. Warn when a single IP connects > 50 times in a window.
        let mut peer_connect_counts: std::collections::HashMap<std::net::IpAddr, u32> =
            std::collections::HashMap::new();
        let mut flap_window_reset = tokio::time::Instant::now() + Duration::from_secs(60);

        loop {
            // Stop accepting when draining.
            if self.lifecycle.is_draining() {
                tracing::info!("TCP server stopping accept loop (draining)");
                break;
            }

            // Reset flapping counters each window
            if tokio::time::Instant::now() >= flap_window_reset {
                peer_connect_counts.clear();
                flap_window_reset = tokio::time::Instant::now() + Duration::from_secs(60);
            }

            tokio::select! {
                result = listener.accept() => {
                    match result {
                        Ok((stream, peer_addr)) => {
                            // Track connection flapping
                            let count = peer_connect_counts
                                .entry(peer_addr.ip())
                                .or_insert(0);
                            *count += 1;
                            if *count == 50 {
                                tracing::warn!(
                                    peer_ip = %peer_addr.ip(),
                                    connections_in_window = *count,
                                    "possible connection flapping: rapid reconnects from same peer"
                                );
                            }
                            // Acquire a permit if connection limit is set.
                            let permit = if let Some(ref sem) = conn_semaphore {
                                match sem.clone().try_acquire_owned() {
                                    Ok(permit) => Some(permit),
                                    Err(_) => {
                                        tracing::warn!(
                                            peer = %peer_addr,
                                            max = self.config.max_connections,
                                            "connection limit reached, rejecting"
                                        );
                                        // Send an error frame before closing so the
                                        // client knows why it was rejected.
                                        let mut framed = Framed::new(
                                            stream,
                                            MessageCodec::new(self.config.max_message_size),
                                        );
                                        let _ = framed.send(Frame::Error {
                                            code: ErrorCode::InternalError,
                                            message: "server at capacity".into(),
                                        }).await;
                                        continue;
                                    }
                                }
                            } else {
                                None
                            };

                            let client_id = ClientId(
                                self.client_id_counter.fetch_add(1, Ordering::Relaxed),
                            );
                            tracing::debug!(
                                client = %client_id, peer = %peer_addr,
                                "accepted connection"
                            );
                            let conn_config = ConnectionConfig {
                                client_id,
                                auth: Arc::clone(&self.auth),
                                message_id_counter: Arc::clone(&self.message_id_counter),
                                max_frame_size: self.config.max_message_size,
                                shared_config: Arc::clone(&self.shared_config),
                                lifecycle: Arc::clone(&self.lifecycle),
                                deploy: self.deploy_config.clone(),
                            };
                            let manager = Arc::clone(&self.manager);
                            let lifecycle = Arc::clone(&self.lifecycle);
                            let tls = self.tls_provider.clone();
                            let span = tracing::info_span!(
                                "conn",
                                client = %client_id,
                                peer = %peer_addr,
                                tls = tracing::field::Empty,
                                tls.version = tracing::field::Empty,
                                tls.cipher = tracing::field::Empty,
                                cert.cn = tracing::field::Empty,
                            );
                            tokio::spawn(async move {
                                match tls {
                                    Some(tls_provider) => {
                                        let handshake_timeout = tls_provider.handshake_timeout();
                                        let start = std::time::Instant::now();
                                        match tokio::time::timeout(
                                            handshake_timeout,
                                            tls_provider.accept(stream),
                                        )
                                        .await
                                        {
                                            Ok(Ok(tls_stream)) => {
                                                let elapsed = start.elapsed();
                                                tracing::Span::current().record("tls", true);

                                                // Extract TLS version and cipher suite from the negotiated connection.
                                                let (_, server_conn) = tls_stream.get_ref();
                                                if let Some(version) = server_conn.protocol_version() {
                                                    let version_str = match version {
                                                        rustls::ProtocolVersion::TLSv1_2 => "TLSv1.2",
                                                        rustls::ProtocolVersion::TLSv1_3 => "TLSv1.3",
                                                        other => {
                                                            tracing::debug!(version = ?other, "unknown TLS version");
                                                            "unknown"
                                                        }
                                                    };
                                                    tracing::Span::current().record("tls.version", version_str);
                                                }
                                                if let Some(cipher) = server_conn.negotiated_cipher_suite() {
                                                    tracing::Span::current().record(
                                                        "tls.cipher",
                                                        tracing::field::debug(cipher.suite()),
                                                    );
                                                }

                                                let cert_info = TlsProvider::client_info(&tls_stream);
                                                if let Some(ref info) = cert_info {
                                                    tracing::Span::current().record(
                                                        "cert.cn",
                                                        info.common_name.as_deref().unwrap_or(""),
                                                    );
                                                    tracing::debug!(
                                                        cert.cn = info.common_name.as_deref().unwrap_or(""),
                                                        cert.serial = %info.serial,
                                                        cert.fingerprint = %info.fingerprint_sha256,
                                                        "mTLS client certificate verified"
                                                    );
                                                }

                                                tracing::debug!(
                                                    elapsed_ms = elapsed.as_millis() as u64,
                                                    "TLS handshake completed"
                                                );
                                                metrics::counter!("fafb_tls_handshake_total", "status" => "success").increment(1);
                                                metrics::histogram!("fafb_tls_handshake_duration_seconds").record(elapsed.as_secs_f64());
                                                lifecycle.inc_connections();
                                                run_connection(tls_stream, conn_config, manager).await;
                                                lifecycle.dec_connections();
                                            }
                                            Ok(Err(e)) => {
                                                tracing::debug!(
                                                    peer = %peer_addr,
                                                    error = %e,
                                                    "TLS handshake failed"
                                                );
                                                metrics::counter!("fafb_tls_handshake_total", "status" => "error").increment(1);
                                                metrics::histogram!("fafb_tls_handshake_duration_seconds").record(start.elapsed().as_secs_f64());
                                            }
                                            Err(_) => {
                                                tracing::debug!(
                                                    peer = %peer_addr,
                                                    timeout_secs = handshake_timeout.as_secs(),
                                                    "TLS handshake timed out"
                                                );
                                                metrics::counter!("fafb_tls_handshake_total", "status" => "timeout").increment(1);
                                                metrics::counter!("fafb_tls_handshake_timeout_total").increment(1);
                                            }
                                        }
                                    }
                                    None => {
                                        tracing::Span::current().record("tls", false);
                                        lifecycle.inc_connections();
                                        run_connection(stream, conn_config, manager).await;
                                        lifecycle.dec_connections();
                                    }
                                }
                                // Permit is dropped here, releasing the semaphore slot.
                                drop(permit);
                            }.instrument(span));
                        }
                        Err(e) => {
                            tracing::warn!(error = %e, "failed to accept connection");
                        }
                    }
                }
                _ = lifecycle_rx.changed() => {
                    let state = *lifecycle_rx.borrow();
                    if state >= ServerState::Draining {
                        tracing::info!("TCP server shutting down");
                        break;
                    }
                }
            }
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::namespace::DefaultNamespaceManager;
    use crate::queue::InMemoryQueueEngine;
    use crate::storage::mock::MockStorage;

    #[test]
    fn server_is_send() {
        fn assert_send<T: Send>() {}
        assert_send::<TcpServer<DefaultNamespaceManager<InMemoryQueueEngine<MockStorage>>>>();
    }
}
