//! Raahi entrypoint: boots the SQLite store, compiles the initial config snapshot,
//! and runs the Pingora data plane (HTTP + optional rustls HTTPS with per-SNI
//! certificate selection) plus the admin API and health checks as background services.

mod cli;

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

/// mimalloc: glibc malloc's arena futex serialized the worker threads under
/// load (~15% of CPU in __lll_lock_*_private in profiles).
#[global_allocator]
static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc;

use anyhow::Context;
use clap::Parser;
use pingora::prelude::*;
use pingora::services::background::background_service;
use raahi_acme::{AcmeHandle, AcmeService};
use raahi_api::ApiService;
use raahi_core::{AlpnChallengeRegistry, RouteSpec, ServiceSpec, TargetSpec};
use raahi_discovery::DiscoveryService;
use raahi_proxy::{
    AlpnChallengeResponder, CertHandle, CertStore, HttpLogService, Metrics, RaahiProxy,
    StreamProxyApp, config_handle, log_channel, sni_tls_settings,
};
use raahi_store::Store;
use tracing::info;

#[derive(Parser, Debug)]
#[command(
    name = "raahi",
    version,
    about = "Raahi — a configurable reverse proxy"
)]
struct Cli {
    /// SQLite database URL (created if missing).
    #[arg(long, env = "RAAHI_DB", default_value = "sqlite://raahi.db")]
    db: String,

    /// Seed a demo service + route on startup (idempotent-ish; skips if a service exists).
    #[arg(long)]
    seed: bool,

    /// Override the proxy HTTP listen address (else taken from settings).
    #[arg(long)]
    http_addr: Option<String>,

    /// Override the proxy HTTPS listen address (else taken from settings).
    #[arg(long)]
    https_addr: Option<String>,

    /// Override the admin API listen address (else taken from settings).
    #[arg(long)]
    admin_addr: Option<String>,

    /// Directory of the built UI to serve from the admin API.
    #[arg(long, env = "RAAHI_UI_DIR", default_value = "ui/build")]
    ui_dir: PathBuf,

    /// Worker threads for the proxy data plane (default: all CPU cores).
    /// Pingora's own default is 1, which caps throughput at one core.
    #[arg(long, env = "RAAHI_THREADS")]
    threads: Option<usize>,

    /// Without a subcommand, Raahi runs the proxy server.
    #[command(subcommand)]
    command: Option<cli::Command>,
}

fn map_pingora<E: std::fmt::Display>(e: E) -> anyhow::Error {
    anyhow::anyhow!("pingora: {e}")
}

/// Insert a demo service (two local targets) + catch-all route if the DB is empty.
async fn seed(store: &Store) -> anyhow::Result<()> {
    if !store.list_services().await?.is_empty() {
        info!("seed: services already present, skipping");
        return Ok(());
    }
    let svc = store
        .create_service(&ServiceSpec {
            kind: Default::default(),
            root: None,
            spa_fallback: false,
            upstream_authority: None,
            name: "demo-service".into(),
            protocol: raahi_core::Protocol::Http,
            connect_timeout_ms: 2000,
            read_timeout_ms: 10000,
            write_timeout_ms: 10000,
            retries: 1,
            lb_algorithm: raahi_core::LbAlgorithm::RoundRobin,
            tls_sni: None,
            health_path: None,
        })
        .await?;
    for port in [9001u16, 9002] {
        store
            .create_target(
                svc.id,
                &TargetSpec {
                    priority: 0,
                    host: "127.0.0.1".into(),
                    port,
                    weight: 100,
                    enabled: true,
                },
            )
            .await?;
    }
    store
        .create_route(&RouteSpec {
            name: "demo-route".into(),
            service_id: svc.id,
            priority: 0,
            hosts: vec![],
            paths: vec!["/".into()],
            methods: vec![],
            headers: Default::default(),
            splits: vec![],
            strip_path: false,
            preserve_host: false,
            enabled: true,
        })
        .await?;
    info!("seed: created demo-service + demo-route");
    Ok(())
}

fn main() -> anyhow::Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()),
        )
        .init();

    let cli = Cli::parse();
    if let Some(command) = cli.command {
        return cli::run(command);
    }

    // One-shot async setup on a current-thread runtime, dropped before the server runs.
    let setup_rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()?;

    let (snapshot, settings, cert_store, store) = setup_rt.block_on(async {
        let store = Store::connect(&cli.db)
            .await
            .with_context(|| format!("open db {}", cli.db))?;
        if cli.seed {
            seed(&store).await?;
        }
        let snapshot = store.build_snapshot().await?;
        let settings = snapshot.settings.clone();
        // Load all certificates into an in-memory SNI store (rustls serves them all).
        let certs = store.list_certificates().await?;
        let cert_store = CertStore::from_certs(certs, settings.active_certificate_id);
        Ok::<_, anyhow::Error>((snapshot, settings, cert_store, store))
    })?;
    drop(setup_rt);

    // Stream (L4) listeners bind at startup only; grab the set before the snapshot
    // moves into the config handle.
    let stream_routes = snapshot.stream_routes.clone();
    let config = config_handle(snapshot);
    let metrics = Arc::new(Metrics::new());
    // Live cert handle shared by the TLS listener (reads it per handshake) and the admin
    // API (swaps it on certificate/default changes — no restart).
    let cert_count = cert_store.len();
    let cert_handle = CertHandle::new(cert_store);
    let alpn_challenges = AlpnChallengeRegistry::default();
    let acme_handle = AcmeHandle::default();
    let (discovery, discovery_handle, publisher, discovery_registry) =
        DiscoveryService::new(cli.db.clone(), store, config.clone());

    let http_addr = cli
        .http_addr
        .unwrap_or_else(|| settings.proxy_http_addr.clone());

    let mut server = Server::new(None).map_err(map_pingora)?;
    server.bootstrap();

    // http-log delivery channel: the proxy's log phase produces, a background
    // service batches + POSTs to collectors.
    let (log_tx, log_rx) = log_channel();

    let proxy = RaahiProxy {
        config: config.clone(),
        metrics: metrics.clone(),
        log_tx,
        split_counter: Default::default(),
    };
    let threads = cli.threads.unwrap_or_else(|| {
        std::thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(1)
    });

    let mut proxy_svc = http_proxy_service(&server.configuration, proxy);
    proxy_svc.threads = Some(threads);
    proxy_svc.add_tcp(&http_addr);
    info!("proxy HTTP listener on {http_addr} ({threads} worker threads)");

    let https_addr = cli.https_addr.or_else(|| settings.proxy_https_addr.clone());
    if let Some(https_addr) = &https_addr {
        proxy_svc.add_tls_with_settings(https_addr, None, sni_tls_settings(cert_handle.clone()));
        proxy_svc
            .endpoints()
            .set_pre_tls_callback(Arc::new(AlpnChallengeResponder::new(
                alpn_challenges.clone(),
            )));
        info!(
            "proxy HTTPS (rustls) listener on {https_addr}; {cert_count} issued cert(s), live per-SNI selection"
        );
    }

    server.add_service(proxy_svc);

    // L4 stream listeners, one raw TCP service per enabled stream route. Adding or
    // removing a stream route requires a restart; retargeting applies live.
    for sr in &stream_routes {
        let app = StreamProxyApp::new(config.clone(), sr.listen_addr.clone());
        let mut stream_svc =
            pingora::services::listening::Service::new(format!("stream-{}", sr.name), app);
        stream_svc.threads = Some(threads);
        stream_svc.add_tcp(&sr.listen_addr);
        server.add_service(stream_svc);
        info!(
            "stream (L4) listener on {} -> service #{} ({})",
            sr.listen_addr, sr.service_id, sr.name
        );
    }

    // Admin API (REST + UI) as a background service sharing the live config + metrics.
    let admin_addr = cli
        .admin_addr
        .unwrap_or_else(|| settings.admin_addr.clone());
    let api = ApiService {
        addr: admin_addr.clone(),
        db_url: cli.db.clone(),
        config: config.clone(),
        metrics: metrics.clone(),
        cert_handle: cert_handle.clone(),
        acme_handle: acme_handle.clone(),
        discovery_handle: discovery_handle.clone(),
        discovery_registry: discovery_registry.clone(),
        publisher: publisher.clone(),
        ui_dir: Some(cli.ui_dir.clone()),
    };
    server.add_service(background_service("admin-api", api));
    info!("admin API on {admin_addr}");

    server.add_service(background_service("http-log", HttpLogService::new(log_rx)));
    server.add_service(background_service("discovery", discovery));

    server.add_service(background_service(
        "acme",
        AcmeService {
            db_url: cli.db.clone(),
            cert_handle: cert_handle.clone(),
            challenges: alpn_challenges,
            handle: acme_handle,
            interval: Duration::from_secs(6 * 60 * 60),
        },
    ));

    // JWKS refresher for jwt plugins in identity-provider mode.
    server.add_service(background_service(
        "jwks",
        raahi_proxy::JwksService {
            config: config.clone(),
        },
    ));

    // Active health checks for upstream targets.
    let health = raahi_proxy::HealthService {
        config: config.clone(),
        interval: Duration::from_secs(5),
        connect_timeout: Duration::from_secs(2),
    };
    server.add_service(background_service("health", health));

    info!("Raahi started");
    server.run_forever();
}