use std::path::PathBuf;
use std::sync::Arc;

use anyhow::Context;
use axum::routing::{get, post};
use axum::{Extension, Router};
use tokio::net::TcpListener;
use tracing::info;
use tracing_subscriber::EnvFilter;

use doh_forwarder::config::AppConfig;
use doh_forwarder::dns::listener::DnsListener;
use doh_forwarder::routes::{dns_query_get, dns_query_post, AppState};
use doh_forwarder::upstream::doh_client::DohClient;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // Initialize tracing (RUST_LOG env filter, defaults to info)
    tracing_subscriber::fmt()
        .with_env_filter(
            EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")),
        )
        .init();

    // Parse CLI args for config path
    let config_path = parse_config_path();
    info!(path = %config_path.display(), "loading configuration");

    let config = AppConfig::load(&config_path)
        .with_context(|| format!("failed to load config from {:?}", config_path))?;

    info!(listen = %config.server.listen, "starting doh-forwarder");
    info!(
        resolvers = config
            .upstream
            .resolvers
            .iter()
            .map(|r| r.name.as_str())
            .collect::<Vec<_>>()
            .join(", "),
        strategy = %config.upstream.strategy,
        "upstream configuration"
    );

    let strategy = config.upstream.strategy()?;
    let doh_client = DohClient::new(config.upstream.resolvers.clone());

    // Start DNS listener (UDP/TCP) in background
    let dns_listener = DnsListener::new(doh_client.clone(), strategy);
    let dns_addr = config.dns.listen.clone();
    tokio::spawn(async move {
        if let Err(e) = dns_listener.run(&dns_addr).await {
            tracing::error!(error = %e, "DNS listener failed");
        }
    });
    info!(addr = %config.dns.listen, "DNS listener starting");

    let state = AppState {
        doh_client,
        rr_counter: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
        strategy,
    };

    let app = Router::new()
        .route("/dns-query", post(dns_query_post).get(dns_query_get))
        .route("/health", get(health))
        .layer(Extension(state));

    let listener = TcpListener::bind(&config.server.listen)
        .await
        .with_context(|| format!("failed to bind to {}", config.server.listen))?;

    info!(addr = %listener.local_addr()?, "listening");

    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal())
        .await
        .context("server error")?;

    info!("shutdown complete");
    Ok(())
}

/// Health-check endpoint.
async fn health() -> &'static str {
    "ok"
}

/// Wait for a shutdown signal (SIGINT or SIGTERM).
async fn shutdown_signal() {
    let ctrl_c = async {
        tokio::signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
            .expect("failed to install signal handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => info!("received SIGINT, shutting down"),
        _ = terminate => info!("received SIGTERM, shutting down"),
    }
}

/// Parse the config file path from CLI args, defaulting to `config/default.toml`.
fn parse_config_path() -> PathBuf {
    let args: Vec<String> = std::env::args().collect();
    if let Some(pos) = args.iter().position(|a| a == "--config") {
        if let Some(path) = args.get(pos + 1) {
            return PathBuf::from(path);
        }
    }
    PathBuf::from("config/default.toml")
}
