use std::{net::SocketAddr, path::PathBuf, time::Duration};

use anyhow::{Context, Result, bail};
use axum::{
    Router,
    routing::{get, post},
};
use clap::Parser;
use tokio_util::sync::CancellationToken;
use tower_http::trace::TraceLayer;
use tracing::{error, info};
use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt};

mod config;
mod constants;
mod downstream;
mod metrics;
mod reload;
mod state;
mod upstream;

use config::{load_config, load_mcp_bearer_token, validate_config};
use downstream::{delete_mcp, get_mcp, post_mcp};
use metrics::{health, metrics};
use reload::{reload_config, reload_http};
use state::{
    AppState, build_state_with_config_path, prune_idle_sessions, stop_accepting_new_sessions,
    terminate_all_sessions,
};
use upstream::{server_names, spawn_server_manager};

#[derive(Parser, Debug)]
struct Args {
    #[arg(short, long, env = "MCPSTEAD_CONFIG")]
    config: Option<PathBuf>,
}

fn resolve_config_path(config: Option<PathBuf>) -> Result<PathBuf> {
    let Some(path) = config else {
        bail!("mcpstead requires --config <path> or MCPSTEAD_CONFIG=<path>");
    };
    Ok(path)
}

#[tokio::main]
async fn main() -> Result<()> {
    let args = Args::parse();
    let config_path = resolve_config_path(args.config)?;
    let config = load_config(&config_path).await?;
    init_logging(&config.logging.level);
    validate_config(&config)?;

    let state =
        build_state_with_config_path(&config, load_mcp_bearer_token(&config)?, config_path)?;

    for name in server_names(&state).await {
        spawn_server_manager(&state, name).await;
    }

    let shutdown_token = CancellationToken::new();
    tokio::spawn(session_gc_task(state.clone(), shutdown_token.clone()));
    tokio::spawn(reload_signal_task(state.clone(), shutdown_token.clone()));

    let app = build_app(state.clone());

    let addr: SocketAddr = format!("{}:{}", config.host, config.port)
        .parse()
        .with_context(|| format!("invalid bind address '{}:{}'", config.host, config.port))?;

    info!(%addr, "starting mcpstead");
    let listener = tokio::net::TcpListener::bind(addr).await?;
    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal(state, shutdown_token))
        .await?;
    Ok(())
}

fn build_app(state: AppState) -> Router {
    Router::new()
        .route("/mcp", post(post_mcp).get(get_mcp).delete(delete_mcp))
        .route("/-/reload", post(reload_http))
        .route("/health", get(health))
        .route("/metrics", get(metrics))
        .layer(TraceLayer::new_for_http())
        .with_state(state)
}

fn init_logging(level: &str) {
    let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(level));
    tracing_subscriber::registry()
        .with(filter)
        .with(tracing_subscriber::fmt::layer())
        .init();
}

async fn session_gc_task(state: AppState, shutdown_token: CancellationToken) {
    loop {
        tokio::select! {
            _ = shutdown_token.cancelled() => break,
            _ = tokio::time::sleep(state.session_gc_interval) => {
                prune_idle_sessions(&state).await;
            }
        }
    }
}

#[cfg(unix)]
async fn reload_signal_task(state: AppState, shutdown_token: CancellationToken) {
    use tokio::signal::unix::{SignalKind, signal};

    let mut hangup = match signal(SignalKind::hangup()) {
        Ok(signal) => signal,
        Err(err) => {
            error!(%err, "failed to install sighup handler");
            return;
        }
    };

    loop {
        tokio::select! {
            _ = shutdown_token.cancelled() => break,
            _ = hangup.recv() => {
                if let Err(err) = reload_config(&state).await {
                    error!(%err, "sighup reload failed");
                }
            }
        }
    }
}

#[cfg(not(unix))]
async fn reload_signal_task(_state: AppState, shutdown_token: CancellationToken) {
    shutdown_token.cancelled().await;
}

#[cfg(unix)]
async fn wait_for_shutdown_signal() {
    use tokio::signal::unix::{SignalKind, signal};

    let mut terminate = match signal(SignalKind::terminate()) {
        Ok(signal) => signal,
        Err(err) => {
            error!(%err, "failed to install sigterm handler");
            let _ = tokio::signal::ctrl_c().await;
            return;
        }
    };

    tokio::select! {
        result = tokio::signal::ctrl_c() => {
            if let Err(err) = result {
                error!(%err, "failed to install ctrl-c handler");
            }
        }
        _ = terminate.recv() => {}
    }
}

#[cfg(not(unix))]
async fn wait_for_shutdown_signal() {
    if let Err(err) = tokio::signal::ctrl_c().await {
        error!(%err, "failed to install ctrl-c handler");
    }
}

async fn shutdown_signal(state: AppState, shutdown_token: CancellationToken) {
    wait_for_shutdown_signal().await;
    shutdown_token.cancel();
    stop_accepting_new_sessions(&state);
    let sweep = terminate_all_sessions(&state, "gateway_shutdown");
    if tokio::time::timeout(state.session_shutdown_grace, sweep)
        .await
        .is_err()
    {
        error!(timeout = ?state.session_shutdown_grace, "session shutdown sweep timed out");
    }
    tokio::time::sleep(Duration::from_millis(10)).await;
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::{
        config::{Config, McpConfig, McpSessionConfig},
        state::build_state,
    };

    #[test]
    fn resolve_config_path_requires_explicit_config() {
        let err = resolve_config_path(None).unwrap_err();
        let message = err.to_string();
        assert!(
            message.contains("mcpstead requires --config <path> or MCPSTEAD_CONFIG=<path>"),
            "{message}"
        );
    }

    #[tokio::test]
    async fn session_gc_task_stops_on_cancel() {
        let config = Config {
            mcp: McpConfig {
                session: McpSessionConfig {
                    idle_ttl_seconds: 60,
                    gc_interval_seconds: 60,
                    shutdown_grace_seconds: 1,
                },
                ..McpConfig::default()
            },
            ..Config::default()
        };
        let state = build_state(&config, None).unwrap();
        let token = CancellationToken::new();
        let handle = tokio::spawn(session_gc_task(state, token.clone()));

        token.cancel();

        tokio::time::timeout(Duration::from_secs(1), handle)
            .await
            .expect("gc task should stop")
            .expect("gc task should not panic");
    }
}
