use anyhow::Result;
use clap::Parser;
use flaredb_proto::kvrpc::kv_cas_server::KvCasServer;
use flaredb_proto::kvrpc::kv_raw_server::KvRawServer;
use flaredb_proto::pdpb::pd_client::PdClient as LegacyPdClient;
use flaredb_proto::pdpb::{ListRegionsRequest, RegisterStoreRequest};
use flaredb_proto::raft_server::raft_service_server::RaftServiceServer;
use flaredb_proto::sqlrpc::sql_service_server::SqlServiceServer;
use flaredb_server::config::{self, Config, NamespaceManager};
use flaredb_storage::rocks_engine::RocksEngine;
use flaredb_types::RegionMeta;
use metrics_exporter_prometheus::PrometheusBuilder;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::time::{sleep, Duration};
use tonic::transport::{Certificate, Channel, Identity, Server, ServerTlsConfig};
use tonic_health::server::health_reporter;
use tracing::{info, warn}; // Import warn
use tracing_subscriber::EnvFilter; // Import anyhow

mod heartbeat;
mod merkle;
mod pd_client;
mod raft_service;
mod rest;
mod service;
mod sql_service;
mod store;

use pd_client::{PdClient as ChainfirePdClient, PdEvent};

const RAFT_GRPC_MESSAGE_SIZE: usize = 64 * 1024 * 1024;

async fn connect_pd_with_retry(
    pd_endpoints: &[String],
    attempts: u32,
    delay: Duration,
) -> Option<ChainfirePdClient> {
    let mut last_error = None;

    for attempt in 1..=attempts {
        match ChainfirePdClient::connect_any(pd_endpoints).await {
            Ok(client) => return Some(client),
            Err(err) => {
                last_error = Some(err.to_string());
                let protocol_mismatch = last_error
                    .as_deref()
                    .map(|msg| msg.contains("Unimplemented"))
                    .unwrap_or(false);
                warn!(
                    attempt,
                    attempts,
                    ?pd_endpoints,
                    error = last_error.as_deref().unwrap_or("unknown"),
                    "Failed to connect to FlareDB PD"
                );
                if protocol_mismatch {
                    warn!(
                        ?pd_endpoints,
                        "PD endpoint does not speak ChainFire; falling back to legacy PD"
                    );
                    return None;
                }
                if attempt < attempts {
                    sleep(delay).await;
                }
            }
        }
    }

    warn!(
        ?pd_endpoints,
        error = last_error.as_deref().unwrap_or("unknown"),
        "Exhausted FlareDB PD connection retries"
    );
    None
}

async fn connect_legacy_pd_with_retry(
    pd_endpoints: &[String],
    attempts: u32,
    delay: Duration,
) -> Option<(String, LegacyPdClient<Channel>)> {
    let mut last_error = None;

    for attempt in 1..=attempts {
        for endpoint in pd_endpoints {
            let transport = if endpoint.starts_with("http") {
                endpoint.clone()
            } else {
                format!("http://{}", endpoint)
            };
            match LegacyPdClient::connect(transport.clone()).await {
                Ok(client) => return Some((endpoint.clone(), client)),
                Err(err) => {
                    last_error = Some(format!("{}: {}", endpoint, err));
                }
            }
        }

        warn!(
            attempt,
            attempts,
            ?pd_endpoints,
            error = last_error.as_deref().unwrap_or("unknown"),
            "Failed to connect to legacy FlareDB PD"
        );

        if attempt < attempts {
            sleep(delay).await;
        }
    }

    warn!(
        ?pd_endpoints,
        error = last_error.as_deref().unwrap_or("unknown"),
        "Exhausted legacy FlareDB PD connection retries"
    );
    None
}

#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// Configuration file path
    #[arg(short, long)]
    config: Option<PathBuf>,

    /// Store ID (overrides config)
    #[arg(long)]
    store_id: Option<u64>,

    /// Listen address for gRPC API (overrides config)
    #[arg(long)]
    addr: Option<String>,

    /// Listen address for HTTP REST API (overrides config)
    #[arg(long)]
    http_addr: Option<String>,

    /// Data directory for RocksDB (overrides config)
    #[arg(long)]
    data_dir: Option<PathBuf>,

    /// FlareDB PD address (overrides config)
    #[arg(long)]
    pd_addr: Option<String>,

    /// Initial cluster peers in id=host:port format (overrides config)
    #[arg(long = "peer")]
    peers: Vec<String>,

    /// Namespace modes in format namespace=strong|eventual (overrides config)
    #[arg(long = "namespace-mode")]
    namespace_modes: Vec<String>,

    /// Default mode for implicitly created namespaces (strong|eventual)
    #[arg(long = "default-namespace-mode")]
    default_namespace_mode: Option<String>,

    /// Log level (e.g., "info", "debug", "trace")
    #[arg(long)]
    log_level: Option<String>,

    /// Metrics port for Prometheus scraping
    #[arg(long, default_value = "9092")]
    metrics_port: u16,
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args = Args::parse();

    // Load configuration
    let mut settings = ::config::Config::builder()
        .add_source(::config::Environment::with_prefix("FLAREDB").separator("_"));

    // Add config file if specified
    if let Some(config_path) = &args.config {
        info!("Loading config from file: {}", config_path.display());
        settings = settings.add_source(::config::File::from(config_path.as_path()));
    }

    let loaded_config: Config = settings
        .build()?
        .try_deserialize()
        .map_err(|e| anyhow::anyhow!("Failed to load configuration: {}", e))?;
    let resolved_pd_endpoints = loaded_config.resolved_pd_endpoints();

    // Apply CLI overrides to the loaded configuration
    let config = Config {
        store_id: args.store_id.unwrap_or(loaded_config.store_id),
        addr: args
            .addr
            .map(|s| s.parse().unwrap_or(loaded_config.addr))
            .unwrap_or(loaded_config.addr),
        http_addr: args
            .http_addr
            .map(|s| s.parse().unwrap_or(loaded_config.http_addr))
            .unwrap_or(loaded_config.http_addr),
        data_dir: args.data_dir.unwrap_or(loaded_config.data_dir),
        pd_addr: args.pd_addr.clone().unwrap_or(loaded_config.pd_addr),
        pd_endpoints: if let Some(pd_addr) = args.pd_addr {
            vec![pd_addr]
        } else {
            resolved_pd_endpoints
        },
        peers: if args.peers.is_empty() {
            loaded_config.peers
        } else {
            let mut peers_map = HashMap::new();
            for p in args.peers {
                if let Some((id_str, addr_str)) = p.split_once('=') {
                    if let Ok(id) = id_str.parse::<u64>() {
                        peers_map.insert(id, addr_str.parse().unwrap());
                    }
                }
            }
            peers_map
        },
        namespace_modes: if args.namespace_modes.is_empty() {
            loaded_config.namespace_modes
        } else {
            config::parse_namespace_modes(&args.namespace_modes)?
        },
        default_namespace_mode: if let Some(mode_str) = args.default_namespace_mode {
            config::parse_mode(&mode_str)?
        } else {
            loaded_config.default_namespace_mode
        },
        log_level: args.log_level.unwrap_or(loaded_config.log_level),
        tls: loaded_config.tls,
    };

    // Initialize tracing
    init_logging(&config.log_level);

    // Initialize Prometheus metrics exporter
    let metrics_addr = format!("0.0.0.0:{}", args.metrics_port);
    let builder = PrometheusBuilder::new();
    builder
        .with_http_listener(metrics_addr.parse::<std::net::SocketAddr>()?)
        .install()
        .expect("Failed to install Prometheus metrics exporter");

    info!(
        "Prometheus metrics available at http://{}/metrics",
        metrics_addr
    );

    info!("FlareDB server listening on {}", config.addr);

    let addr = config.addr;
    let server_config = Arc::new(config);
    let namespace_manager = Arc::new(NamespaceManager::from_config(&server_config));

    // Parse peer addresses for cluster membership
    let mut voters = vec![server_config.store_id];
    let mut peer_addrs: HashMap<u64, String> = HashMap::new();
    // Add self address
    peer_addrs.insert(server_config.store_id, server_config.addr.to_string());
    for (id, addr) in server_config.peers.clone() {
        if id != server_config.store_id {
            voters.push(id);
            peer_addrs.insert(id, addr.to_string());
        }
    }
    let peer_addrs = Arc::new(peer_addrs);

    let engine = Arc::new(RocksEngine::new(server_config.data_dir.to_str().unwrap())?);
    let store = Arc::new(store::Store::new(
        server_config.store_id,
        engine.clone(),
        server_config.clone(),     // Pass server_config
        namespace_manager.clone(), // Pass namespace manager
        peer_addrs.clone(),
    ));

    let service =
        service::KvServiceImpl::new(engine.clone(), namespace_manager.clone(), store.clone());
    let raft_service = raft_service::RaftServiceImpl::new(store.clone(), server_config.store_id);
    let pd_endpoints = server_config.resolved_pd_endpoints();

    println!("Connecting to FlareDB PD at {:?}...", pd_endpoints);
    if let Some(mut pd_client) =
        connect_pd_with_retry(&pd_endpoints, 30, Duration::from_secs(2)).await
    {
        println!(
            "Connected to ChainFire. Cluster ID: {}",
            pd_client.cluster_id()
        );

        // Register this store with the PD
        if let Err(e) = pd_client
            .register_store(server_config.store_id, server_config.addr.to_string())
            .await
        {
            eprintln!("Failed to register store: {}", e);
        } else {
            println!("Store {} registered with PD", server_config.store_id);
        }

        // Start watching for metadata changes from ChainFire
        let mut event_rx = pd_client.subscribe();
        if let Err(e) = pd_client.start_watch().await {
            eprintln!("Failed to start PD watch: {}", e);
        } else {
            println!("Started watching PD for metadata changes");

            // Spawn a background task to handle PD events
            let store_clone_for_events = store.clone();
            let server_config_for_events = server_config.clone(); // Capture server_config
            tokio::spawn(async move {
                while let Ok(event) = event_rx.recv().await {
                    match event {
                        PdEvent::RegionUpdated(region) => {
                            info!("Region {} updated via watch", region.id);
                            let metas = vec![(
                                RegionMeta {
                                    id: region.id,
                                    start_key: region.start_key,
                                    end_key: region.end_key,
                                },
                                if region.peers.is_empty() {
                                    vec![server_config_for_events.store_id]
                                } else {
                                    region.peers
                                },
                            )];
                            if let Err(e) = store_clone_for_events.refresh_regions(metas).await {
                                warn!("Failed to refresh region from event: {}", e);
                            }
                        }
                        PdEvent::RegionRemoved(id) => {
                            info!("Region {} removed via watch", id);
                        }
                        PdEvent::StoreUpdated(store_info) => {
                            info!("Store {} updated via watch", store_info.id);
                        }
                        PdEvent::StoreRemoved(id) => {
                            info!("Store {} removed via watch", id);
                        }
                    }
                }
            });
        }

        // Initialize default region if this is the first node
        if let Err(e) = pd_client.init_default_region(voters.clone()).await {
            warn!("Failed to init default region: {}", e);
        }

        // Fetch initial region metadata from PD (from cache)
        let regions = pd_client.list_regions().await;
        let mut region_metas = Vec::new();
        for r in regions {
            let region_voters = if r.peers.is_empty() {
                voters.clone()
            } else {
                r.peers.clone()
            };
            region_metas.push((
                RegionMeta {
                    id: r.id,
                    start_key: r.start_key,
                    end_key: r.end_key,
                },
                region_voters,
            ));
        }
        if region_metas.is_empty() {
            region_metas.push((
                RegionMeta {
                    id: 1,
                    start_key: Vec::new(),
                    end_key: Vec::new(),
                },
                voters.clone(),
            ));
        }
        if let Err(e) = store.bootstrap_regions(region_metas.clone()).await {
            warn!("failed to bootstrap regions: {}", e);
        }

        // Background task: heartbeat and refresh regions from PD
        let store_clone = store.clone();
        let pd_endpoints_for_task = pd_endpoints.clone();
        let store_id = server_config.store_id;
        let server_addr_string = server_config.addr.to_string();
        tokio::spawn(async move {
            let mut last_reported_leaders: HashMap<u64, u64> = HashMap::new();
            let client = Arc::new(Mutex::new(
                ChainfirePdClient::connect_any(&pd_endpoints_for_task)
                    .await
                    .ok(),
            ));

            loop {
                let mut guard = client.lock().await;
                if let Some(ref mut c) = *guard {
                    // Send heartbeat
                    let heartbeat_ok = match c.heartbeat(store_id, server_addr_string.clone()).await
                    {
                        Ok(_) => true,
                        Err(e) => {
                            warn!("Heartbeat failed: {}", e);
                            false
                        }
                    };

                    // If heartbeat failed, try to reconnect on next cycle
                    if !heartbeat_ok {
                        *guard = None;
                        continue;
                    }

                    // Report observed leader status so routing metadata converges
                    // even when followers are the first nodes to notice a leadership change.
                    let region_ids = store_clone.list_region_ids().await;
                    let mut observed_regions = HashMap::new();
                    for region_id in region_ids {
                        if let Some(node) = store_clone.get_raft_node(region_id).await {
                            if let Some(observed_leader) = node.leader_id().await {
                                observed_regions.insert(region_id, observed_leader);
                                let previous = last_reported_leaders.get(&region_id).copied();
                                if previous != Some(observed_leader) {
                                    info!(
                                        region_id,
                                        previous_leader = ?previous,
                                        observed_leader,
                                        "Reporting FlareDB region leader to PD"
                                    );
                                }
                                match c.report_leader(region_id, observed_leader).await {
                                    Ok(_) => {
                                        last_reported_leaders.insert(region_id, observed_leader);
                                    }
                                    Err(e) => {
                                        warn!(
                                            region_id,
                                            observed_leader,
                                            error = %e,
                                            "Report leader failed"
                                        );
                                    }
                                }
                            }
                        }
                    }
                    last_reported_leaders
                        .retain(|region_id, _| observed_regions.contains_key(region_id));

                    // Refresh regions from PD (from cache, updated via watch)
                    let regions = c.list_regions().await;
                    let metas: Vec<_> = regions
                        .into_iter()
                        .map(|r| {
                            let region_voters = if r.peers.is_empty() {
                                voters.clone()
                            } else {
                                r.peers.clone()
                            };
                            (
                                RegionMeta {
                                    id: r.id,
                                    start_key: r.start_key,
                                    end_key: r.end_key,
                                },
                                region_voters,
                            )
                        })
                        .collect();
                    if !metas.is_empty() {
                        if let Err(e) = store_clone.refresh_regions(metas).await {
                            warn!("refresh regions failed: {}", e);
                        }
                    }
                } else {
                    // Try to reconnect
                    if let Ok(new_client) =
                        ChainfirePdClient::connect_any(&pd_endpoints_for_task).await
                    {
                        info!("Reconnected to PD");
                        *guard = Some(new_client);
                    }
                }

                sleep(Duration::from_secs(10)).await;
            }
        });
    } else if let Some((legacy_pd_addr, mut legacy_pd_client)) =
        connect_legacy_pd_with_retry(&pd_endpoints, 3, Duration::from_secs(1)).await
    {
        info!(pd_addr = %legacy_pd_addr, "Connected to legacy FlareDB PD");

        match legacy_pd_client
            .register_store(RegisterStoreRequest {
                addr: server_config.addr.to_string(),
                store_id: server_config.store_id,
            })
            .await
        {
            Ok(resp) => {
                let resp = resp.into_inner();
                if resp.store_id != 0 && resp.store_id != server_config.store_id {
                    warn!(
                        expected_store_id = server_config.store_id,
                        assigned_store_id = resp.store_id,
                        "legacy PD assigned a different store id than local config"
                    );
                }
            }
            Err(err) => warn!("failed to register with legacy PD: {}", err),
        }

        let mut region_metas = Vec::new();
        match legacy_pd_client.list_regions(ListRegionsRequest {}).await {
            Ok(resp) => {
                for region in resp.into_inner().regions {
                    let voters = if region.peers.is_empty() || region.peers.len() < voters.len() {
                        voters.clone()
                    } else {
                        region.peers.clone()
                    };
                    region_metas.push((
                        RegionMeta {
                            id: region.id,
                            start_key: region.start_key,
                            end_key: region.end_key,
                        },
                        voters,
                    ));
                }
            }
            Err(err) => warn!("failed to list regions from legacy PD: {}", err),
        }

        if region_metas.is_empty() {
            region_metas.push((
                RegionMeta {
                    id: 1,
                    start_key: Vec::new(),
                    end_key: Vec::new(),
                },
                voters.clone(),
            ));
        }

        if let Err(e) = store.bootstrap_regions(region_metas).await {
            warn!("failed to bootstrap regions from legacy PD: {}", e);
        }

        heartbeat::start_heartbeat(
            legacy_pd_addr,
            store.clone(),
            server_config.addr.to_string(),
            server_config.store_id,
        )
        .await;
    } else {
        info!("Starting in standalone mode with default region...");
        let _ = store
            .bootstrap_regions(vec![(
                RegionMeta {
                    id: 1,
                    start_key: Vec::new(),
                    end_key: Vec::new(),
                },
                voters.clone(),
            )])
            .await;
    }

    // Health check service for K8s liveness/readiness probes
    let (mut health_reporter, health_service) = health_reporter();
    health_reporter
        .set_serving::<KvRawServer<service::KvServiceImpl>>()
        .await;
    health_reporter
        .set_serving::<KvCasServer<service::KvServiceImpl>>()
        .await;

    info!("FlareDB server starting with health checks enabled");

    // Create SQL service
    let sql_service =
        sql_service::SqlServiceImpl::new(server_config.addr.to_string(), pd_endpoints.clone());

    // Configure TLS if enabled
    let mut server = Server::builder();

    if let Some(tls_config) = &server_config.tls {
        info!("TLS enabled, loading certificates...");
        let cert = tokio::fs::read(&tls_config.cert_file)
            .await
            .map_err(|e| anyhow::anyhow!("Failed to read cert file: {}", e))?;
        let key = tokio::fs::read(&tls_config.key_file)
            .await
            .map_err(|e| anyhow::anyhow!("Failed to read key file: {}", e))?;
        let server_identity = Identity::from_pem(cert, key);

        let tls = if tls_config.require_client_cert {
            info!("mTLS enabled, requiring client certificates");
            let ca_cert = tokio::fs::read(tls_config.ca_file.as_ref().ok_or_else(|| {
                anyhow::anyhow!("ca_file required when require_client_cert=true")
            })?)
            .await
            .map_err(|e| anyhow::anyhow!("Failed to read CA file: {}", e))?;
            let ca = Certificate::from_pem(ca_cert);

            ServerTlsConfig::new()
                .identity(server_identity)
                .client_ca_root(ca)
        } else {
            info!("TLS-only mode, client certificates not required");
            ServerTlsConfig::new().identity(server_identity)
        };

        server = server
            .tls_config(tls)
            .map_err(|e| anyhow::anyhow!("Failed to configure TLS: {}", e))?;
        info!("TLS configuration applied successfully");
    } else {
        info!("TLS disabled, running in plain-text mode");
    }

    // gRPC server
    let grpc_server = server
        .add_service(health_service)
        .add_service(KvRawServer::new(service.clone()))
        .add_service(KvCasServer::new(service))
        .add_service(
            RaftServiceServer::new(raft_service)
                .max_decoding_message_size(RAFT_GRPC_MESSAGE_SIZE)
                .max_encoding_message_size(RAFT_GRPC_MESSAGE_SIZE),
        )
        .add_service(SqlServiceServer::new(sql_service))
        .serve(addr);

    // HTTP REST API server
    let http_addr = server_config.http_addr;
    let rest_state = rest::RestApiState {
        server_addr: server_config.addr.to_string(),
        pd_endpoints: pd_endpoints.clone(),
        store_id: server_config.store_id,
        configured_peers: (*peer_addrs).clone(),
    };
    let rest_app = rest::build_router(rest_state);
    let http_listener = tokio::net::TcpListener::bind(&http_addr).await?;

    info!(http_addr = %http_addr, "HTTP REST API server starting");

    let http_server = async move {
        axum::serve(http_listener, rest_app)
            .await
            .map_err(|e| anyhow::anyhow!("HTTP server error: {}", e))
    };

    // Run both servers concurrently
    tokio::select! {
        result = grpc_server => {
            result?;
        }
        result = http_server => {
            result?;
        }
    }

    Ok(())
}

fn init_logging(level: &str) {
    tracing_subscriber::fmt()
        .with_env_filter(
            EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(level)),
        )
        .init();
}
