//! Episteme (StemeDB) API server binary.

use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::Arc;
use tracing::{error, info, warn};
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

use axum::Extension;
use metrics_exporter_prometheus::PrometheusBuilder;
use stemedb_api::{
    bootstrap, create_router_config, create_router_full_protection_full_config,
    create_router_with_meter_config, ApiKeyAuthConfig, AppState, SecurityConfig,
};
use stemedb_ingest::worker::IngestWorker;
use stemedb_storage::HybridStore;
use stemedb_wal::Journal;

use axum_server::tls_rustls::RustlsConfig;
use std::path::Path;

/// Server configuration.
#[derive(Debug, Clone)]
struct Config {
    wal_dir: PathBuf,
    db_dir: PathBuf,
    bind_addr: String,
    meter_enabled: bool,
    corpus_db_dir: Option<PathBuf>,

    /// TLS certificate path (optional - enables HTTPS)
    tls_cert_path: Option<PathBuf>,

    /// TLS private key path (optional - enables HTTPS)
    tls_key_path: Option<PathBuf>,

    // P5.1: Security Configuration
    /// Write endpoint body limit in bytes (default: 1MB)
    write_body_limit: usize,
    /// Read endpoint body limit in bytes (default: 64KB)
    read_body_limit: usize,
    /// HTTP request timeout in seconds (default: 30)
    http_timeout_secs: u64,

    /// Skip Ed25519 signature verification (unsafe, for dev/testing only)
    unsafe_skip_signatures: bool,
}

impl Default for Config {
    fn default() -> Self {
        Self {
            wal_dir: PathBuf::from("data/wal"),
            db_dir: PathBuf::from("data/db"),
            bind_addr: "127.0.0.1:18180".to_string(),
            meter_enabled: true,
            corpus_db_dir: None,
            tls_cert_path: None,
            tls_key_path: None,
            // P5.1: Security defaults
            write_body_limit: 1024 * 1024, // 1MB
            read_body_limit: 64 * 1024,    // 64KB
            http_timeout_secs: 30,
            unsafe_skip_signatures: false,
        }
    }
}

impl Config {
    /// Convert to SecurityConfig for router configuration.
    fn to_security_config(&self) -> SecurityConfig {
        SecurityConfig {
            write_body_limit: self.write_body_limit,
            read_body_limit: self.read_body_limit,
            http_timeout_secs: self.http_timeout_secs,
        }
    }
}

impl Config {
    fn from_env() -> Self {
        let mut config = Self::default();
        if let Ok(wal_dir) = std::env::var("STEMEDB_WAL_DIR") {
            config.wal_dir = PathBuf::from(wal_dir);
        }
        if let Ok(db_dir) = std::env::var("STEMEDB_DB_DIR") {
            config.db_dir = PathBuf::from(db_dir);
        }
        if let Ok(bind_addr) = std::env::var("STEMEDB_BIND_ADDR") {
            config.bind_addr = bind_addr;
        }
        if let Ok(meter_enabled) = std::env::var("STEMEDB_METER_ENABLED") {
            config.meter_enabled = meter_enabled.to_lowercase() != "false" && meter_enabled != "0";
        }
        if let Ok(corpus_db_dir) = std::env::var("STEMEDB_CORPUS_DB_DIR") {
            config.corpus_db_dir = Some(PathBuf::from(corpus_db_dir));
        }

        if let Ok(tls_cert_path) = std::env::var("STEMEDB_TLS_CERT_PATH") {
            config.tls_cert_path = Some(PathBuf::from(tls_cert_path));
        }

        if let Ok(tls_key_path) = std::env::var("STEMEDB_TLS_KEY_PATH") {
            config.tls_key_path = Some(PathBuf::from(tls_key_path));
        }

        // P5.1: Security Configuration
        if let Ok(limit) = std::env::var("STEMEDB_WRITE_BODY_LIMIT") {
            if let Ok(parsed) = limit.parse::<usize>() {
                config.write_body_limit = parsed;
            }
        }

        if let Ok(limit) = std::env::var("STEMEDB_READ_BODY_LIMIT") {
            if let Ok(parsed) = limit.parse::<usize>() {
                config.read_body_limit = parsed;
            }
        }

        if let Ok(timeout) = std::env::var("STEMEDB_HTTP_TIMEOUT_SECS") {
            if let Ok(parsed) = timeout.parse::<u64>() {
                config.http_timeout_secs = parsed;
            }
        }

        if let Ok(val) = std::env::var("STEMEDB_UNSAFE_SKIP_SIGNATURES") {
            config.unsafe_skip_signatures = val.to_lowercase() == "true" || val == "1";
        }

        config
    }
}

/// Load TLS configuration from certificate and key files.
///
/// Returns an axum-server RustlsConfig.
async fn load_tls_config(
    cert_path: &Path,
    key_path: &Path,
) -> Result<RustlsConfig, Box<dyn std::error::Error>> {
    let config = RustlsConfig::from_pem_file(cert_path, key_path)
        .await
        .map_err(|e| format!("Failed to load TLS config: {}", e))?;

    Ok(config)
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let env_filter = tracing_subscriber::EnvFilter::try_from_default_env()
        .unwrap_or_else(|_| "stemedb_api=debug,tower_http=debug".into());
    tracing_subscriber::registry().with(env_filter).with(tracing_subscriber::fmt::layer()).init();

    let prometheus_handle = Arc::new(PrometheusBuilder::new().install_recorder()?);
    let config = Config::from_env();

    std::fs::create_dir_all(&config.wal_dir)?;
    std::fs::create_dir_all(&config.db_dir)?;

    let write_journal = Journal::open(&config.wal_dir)?;
    let read_journal = Journal::open(&config.wal_dir)?;
    let store = Arc::new(HybridStore::open(&config.db_dir)?);
    let corpus_store = config
        .corpus_db_dir
        .as_ref()
        .map(|d| {
            let _ = std::fs::create_dir_all(d);
            HybridStore::open(d).map(Arc::new)
        })
        .transpose()?;

    let state = AppState::new(write_journal, read_journal, Arc::clone(&store), corpus_store);

    // Bootstrap root API key from env (idempotent: no-op if key already exists).
    if let Err(e) = bootstrap::bootstrap_root_api_key(&*state.api_key_store).await {
        error!("Failed to bootstrap root API key: {}", e);
        std::process::exit(1);
    }

    // Cluster mode: join SWIM membership, start gRPC server, Gateway, gossip, and anti-entropy.
    // Requires the `cluster` feature to be enabled at compile time.
    #[cfg(feature = "cluster")]
    let cluster_gossip: Option<Arc<dyn stemedb_ingest::gossip::GossipBroadcast>> = {
        let cluster_mode = std::env::var("STEMEDB_CLUSTER_MODE")
            .map(|v| v.to_lowercase() == "true" || v == "1")
            .unwrap_or(false);

        if cluster_mode {
            use stemedb_cluster::{
                stable_node_id, Gateway, NodeInfo, RangeManager, RangeRouter, ShardingConfig,
                SwimConfig, SwimMembership,
            };
            use stemedb_rpc::proto::sync_service_server::SyncServiceServer;
            use stemedb_rpc::SyncServiceHandler;

            let node_id = stable_node_id();

            let rpc_addr: SocketAddr = std::env::var("STEMEDB_NODE_RPC_ADDR")
                .unwrap_or_else(|_| "0.0.0.0:18182".to_string())
                .parse()
                .unwrap_or_else(|_| SocketAddr::from(([0, 0, 0, 0], 18182)));

            let gateway_addr: SocketAddr = std::env::var("STEMEDB_NODE_API_ADDR")
                .unwrap_or_else(|_| "0.0.0.0:18181".to_string())
                .parse()
                .unwrap_or_else(|_| SocketAddr::from(([0, 0, 0, 0], 18181)));

            let api_addr: SocketAddr = config
                .bind_addr
                .parse()
                .unwrap_or_else(|_| SocketAddr::from(([0, 0, 0, 0], 18180)));

            // --- Membership ---
            let local_info = NodeInfo::new(node_id, rpc_addr, api_addr);
            let membership = Arc::new(SwimMembership::new(local_info, SwimConfig::default()));

            // Resolve seeds via DNS (for k8s headless service names)
            let raw_seeds = std::env::var("STEMEDB_SEED_NODES").unwrap_or_default();
            let mut seeds = Vec::new();
            for entry in raw_seeds.split(',').filter(|s| !s.trim().is_empty()) {
                let entry = entry.trim();
                if let Ok(addr) = entry.parse::<SocketAddr>() {
                    seeds.push(addr);
                } else {
                    match tokio::net::lookup_host(entry).await {
                        Ok(mut addrs) => {
                            if let Some(addr) = addrs.next() {
                                info!(seed = entry, resolved = %addr, "Resolved seed node");
                                seeds.push(addr);
                            }
                        }
                        Err(e) => {
                            info!(seed = entry, error = %e, "Failed to resolve seed");
                        }
                    }
                }
            }

            if let Err(e) = membership.join(seeds).await {
                warn!("Cluster join failed (continuing as solo node): {}", e);
            }
            membership.start();
            membership.spawn_background_tasks();

            // --- Sharding ---
            let num_shards = std::env::var("STEMEDB_NUM_SHARDS")
                .ok()
                .and_then(|s| s.parse().ok())
                .unwrap_or(4u32);
            let replication_factor = std::env::var("STEMEDB_REPLICATION_FACTOR")
                .ok()
                .and_then(|s| s.parse().ok())
                .unwrap_or(1u32);

            let router = Arc::new(RangeRouter::new(node_id));
            let sharding_config = ShardingConfig::new()
                .with_num_shards(num_shards)
                .with_replication_factor(replication_factor);

            let range_manager = RangeManager::new(
                Arc::clone(&router),
                Arc::clone(&membership),
                sharding_config,
                node_id,
            );
            if let Err(e) = range_manager.initialize_shards() {
                error!("Failed to initialize shards: {}", e);
            }

            // --- gRPC server (SyncStorageBridge) ---
            let merkle_manager = Arc::new(
                stemedb_sync::MerkleTreeManager::load_or_create(Arc::clone(&store))
                    .await
                    .map_err(|e| format!("Failed to load Merkle tree: {e}"))?,
            );

            let bridge = Arc::new(stemedb_api::sync_storage_bridge::SyncStorageBridge::new(
                Arc::clone(&store),
                Arc::clone(&merkle_manager),
                *node_id.as_bytes(),
                rpc_addr.to_string(),
                api_addr.to_string(),
            ));

            let grpc_service = SyncServiceServer::new(SyncServiceHandler::new(bridge));
            let grpc_rpc_addr = rpc_addr;
            tokio::spawn(async move {
                info!(addr = %grpc_rpc_addr, "gRPC server listening");
                if let Err(e) = tonic::transport::Server::builder()
                    .add_service(grpc_service)
                    .serve(grpc_rpc_addr)
                    .await
                {
                    error!("gRPC server error: {}", e);
                }
            });

            // --- Gossip broadcaster ---
            let peer_addrs: Vec<String> =
                membership.members().iter().map(|m| format!("http://{}", m.rpc_addr)).collect();
            let gossip_broadcaster: Option<Arc<dyn stemedb_ingest::gossip::GossipBroadcast>> =
                match stemedb_sync::GossipBroadcaster::new(peer_addrs).await {
                    Ok(b) => Some(Arc::new(b)),
                    Err(e) => {
                        warn!("Failed to create gossip broadcaster: {}", e);
                        None
                    }
                };

            // --- Gateway ---
            let gateway = Gateway::new(Arc::clone(&router), Arc::clone(&membership), gateway_addr)
                .with_prometheus_handle(Arc::clone(&prometheus_handle));

            tokio::spawn(async move {
                if let Err(e) = gateway.serve().await {
                    error!("Gateway error: {}", e);
                }
            });

            // --- Shard rebalance listener ---
            let rebalance_membership = Arc::clone(&membership);
            let rebalance_router = Arc::clone(&router);
            tokio::spawn(async move {
                let mut rx = rebalance_membership.subscribe();
                let sharding_cfg = ShardingConfig::new()
                    .with_num_shards(num_shards)
                    .with_replication_factor(replication_factor);
                loop {
                    match rx.recv().await {
                        Ok(event) => {
                            if event.is_join() || event.is_failure() || event.is_leave() {
                                // Debounce: wait 5s for cluster to stabilize
                                tokio::time::sleep(std::time::Duration::from_secs(5)).await;
                                let rm = RangeManager::new(
                                    Arc::clone(&rebalance_router),
                                    Arc::clone(&rebalance_membership),
                                    sharding_cfg.clone(),
                                    node_id,
                                );
                                if let Err(e) = rm.rebalance_shards() {
                                    warn!("Shard rebalance failed: {}", e);
                                }
                            }
                        }
                        Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
                            warn!(skipped = n, "Membership event receiver lagged");
                        }
                        Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
                    }
                }
            });

            info!(
                node_id = %node_id.short_hex(),
                rpc_addr = %rpc_addr,
                gateway_addr = %gateway_addr,
                "Cluster mode fully active"
            );

            gossip_broadcaster
        } else {
            None
        }
    };

    #[cfg(not(feature = "cluster"))]
    let cluster_gossip: Option<Arc<dyn stemedb_ingest::gossip::GossipBroadcast>> = None;

    // Spawn IngestWorker (after cluster block so gossip broadcaster can be injected)
    let worker_journal = state.journal.clone();
    let worker_store = store;
    let worker_flush_notify = Arc::clone(&state.flush_notify);
    let skip_sigs = config.unsafe_skip_signatures;
    let worker_memtable = Arc::clone(&state.memtable);
    tokio::spawn(async move {
        match IngestWorker::new(worker_journal, worker_store).await {
            Ok(worker) => {
                let mut worker = worker
                    .with_flush_notify(worker_flush_notify)
                    .with_memtable(worker_memtable)
                    .with_skip_signature_verification(skip_sigs);
                if let Some(gossip) = cluster_gossip {
                    worker = worker.with_gossip_broadcaster(gossip);
                }
                info!(skip_signatures = skip_sigs, "IngestWorker started");
                worker.run().await;
            }
            Err(e) => error!("Failed to create IngestWorker: {:?}", e),
        }
    });

    // Startup guard: unsafe skip + auth enabled is a fatal misconfiguration.
    if config.unsafe_skip_signatures && bootstrap::is_auth_enabled() {
        error!(
            "FATAL: STEMEDB_UNSAFE_SKIP_SIGNATURES=true conflicts with \
             STEMEDB_AUTH_ENABLED=true. Signature verification must be enabled \
             when auth is enforced. Unset STEMEDB_UNSAFE_SKIP_SIGNATURES or \
             disable STEMEDB_AUTH_ENABLED."
        );
        std::process::exit(1);
    }

    // Build router (with or without metering) with security config
    let security_config = config.to_security_config();
    info!(
        "P5.1 Security: write_limit={}KB, read_limit={}KB, http_timeout={}s",
        security_config.write_body_limit / 1024,
        security_config.read_body_limit / 1024,
        security_config.http_timeout_secs,
    );

    let app = if bootstrap::is_auth_enabled() {
        info!(
            require_all = bootstrap::is_auth_require_all(),
            "Auth enforced (STEMEDB_AUTH_ENABLED=true) — full protection stack active"
        );
        create_router_full_protection_full_config(
            state,
            ApiKeyAuthConfig {
                enabled: true,
                require_for_all: bootstrap::is_auth_require_all(),
                ..ApiKeyAuthConfig::default()
            },
            security_config,
        )
    } else if config.meter_enabled {
        info!("The Meter enabled: economic throttling active (10K tokens/agent/hour)");
        create_router_with_meter_config(state, security_config)
    } else {
        info!("The Meter disabled: no quota enforcement");
        create_router_config(state, security_config)
    };

    // Add Prometheus handle extension and /metrics route
    let app = app.layer(Extension(prometheus_handle));

    // Start server with or without TLS
    if let (Some(cert_path), Some(key_path)) = (&config.tls_cert_path, &config.tls_key_path) {
        info!("TLS enabled - loading certificate and key");
        let tls_config = load_tls_config(cert_path, key_path).await?;

        info!("API server listening on {} (TLS enabled)", config.bind_addr);
        info!("Swagger UI available at https://{}/swagger-ui", config.bind_addr);

        axum_server::bind_rustls(config.bind_addr.parse()?, tls_config)
            .serve(app.into_make_service_with_connect_info::<SocketAddr>())
            .await?;
    } else {
        warn!("TLS not configured - running in plaintext mode (NOT for production)");

        let listener = tokio::net::TcpListener::bind(&config.bind_addr).await?;
        info!("API server listening on {} (plaintext)", config.bind_addr);
        info!("Swagger UI available at http://{}/swagger-ui", config.bind_addr);

        axum::serve(listener, app.into_make_service_with_connect_info::<SocketAddr>()).await?;
    }

    Ok(())
}
