use std::env;
use std::fs;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;

use anyhow::{Context, Result};
use clap::Parser;
use pd_proto_grpc::{
    admin::v1::admin_server::AdminServer, pdpb::pd_server::PdServer, tsopb::tso_server::TsoServer,
};
use pd_proto_tso::KeyspaceGroup;
use pd_service_tso::{
    AdminConfigState, AdminGrpcService, AllocatorConfig, EtcdTimestampStore, InMemoryEtcdStore,
    KeyspaceGroupManager, LegacyPdGrpcService, TimestampStore, TsoGrpcService, TsoRuntime,
    TsoRuntimeConfig, TsoService,
};
use tokio::sync::{Mutex, RwLock};
use tonic::transport::{Certificate, Identity, Server, ServerTlsConfig};

const DEFAULT_BUILD_METADATA: &str = "None";

#[derive(Debug, Parser)]
#[command(name = "pd-tso", about = "PD TSO server", disable_version_flag = true)]
struct Cli {
    #[arg(
        short = 'V',
        long = "version",
        help = "Print version information and exit"
    )]
    version: bool,
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();
    if cli.version {
        print_version();
        return Ok(());
    }

    let log_level = env::var("PD_TSO_LOG_LEVEL").unwrap_or_else(|_| "info".to_string());
    pd_service_tso::init_runtime_logging(&log_level)
        .map_err(anyhow::Error::msg)
        .context("initialize runtime logging failed")?;

    let listen_addr = env::var("PD_TSO_LISTEN_ADDR")
        .unwrap_or_else(|_| "127.0.0.1:3379".to_string())
        .parse::<SocketAddr>()
        .context("invalid PD_TSO_LISTEN_ADDR")?;
    let cluster_id = env::var("PD_TSO_CLUSTER_ID")
        .unwrap_or_else(|_| "1".to_string())
        .parse::<u64>()
        .context("invalid PD_TSO_CLUSTER_ID")?;
    let member_name = env::var("PD_TSO_MEMBER_NAME").unwrap_or_else(|_| "pd-tso".to_string());
    let store_kind = env::var("PD_TSO_STORE").unwrap_or_else(|_| "etcd".to_string());
    let etcd_endpoints = parse_endpoints(env::var("PD_TSO_ETCD_ENDPOINTS").ok());
    let campaign_interval = parse_duration_ms_env("PD_TSO_CAMPAIGN_INTERVAL_MS", 500)?;
    let keepalive_interval = parse_duration_ms_env("PD_TSO_KEEPALIVE_INTERVAL_MS", 200)?;
    let watch_retry_backoff = parse_duration_ms_env("PD_TSO_WATCH_RETRY_BACKOFF_MS", 500)?;
    let lease_ttl_secs = parse_i64_env("PD_TSO_LEASE_TTL_SECS", 5)?;
    let external_ntp_server = env::var("PD_TSO_EXTERNAL_NTP_SERVER")
        .ok()
        .map(|value| value.trim().to_string())
        .filter(|value| !value.is_empty());
    let external_ntp_refresh_interval =
        parse_duration_ms_env("PD_TSO_EXTERNAL_NTP_REFRESH_MS", 30_000)?;
    let external_ntp_timeout = parse_duration_ms_env("PD_TSO_EXTERNAL_NTP_TIMEOUT_MS", 1_000)?;
    let allocator_lease_ms = allocator_lease_ms_from_ttl_secs(lease_ttl_secs);
    let startup_max_concurrent_tso_proxy_streamings = parse_legacy_proxy_max_concurrent_env(
        "PD_TSO_MAX_CONCURRENT_TSO_PROXY_STREAMINGS",
        pd_service_tso::runtime_config::DEFAULT_MAX_CONCURRENT_TSO_PROXY_STREAMINGS,
    )?;
    let startup_tso_proxy_recv_from_client_timeout = parse_legacy_proxy_recv_timeout_env(
        "PD_TSO_PROXY_RECV_FROM_CLIENT_TIMEOUT_MS",
        pd_service_tso::runtime_config::DEFAULT_TSO_PROXY_RECV_FROM_CLIENT_TIMEOUT,
    )?;
    let startup_max_gap_reset_ts_ms = parse_optional_positive_i64_env(&[
        "PD_TSO_MAX_GAP_RESET_TS",
        "PD_TSO_MAX_GAP_RESET_TS_MS",
    ])?;

    let store: Arc<dyn TimestampStore> = if store_kind.eq_ignore_ascii_case("in-memory") {
        Arc::new(InMemoryEtcdStore::new())
    } else {
        Arc::new(
            EtcdTimestampStore::connect(etcd_endpoints.clone())
                .context("connect etcd timestamp store failed")?,
        )
    };

    let mut manager = KeyspaceGroupManager::new(
        cluster_id,
        store,
        AllocatorConfig {
            lease_ms: allocator_lease_ms,
            ..AllocatorConfig::default()
        },
        member_name,
    );
    manager
        .register_group(KeyspaceGroup {
            id: 0,
            user_kind: "basic".to_string(),
            split_state: None,
            merge_state: None,
            members: Vec::new(),
            keyspaces: vec![0],
        })
        .context("register default keyspace group failed")?;
    if store_kind.eq_ignore_ascii_case("in-memory") {
        // In-memory test mode has no external election loop, so it self-campaigns.
        manager
            .campaign_group_primary(0, now_ms())
            .context("campaign default group primary failed")?;
    }

    let service = Arc::new(RwLock::new(TsoService::new(cluster_id, manager)));
    {
        let guard = service.read().await;
        let mut runtime = guard.runtime_config();
        runtime.max_concurrent_tso_proxy_streamings = startup_max_concurrent_tso_proxy_streamings;
        runtime.tso_proxy_recv_from_client_timeout = startup_tso_proxy_recv_from_client_timeout;
        if let Some(max_gap_reset_ts_ms) = startup_max_gap_reset_ts_ms {
            runtime.allocator.oracle.max_reset_ts_gap_ms = max_gap_reset_ts_ms;
        }
        guard.apply_runtime_config(runtime);
    }
    let admin_config = Arc::new(Mutex::new(AdminConfigState::default()));
    let tso_grpc = TsoGrpcService::new(Arc::clone(&service));
    let legacy_pd_grpc = LegacyPdGrpcService::new(Arc::clone(&service));
    let admin_grpc = AdminGrpcService::new(Arc::clone(&service), admin_config);

    let _runtime = if !store_kind.eq_ignore_ascii_case("in-memory") {
        Some(TsoRuntime::start(
            Arc::clone(&service),
            TsoRuntimeConfig {
                cluster_id,
                etcd_endpoints,
                campaign_interval,
                keepalive_interval,
                watch_retry_backoff,
                lease_ttl_secs,
                external_ntp_server,
                external_ntp_refresh_interval,
                external_ntp_timeout,
            },
        ))
    } else {
        None
    };

    let tls_ca_configured = env::var("PD_TSO_TLS_CA_PATH")
        .ok()
        .is_some_and(|value| !value.trim().is_empty());
    let tls_config = load_server_tls_config().context("load server TLS configuration failed")?;
    let require_admin_mtls = parse_bool_env(
        "PD_TSO_ADMIN_REQUIRE_MTLS",
        !store_kind.eq_ignore_ascii_case("in-memory"),
    );
    if require_admin_mtls {
        let Some(_) = tls_config.as_ref() else {
            anyhow::bail!(
                "PD_TSO_ADMIN_REQUIRE_MTLS requires PD_TSO_TLS_CERT_PATH and PD_TSO_TLS_KEY_PATH"
            );
        };
        if !tls_ca_configured {
            anyhow::bail!(
                "PD_TSO_ADMIN_REQUIRE_MTLS requires PD_TSO_TLS_CA_PATH for client cert verification"
            );
        }
    }

    println!("pd-tso listening on {listen_addr}");
    let mut server = Server::builder();
    if let Some(tls_config) = tls_config {
        server = server
            .tls_config(tls_config)
            .context("apply TLS server config failed")?;
    }
    server
        .add_service(PdServer::new(legacy_pd_grpc))
        .add_service(TsoServer::new(tso_grpc))
        .add_service(AdminServer::new(admin_grpc))
        .serve(listen_addr)
        .await
        .context("tonic server failed")?;
    Ok(())
}

fn print_version() {
    let release_version = env!("CARGO_PKG_VERSION");
    let git_hash = option_env!("PD_TSO_GIT_HASH").unwrap_or(DEFAULT_BUILD_METADATA);
    let git_ref = option_env!("PD_TSO_GIT_REF").unwrap_or(DEFAULT_BUILD_METADATA);
    let build_time = option_env!("PD_TSO_BUILD_TIME").unwrap_or(DEFAULT_BUILD_METADATA);

    println!("Release Version: {release_version}");
    println!("Git Commit Hash: {git_hash}");
    println!("Git Ref: {git_ref}");
    println!("UTC Build Time: {build_time}");
}

fn parse_endpoints(raw: Option<String>) -> Vec<String> {
    raw.unwrap_or_else(|| "http://127.0.0.1:2379".to_string())
        .split(',')
        .map(str::trim)
        .filter(|endpoint| !endpoint.is_empty())
        .map(ToString::to_string)
        .collect()
}

fn now_ms() -> i64 {
    use std::time::{SystemTime, UNIX_EPOCH};
    let duration = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default();
    (duration.as_millis().min(i64::MAX as u128)) as i64
}

fn parse_duration_ms_env(key: &str, default_ms: u64) -> Result<Duration> {
    let value = env::var(key)
        .ok()
        .map(|raw| raw.parse::<u64>())
        .transpose()
        .with_context(|| format!("invalid {key}"))?
        .unwrap_or(default_ms);
    Ok(Duration::from_millis(value))
}

fn parse_i64_env(key: &str, default_value: i64) -> Result<i64> {
    env::var(key)
        .ok()
        .map(|raw| raw.parse::<i64>())
        .transpose()
        .with_context(|| format!("invalid {key}"))?
        .map(|value| value.max(1))
        .map_or(Ok(default_value), Ok)
}

fn parse_legacy_proxy_max_concurrent_env(key: &str, default_value: i32) -> Result<i32> {
    let value = env::var(key)
        .ok()
        .map(|raw| raw.parse::<i32>())
        .transpose()
        .with_context(|| format!("invalid {key}"))?;
    Ok(match value {
        None => default_value,
        Some(0) => default_value,
        Some(parsed) => parsed,
    })
}

fn parse_legacy_proxy_recv_timeout_env(key: &str, default_value: Duration) -> Result<Duration> {
    let value_ms = env::var(key)
        .ok()
        .map(|raw| raw.parse::<i64>())
        .transpose()
        .with_context(|| format!("invalid {key}"))?;
    let default_ms = default_value.as_millis().min(i64::MAX as u128) as i64;
    let normalized_ms = match value_ms {
        None => default_ms,
        Some(parsed) if parsed <= 0 => default_ms,
        Some(parsed) => parsed,
    };
    Ok(Duration::from_millis(normalized_ms as u64))
}

fn parse_optional_positive_i64_env(keys: &[&str]) -> Result<Option<i64>> {
    for key in keys {
        if let Ok(raw) = env::var(key) {
            let value = raw
                .parse::<i64>()
                .with_context(|| format!("invalid {key}"))?;
            return Ok(Some(value.max(1)));
        }
    }
    Ok(None)
}

fn allocator_lease_ms_from_ttl_secs(lease_ttl_secs: i64) -> i64 {
    lease_ttl_secs.saturating_mul(1_000).max(1_000)
}

fn parse_bool_env(key: &str, default_value: bool) -> bool {
    env::var(key)
        .ok()
        .and_then(|raw| match raw.trim().to_ascii_lowercase().as_str() {
            "1" | "true" | "yes" | "on" => Some(true),
            "0" | "false" | "no" | "off" => Some(false),
            _ => None,
        })
        .unwrap_or(default_value)
}

fn load_server_tls_config() -> Result<Option<ServerTlsConfig>> {
    let cert_path = env::var("PD_TSO_TLS_CERT_PATH")
        .ok()
        .filter(|value| !value.trim().is_empty());
    let key_path = env::var("PD_TSO_TLS_KEY_PATH")
        .ok()
        .filter(|value| !value.trim().is_empty());
    let ca_path = env::var("PD_TSO_TLS_CA_PATH")
        .ok()
        .filter(|value| !value.trim().is_empty());

    let (cert_path, key_path) = match (cert_path, key_path) {
        (Some(cert_path), Some(key_path)) => (cert_path, key_path),
        (None, None) => {
            if ca_path.is_some() {
                anyhow::bail!(
                    "PD_TSO_TLS_CERT_PATH and PD_TSO_TLS_KEY_PATH must both be set when TLS CA is configured"
                );
            }
            return Ok(None);
        }
        _ => {
            anyhow::bail!(
                "PD_TSO_TLS_CERT_PATH and PD_TSO_TLS_KEY_PATH must both be set when TLS is configured"
            );
        }
    };

    let cert_pem = fs::read(&cert_path).with_context(|| format!("read cert PEM {cert_path}"))?;
    let key_pem = fs::read(&key_path).with_context(|| format!("read key PEM {key_path}"))?;

    let mut config = ServerTlsConfig::new().identity(Identity::from_pem(cert_pem, key_pem));
    if let Some(ca_path) = ca_path {
        let ca_pem = fs::read(&ca_path).with_context(|| format!("read CA PEM {ca_path}"))?;
        config = config.client_ca_root(Certificate::from_pem(ca_pem));
    }
    Ok(Some(config))
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use pd_test_util::env::{EnvVarRestore, async_env_lock, remove_env, set_env};

    use super::{
        allocator_lease_ms_from_ttl_secs, parse_legacy_proxy_max_concurrent_env,
        parse_legacy_proxy_recv_timeout_env, parse_optional_positive_i64_env,
    };

    #[test]
    fn allocator_lease_tracks_runtime_ttl() {
        assert_eq!(allocator_lease_ms_from_ttl_secs(1), 1_000);
        assert_eq!(allocator_lease_ms_from_ttl_secs(5), 5_000);
        assert_eq!(allocator_lease_ms_from_ttl_secs(30), 30_000);
    }

    #[test]
    fn startup_max_concurrent_proxy_streams_normalizes_zero_to_default() {
        let _guard = async_env_lock().blocking_lock();
        let _restore = EnvVarRestore::capture("PD_TSO_MAX_CONCURRENT_TSO_PROXY_STREAMINGS");

        set_env("PD_TSO_MAX_CONCURRENT_TSO_PROXY_STREAMINGS", "0");
        let parsed = parse_legacy_proxy_max_concurrent_env(
            "PD_TSO_MAX_CONCURRENT_TSO_PROXY_STREAMINGS",
            5_000,
        )
        .expect("parse max concurrent");
        assert_eq!(parsed, 5_000);
    }

    #[test]
    fn startup_proxy_recv_timeout_normalizes_non_positive_to_default() {
        let _guard = async_env_lock().blocking_lock();
        let _restore = EnvVarRestore::capture("PD_TSO_PROXY_RECV_FROM_CLIENT_TIMEOUT_MS");

        set_env("PD_TSO_PROXY_RECV_FROM_CLIENT_TIMEOUT_MS", "0");
        let parsed = parse_legacy_proxy_recv_timeout_env(
            "PD_TSO_PROXY_RECV_FROM_CLIENT_TIMEOUT_MS",
            Duration::from_secs(60 * 60),
        )
        .expect("parse recv timeout");
        assert_eq!(parsed, Duration::from_secs(60 * 60));
    }

    #[test]
    fn startup_max_gap_reset_ts_reads_canonical_and_legacy_env_keys() {
        let _guard = async_env_lock().blocking_lock();
        let [_canonical_restore, _legacy_restore] =
            EnvVarRestore::capture_many(["PD_TSO_MAX_GAP_RESET_TS", "PD_TSO_MAX_GAP_RESET_TS_MS"]);

        remove_env("PD_TSO_MAX_GAP_RESET_TS");
        set_env("PD_TSO_MAX_GAP_RESET_TS_MS", "4500");
        let legacy_only = parse_optional_positive_i64_env(&[
            "PD_TSO_MAX_GAP_RESET_TS",
            "PD_TSO_MAX_GAP_RESET_TS_MS",
        ])
        .expect("parse legacy key");
        assert_eq!(legacy_only, Some(4_500));

        set_env("PD_TSO_MAX_GAP_RESET_TS", "3300");
        let canonical_preferred = parse_optional_positive_i64_env(&[
            "PD_TSO_MAX_GAP_RESET_TS",
            "PD_TSO_MAX_GAP_RESET_TS_MS",
        ])
        .expect("parse canonical key");
        assert_eq!(canonical_preferred, Some(3_300));
    }

    #[test]
    fn startup_max_gap_reset_ts_normalizes_non_positive_to_one() {
        let _guard = async_env_lock().blocking_lock();
        let _restore = EnvVarRestore::capture("PD_TSO_MAX_GAP_RESET_TS");

        set_env("PD_TSO_MAX_GAP_RESET_TS", "0");
        let parsed = parse_optional_positive_i64_env(&[
            "PD_TSO_MAX_GAP_RESET_TS",
            "PD_TSO_MAX_GAP_RESET_TS_MS",
        ])
        .expect("parse max gap reset ts");
        assert_eq!(parsed, Some(1));
    }
}