use std::net::SocketAddr;
use std::path::PathBuf;
use std::time::Duration;

use clap::Parser;
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;
use tracing::{info, warn};

use lightrail::error::{Error, Result};
use lightrail::identity;
use lightrail::storage;
use lightrail::sync::discovery_queue::DiscoveryQueue;
use lightrail::sync::{self, backfill, firehose, resync};
use lightrail::util::TokenExt;

#[derive(Parser, Debug)]
#[command(name = "lightrail", about = "listReposByCollection indexing service")]
struct Args {
    /// ATProto relay or PDS host to subscribe to (e.g. bsky.network).
    #[arg(long, env = "LIGHTRAIL_SUBSCRIBE")]
    subscribe: jacquard_common::url::Url,

    /// Path to the fjall database directory.
    #[arg(long, env = "LIGHTRAIL_DB_PATH", default_value = "lightrail.db")]
    db_path: PathBuf,

    /// TCP address for the XRPC API server.
    #[arg(long, env = "LIGHTRAIL_LISTEN", default_value = "0.0.0.0:2511")]
    listen: SocketAddr,

    /// PLC directory URL for did:plc resolution.
    /// Defaults to https://plc.directory. When used together with
    /// --slingshot-url, acts as a fallback if the slingshot resolver fails.
    #[arg(long, env = "LIGHTRAIL_PLC_URL")]
    plc_url: Option<jacquard_common::url::Url>,

    /// Slingshot URL for DID resolution.
    /// Slingshot covers did:plc (primary) and did:web (https well-known, then
    /// slingshot mini-doc). Omit the URL to use https://slingshot.microcosm.blue.
    #[arg(
        long,
        env = "LIGHTRAIL_SLINGSHOT_URL",
        num_args = 0..=1,
        default_missing_value = "https://slingshot.microcosm.blue"
    )]
    slingshot_url: Option<jacquard_common::url::Url>,

    /// Max identities kept in in-process identity cache.
    #[arg(long, env = "LIGHTRAIL_IDENT_CACHE_SIZE", default_value_t = 2_000_000)]
    ident_cache_size: u64,

    /// Global identity resolution rate limit (requests/sec).
    /// Backfill tasks wait for tokens; firehose/resync proceed freely but
    /// subtract from the budget so backfill slows when the firehose is busy.
    #[arg(long, env = "LIGHTRAIL_IDENTITY_RESOLUTION_QPS")]
    identity_resolution_qps: Option<std::num::NonZeroU32>,

    /// Maximum concurrent firehose commit worker tasks.
    #[arg(long, env = "LIGHTRAIL_MAX_FIREHOSE_WORKERS", default_value_t = 6)]
    max_firehose_workers: usize,

    /// Maximum concurrent resync worker tasks.
    #[arg(long, env = "LIGHTRAIL_MAX_RESYNC_WORKERS", default_value_t = 16)]
    max_resync_workers: usize,

    /// How often to flush the firehose cursor watermark to storage, in seconds.
    #[arg(long, env = "LIGHTRAIL_CURSOR_SAVE_INTERVAL", default_value_t = 1)]
    cursor_save_interval_secs: u64,

    /// HTTP timeout for describeRepo + getLatestCommit during resync, in seconds.
    #[arg(
        long,
        env = "LIGHTRAIL_DESCRIBE_REPO_FETCH_TIMEOUT",
        default_value_t = 30
    )]
    describe_repo_fetch_timeout_secs: u64,

    /// HTTP timeout for getRepo (full CAR download) during resync, in seconds.
    #[arg(long, env = "LIGHTRAIL_GET_REPO_FETCH_TIMEOUT", default_value_t = 300)]
    get_repo_fetch_timeout_secs: u64,

    /// TCP address for the Prometheus metrics HTTP endpoint.
    /// If not set, metrics are not exported.
    #[arg(long, env = "LIGHTRAIL_METRICS_LISTEN", num_args = 0..=1, default_missing_value = "0.0.0.0:6789")]
    metrics_listen: Option<SocketAddr>,

    /// Admin password for privileged API endpoints.
    #[arg(long, env = "LIGHTRAIL_ADMIN_PASSWORD")]
    admin_password: Option<String>,

    /// Enable deep crawl: discover PDS hosts via listHosts and crawl each one's repos.
    #[arg(long, action, env = "LIGHTRAIL_DEEP_CRAWL")]
    deep_crawl: bool,

    /// Heavy mode: always fetch the full repo CAR via getRepo for resync,
    /// skipping the cheaper describeRepo fast path.
    #[arg(long, action, env = "LIGHTRAIL_HEAVY")]
    heavy: bool,

    /// Per-PDS HTTP rate limit for crawl/resync requests, in requests per second.
    #[arg(long, env = "LIGHTRAIL_CRAWL_QPS", default_value_t = std::num::NonZeroU32::new(10).unwrap())]
    crawl_qps: std::num::NonZeroU32,

    /// fjall block cache size in MiB.
    #[arg(long, env = "LIGHTRAIL_FJALL_CACHE_MB", default_value_t = 256)]
    fjall_cache_mb: u64,

    /// Number of fjall background worker threads (flush + compaction).
    /// Defaults to fjall's own heuristic (min(CPU cores, 4)).
    #[arg(long, env = "LIGHTRAIL_FJALL_WORKER_THREADS")]
    fjall_worker_threads: Option<usize>,

    /// Max concurrent per-PDS listRepos workers during deep crawl.
    #[arg(
        long,
        env = "LIGHTRAIL_MAX_DEEP_CRAWL_WORKERS",
        requires("deep_crawl"),
        default_value_t = 4
    )]
    max_deep_crawl_workers: usize,
}

fn main() {
    rustls::crypto::aws_lc_rs::default_provider()
        .install_default()
        .expect("failed to install rustls crypto provider");

    tracing_subscriber::fmt()
        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
        .init();

    let rt = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .expect("failed to build tokio runtime");

    let result = rt.block_on(run());

    // Force-shutdown the runtime after a bounded wait. Without this, a
    // `spawn_blocking` task genuinely stuck in fjall (e.g. after Poisoned)
    // holds a blocking-pool thread, and since those threads are non-daemon
    // they'd prevent process exit. `shutdown_timeout` detaches any remaining
    // tasks after the deadline; the explicit `process::exit` below then
    // guarantees we don't wait for detached blocking threads either.
    rt.shutdown_timeout(Duration::from_secs(10));

    match result {
        Ok(()) => std::process::exit(0),
        Err(e) => {
            eprintln!("fatal: {e}");
            std::process::exit(1);
        }
    }
}

async fn run() -> Result<()> {
    let args = Args::parse();

    let subscribe_host = args
        .subscribe
        .host()
        .map(|h| h.to_owned())
        .ok_or(crate::Error::Other(
            "could not get host from --upstream".to_string(),
        ))?;

    let token = CancellationToken::new();

    let slingshot_url = args.slingshot_url;
    let plc_url = args.plc_url;
    let ident_cache_size = args.ident_cache_size;
    let resolver = std::sync::Arc::new(identity::build_resolver(
        slingshot_url,
        plc_url,
        ident_cache_size,
        token.clone(),
        args.identity_resolution_qps,
    ));

    if let Some(addr) = args.metrics_listen {
        install_metrics(addr)?;
    }

    let db = storage::open(
        &args.db_path,
        args.fjall_cache_mb,
        args.fjall_worker_threads,
    )?;
    let client = lightrail::http::build_client(args.crawl_qps);

    let dispatcher_state: resync::dispatcher::DispatcherState = std::sync::Arc::new(
        std::sync::Mutex::new(resync::dispatcher::DispatcherSnapshot::default()),
    );

    let discovery_queue = std::sync::Arc::new(DiscoveryQueue::new(8192, 32));

    let mut tasks: JoinSet<Result<()>> = JoinSet::new();

    tasks.spawn({
        let token = token.clone();
        let db = db.clone();
        let host = subscribe_host.clone();
        let resolver = resolver.clone();
        let client = client.clone();
        async move {
            let mut sub = firehose::Subscriber::new(
                host,
                db,
                resolver,
                args.max_firehose_workers,
                Duration::from_secs(args.cursor_save_interval_secs),
                client,
            );
            sub.run(token)
                .await
                .inspect(|_| info!("firehose subscriber done."))
                .inspect_err(|e| warn!(error = %e, "firehose exited"))
        }
    });

    tasks.spawn({
        let token = token.clone();
        let db = db.clone();
        let client = client.clone();
        let host = subscribe_host.clone();
        let resolver = resolver.clone();
        let discovery_queue = discovery_queue.clone();
        async move {
            match backfill::run(
                host,
                db,
                client,
                token.clone(),
                resolver,
                backfill::BackfillMode::Relay,
                discovery_queue,
            )
            .await
            {
                Ok(true) => {
                    info!("backfill complete; idling task");
                    token.cancelled().await;
                    Ok(())
                }
                Ok(false) => {
                    warn!("backfill ended without finishing, exiting");
                    Ok(())
                }
                Err(e) => {
                    warn!(error = %e, "backfill errored, exiting");
                    Err(e)
                }
            }
        }
    });

    tasks.spawn({
        let token = token.clone();
        let db = db.clone();
        let client = client.clone();
        let resolver = resolver.clone();
        let dispatcher_state = dispatcher_state.clone();
        let discovery_queue = discovery_queue.clone();
        async move {
            resync::dispatcher::run(resync::DispatcherConfig {
                resolver,
                db,
                client,
                max_concurrent: args.max_resync_workers,
                describe_timeout: Duration::from_secs(args.describe_repo_fetch_timeout_secs),
                get_repo_timeout: Duration::from_secs(args.get_repo_fetch_timeout_secs),
                token,
                force_get_repo: args.heavy,
                state: dispatcher_state,
                discovery_queue,
            })
            .await
            .inspect(|_| info!("resync done."))
            .inspect_err(|e| warn!(error = %e, "resync exited"))
        }
    });

    tasks.spawn({
        let token = token.clone();
        let db = db.clone();
        let addr = args.listen;
        let dispatcher_state = dispatcher_state.clone();
        let client = client.clone();
        let admin_config = args
            .admin_password
            .map(|pw| lightrail::server::AdminConfig {
                subscribe_host: subscribe_host.clone(),
                admin_password: pw,
            });
        async move {
            lightrail::server::serve(
                addr,
                db,
                token,
                admin_config,
                Some(dispatcher_state),
                Some(client),
            )
            .await
            .inspect(|_| info!("server done."))
            .inspect_err(|e| warn!(error = %e, "server exited"))
        }
    });

    tasks.spawn({
        let db = db.clone();
        let token = token.clone();
        async move {
            let mut interval = tokio::time::interval(Duration::from_secs(60));
            interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
            while token.run(interval.tick()).await.is_some() {
                // Flush the journal BufWriter to the OS. We use
                // manual_journal_persist so individual writes skip this;
                // doing it here batches many writes into one flush.
                if let Err(e) = db.persist_journal() {
                    warn!(error = %e, "failed to persist journal buffer");
                }
                if let Err(e) = storage::meta::save(&db) {
                    warn!(error = %e, "failed to periodically save meta stats");
                }
                // Emit fjall storage gauges.
                let ss = db.storage_stats();
                metrics::gauge!("lightrail_db_disk_bytes", "keyspace" => "total")
                    .set(ss.disk_bytes as f64);
                metrics::gauge!("lightrail_db_disk_bytes", "keyspace" => "default")
                    .set(ss.default_ks_disk_bytes as f64);
                metrics::gauge!("lightrail_db_disk_bytes", "keyspace" => "index")
                    .set(ss.index_ks_disk_bytes as f64);
                metrics::gauge!("lightrail_db_journal_count").set(ss.journal_count as f64);
                metrics::gauge!("lightrail_db_active_compactions")
                    .set(ss.active_compactions as f64);
                metrics::gauge!("lightrail_db_compactions_completed")
                    .set(ss.compactions_completed as f64);
                metrics::gauge!("lightrail_db_time_compacting_seconds")
                    .set(ss.time_compacting.as_secs_f64());
            }
            info!("meta stats done.");
            Ok(())
        }
    });

    if args.deep_crawl {
        tasks.spawn({
            let token = token.clone();
            let db = db.clone();
            let client = client.clone();
            let host = subscribe_host.clone();
            let resolver = resolver.clone();
            let discovery_queue = discovery_queue.clone();
            async move {
                sync::deep_crawl::run(
                    host,
                    db,
                    client,
                    args.max_deep_crawl_workers,
                    token,
                    resolver,
                    discovery_queue,
                )
                .await
                .inspect(|_| info!("deep crawl done."))
                .inspect_err(|e| warn!(error = %e, "deep crawl exited"))
            }
        });
    }

    // Wait for a shutdown trigger: ctrl-c or any task exiting (including via
    // panic, surfaced as a JoinError).
    let first = tokio::select! {
        _ = tokio::signal::ctrl_c() => {
            eprintln!("Shutting down...");
            None
        }
        r = tasks.join_next() => {
            eprintln!("=== a task exited with: {r:?} ===");
            r
        },
    };
    token.cancel();

    // Drain remaining tasks and surface the first error encountered.
    let mut error = first.and_then(into_error);
    while let Some(r) = tasks.join_next().await {
        if error.is_none() {
            error = into_error(r);
        }
    }

    if let Err(e) = storage::meta::save(&db) {
        warn!(error = %e, "failed to save meta stats on shutdown");
    }

    error.map_or(Ok(()), Err)
}

/// Flatten a task join result into an optional error.
/// Panics (JoinError) are treated as errors.
///
/// If the error indicates an unrecoverable database state
/// ([`Error::is_db_fatal`]), this immediately force-exits the process rather
/// than returning. Graceful shutdown isn't safe in that state because other
/// tasks may be stuck in blocking fjall calls that will never return.
fn into_error(r: std::result::Result<Result<()>, tokio::task::JoinError>) -> Option<Error> {
    let err = match r {
        Ok(Ok(())) => return None,
        Ok(Err(e)) => e,
        Err(e) => Error::TaskPanic(e),
    };
    if err.is_db_fatal() {
        eprintln!("FATAL: database poisoned, force-exiting: {err}");
        std::process::exit(2);
    }
    Some(err)
}

fn install_metrics(addr: SocketAddr) -> Result<()> {
    use metrics_exporter_prometheus::Matcher;
    metrics_exporter_prometheus::PrometheusBuilder::new()
        .with_http_listener(addr)
        // commit car payload sizes.. probably needs more small buckets
        .set_buckets_for_metric(
            Matcher::Full("lightrail_commit_car_bytes".to_string()),
            &[
                1_024., 4_096., 16_384., 65_536., 262_144., 1_048_576., 4_194_304.,
            ],
        )
        // ops per commit: vibes
        .and_then(|b| {
            b.set_buckets_for_metric(
                Matcher::Full("lightrail_commit_ops".to_string()),
                &[1., 2., 5., 10., 25., 50., 100., 200.],
            )
        })
        // retry delay seconds: matches the backoff_secs() ladder
        // (except for NotFound backoffs)
        .and_then(|b| {
            b.set_buckets_for_metric(
                Matcher::Full("lightrail_resync_retry_delay_seconds".to_string()),
                &[60., 120., 240., 480., 960., 1_920., 3_600.],
            )
        })
        // XRPC request latency: sub-ms cache hits up to multi-second scans
        .and_then(|b| {
            b.set_buckets_for_metric(
                Matcher::Full("lightrail_http_server_request_duration_seconds".to_string()),
                &[0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1., 5.],
            )
        })
        .and_then(|b| b.install())
        .map_err(|e| Error::Other(format!("failed to install metrics exporter: {e}")))?;
    info!(%addr, "metrics exporter listening");
    Ok(())
}