use std::sync::Arc;
use std::sync::atomic::AtomicU64;
use std::time::Duration;

use futures::SinkExt;
use futures::StreamExt;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::sync::mpsc;
use tokio_util::codec::Framed;
use tracing;

use crate::auth::Authenticator;
use crate::codec::frame::{
    CAP_GOING_AWAY, CAP_HEADERS, CAP_MESSAGE_TTL, ErrorCode, GoingAwayReason, HeaderLimits,
};
use crate::codec::{Frame, MessageCodec};
use crate::config::{DeployConfig, SharedConfig};
use crate::namespace::NamespaceManager;
use crate::server::state::{ServerLifecycle, ServerState};
use crate::types::{ClientId, PublishedMessage};

use super::session::{FrameOutcome, Session};

/// Configuration for a single connection handler, bundled to avoid
/// too many function parameters.
pub(crate) struct ConnectionConfig {
    pub client_id: ClientId,
    pub auth: Arc<dyn Authenticator>,
    pub message_id_counter: Arc<AtomicU64>,
    pub max_frame_size: u32,
    pub shared_config: SharedConfig,
    pub lifecycle: Arc<ServerLifecycle>,
    pub deploy: DeployConfig,
}

/// Run a single client connection to completion.
///
/// Reads frames from the TCP stream, dispatches them through the session
/// state machine, and writes replies back. Also multiplexes delivered
/// messages from all subscriptions into the outbound stream.
///
/// `idle_timeout` of `None` means no timeout. The connection watches the
/// server lifecycle for drain/shutdown transitions.
pub(crate) async fn run_connection<S, M>(stream: S, config: ConnectionConfig, manager: Arc<M>)
where
    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
    M: NamespaceManager + 'static,
{
    let ConnectionConfig {
        client_id,
        auth,
        message_id_counter,
        max_frame_size,
        shared_config,
        lifecycle,
        deploy,
    } = config;

    /// Read the current idle timeout from SharedConfig. Returns None if
    /// idle_timeout_secs is 0 (disabled).
    fn current_idle_timeout(cfg: &SharedConfig) -> Option<Duration> {
        let secs = cfg.load().idle_timeout_secs;
        if secs > 0 {
            Some(Duration::from_secs(secs))
        } else {
            None
        }
    }

    let mut framed = Framed::new(stream, MessageCodec::new(max_frame_size));
    let (delivery_tx, mut delivery_rx) = mpsc::channel::<PublishedMessage>(256);
    let mut session = Session::new(
        client_id,
        message_id_counter,
        max_frame_size,
        Arc::clone(&shared_config),
    );

    tracing::info!(client = %client_id, "connection opened");

    // Idle timer: dynamically re-read from SharedConfig on each activity.
    // When idle_timeout_secs is 0, idle_deadline stays None (disabled).
    let mut idle_deadline: Option<tokio::time::Instant> =
        current_idle_timeout(&shared_config).map(|d| tokio::time::Instant::now() + d);

    // GoingAway deadline: set when we send GoingAway to a capable client.
    // The connection force-closes when this expires, independent of the
    // lifecycle watcher, as a safety net.
    let mut drain_deadline: Option<tokio::time::Instant> = None;

    // Subscribe to lifecycle state changes for drain/shutdown detection.
    let mut lifecycle_rx = lifecycle.subscribe();

    loop {
        // Build idle sleep future for this iteration.
        let idle_fut = async {
            match idle_deadline {
                Some(deadline) => tokio::time::sleep_until(deadline).await,
                None => std::future::pending().await,
            }
        };

        // Build drain deadline future.
        let drain_fut = async {
            match drain_deadline {
                Some(deadline) => tokio::time::sleep_until(deadline).await,
                None => std::future::pending().await,
            }
        };

        tokio::select! {
            frame_result = framed.next() => {
                // Any inbound frame resets the idle timer.
                // Re-read timeout from SharedConfig so hot-reload changes
                // affect existing connections.
                idle_deadline = current_idle_timeout(&shared_config)
                    .map(|d| tokio::time::Instant::now() + d);

                match frame_result {
                    Some(Ok(frame)) => {
                        // Belt-and-suspenders: reject subscribes if the lifecycle
                        // has transitioned to Draining but the session hasn't been
                        // transitioned yet (race: subscribe frame read from buffer
                        // in the same select iteration as the lifecycle change).
                        if lifecycle.is_draining()
                            && matches!(
                                frame,
                                Frame::Subscribe { .. }
                                    | Frame::SubscribeAt { .. }
                                    | Frame::SubscribeFrom { .. }
                                    | Frame::Reconnect { .. }
                            )
                        {
                            let _ = framed.send(Frame::Error {
                                code: ErrorCode::ServerDraining,
                                message: "server is draining, new subscriptions rejected".into(),
                            }).await;
                            continue;
                        }

                        let outcome = session.process_frame(frame, auth.as_ref(), manager.as_ref(), &delivery_tx).await;

                        // After successful Connect/Reconnect, configure the codec
                        // for header support based on negotiated capabilities.
                        if let FrameOutcome::Reply(ref f) = outcome {
                            let is_auth_success = matches!(
                                f,
                                Frame::ConnectAck { success: true, .. }
                                    | Frame::ReconnectAck { .. }
                            );
                            if is_auth_success {
                                let caps = session.capabilities();

                                if caps & CAP_HEADERS != 0 {
                                    let cfg = shared_config.load();
                                    let codec = framed.codec_mut();
                                    codec.set_headers_enabled(true);
                                    codec.set_header_limits(HeaderLimits {
                                        max_count: cfg.max_headers_per_message,
                                        max_key_bytes: cfg.max_header_key_bytes,
                                        max_value_bytes: cfg.max_header_value_bytes,
                                        max_total_bytes: cfg.max_header_bytes_total,
                                    });
                                }

                                if caps & CAP_MESSAGE_TTL != 0 {
                                    framed.codec_mut().set_ttl_enabled(true);
                                }
                            }
                        }

                        match outcome {
                            FrameOutcome::Reply(f) => {
                                if framed.send(f).await.is_err() {
                                    break;
                                }
                            }
                            FrameOutcome::ReplyMany(frames) => {
                                for f in frames {
                                    if framed.send(f).await.is_err() {
                                        break;
                                    }
                                }
                            }
                            FrameOutcome::ReplyAndClose(f) => {
                                let _ = framed.send(f).await;
                                break;
                            }
                            FrameOutcome::None => {}
                        }
                    }
                    Some(Err(e)) => {
                        let code = ErrorCode::from_error(&e);
                        tracing::warn!(
                            client = %client_id,
                            error = %e,
                            fatal = code.is_fatal(),
                            "codec error on inbound frame"
                        );
                        let _ = framed.send(Frame::Error {
                            code,
                            message: e.to_string(),
                        }).await;
                        if code.is_fatal() {
                            break;
                        }
                    }
                    None => break, // TCP FIN
                }
            }
            Some(msg) = delivery_rx.recv() => {
                // Deliveries also count as activity — re-read timeout dynamically.
                idle_deadline = current_idle_timeout(&shared_config)
                    .map(|d| tokio::time::Instant::now() + d);

                let frame = session.build_deliver_frame(&msg);
                if framed.send(frame).await.is_err() {
                    break;
                }
            }
            () = idle_fut, if idle_deadline.is_some() => {
                tracing::info!(client = %client_id, "closing idle connection");
                let _ = framed.send(Frame::Error {
                    code: ErrorCode::InternalError,
                    message: "idle timeout".into(),
                }).await;
                break;
            }
            () = drain_fut, if drain_deadline.is_some() => {
                tracing::info!(client = %client_id, "GoingAway deadline expired, force-closing");
                let _ = framed.send(Frame::Error {
                    code: ErrorCode::ServerDraining,
                    message: "drain deadline expired".into(),
                }).await;
                break;
            }
            _ = lifecycle_rx.changed() => {
                let state = *lifecycle_rx.borrow();
                match state {
                    ServerState::Draining => {
                        let caps = session.capabilities();
                        if caps & CAP_GOING_AWAY != 0 {
                            // Capable client: send GoingAway with remaining deadline.
                            let remaining_ms = lifecycle.drain_started_at()
                                .map(|started| {
                                    let elapsed = started.elapsed().as_millis() as u64;
                                    (deploy.drain_timeout_secs * 1000).saturating_sub(elapsed) as u32
                                })
                                .unwrap_or(deploy.drain_timeout_secs as u32 * 1000);
                            let reason = GoingAwayReason::from_config(&deploy.going_away_reason);
                            tracing::debug!(
                                client = %client_id,
                                deadline_ms = remaining_ms,
                                "sending GoingAway to capable client"
                            );
                            let _ = framed.send(Frame::GoingAway {
                                reason,
                                deadline_ms: remaining_ms,
                                message: "server shutting down".into(),
                            }).await;
                            metrics::counter!("fafb_going_away_sent_total").increment(1);
                            // Set a local deadline so this connection force-closes
                            // even if the lifecycle watcher fails to transition.
                            // Add 1s buffer so the lifecycle's ShuttingDown transition
                            // wins in the normal case — this only fires as a safety net.
                            drain_deadline = Some(
                                tokio::time::Instant::now()
                                    + Duration::from_millis(remaining_ms as u64)
                                    + Duration::from_secs(1),
                            );
                            // Transition session to Draining — rejects new subscribes
                            // but allows publish/ack/unsub/ping/pong.
                            session.transition_to_draining();
                        } else {
                            // No GoingAway capability: send error and close.
                            tracing::debug!(client = %client_id, "server draining, closing connection (no CAP_GOING_AWAY)");
                            let _ = framed.send(Frame::Error {
                                code: ErrorCode::ServerDraining,
                                message: "server shutting down".into(),
                            }).await;
                            break;
                        }
                    }
                    ServerState::ShuttingDown => {
                        tracing::debug!(client = %client_id, "server shutting down, force-closing");
                        let _ = framed.send(Frame::Error {
                            code: ErrorCode::InternalError,
                            message: "server shutting down".into(),
                        }).await;
                        break;
                    }
                    ServerState::Running => {
                        // No action needed.
                    }
                }
            }
        }
    }

    // Cleanup: always runs.
    session.abort_forwarders();
    let _ = manager.disconnect(&session.client_id()).await;

    tracing::info!(client = %client_id, "connection closed");
}
