use anyhow::{Context, Result, ensure};
use clap::Parser;
use log::{debug, error, info};
use tici::{config::MetaServiceConfig, meta::server::MetaServiceServer};
use tici_common::{config::LogConfig, logger::init_logger};
use tici_shard::config::WorkerNodeConfig;
use tici_system::readable::ReadableSize;

#[derive(Parser, Debug)]
#[command(name = "TiCI Server", about = "TiCI Server", version = tici_shard::utils::build_info())]
struct TiCICommand {
    #[command(subcommand)]
    command: TiCISubCommands,
    /// The log file. If this configuration item is not set, logs are output to
    /// "stderr" by default. If this configuration item is set, logs are output
    /// to the corresponding file.
    #[arg(long = "log-file", global = true, value_name = "FILE")]
    log_file: Option<String>,
    /// The log level, such as "info", "warn", "error".
    #[arg(long = "log-level", global = true, value_name = "LEVEL")]
    log_level: Option<String>,
    /// The maximum size for each rotated log file, such as "256MB".
    #[arg(long = "log-size", global = true, value_name = "SIZE")]
    log_size: Option<ReadableSize>,
    /// The number of rotated log files to keep.
    #[arg(long = "log-count", global = true, value_name = "COUNT")]
    log_count: Option<usize>,
}

impl TiCICommand {
    fn overwrite_log_by_cmdline(&self, log_config: &mut LogConfig) {
        log_config.overwrite_by_cmdline(
            &self.log_file,
            &self.log_level,
            &self.log_size,
            &self.log_count,
        );
    }

    fn load_log_config(&self) -> Result<LogConfig> {
        let mut log_config = match &self.command {
            TiCISubCommands::Meta(m) => {
                let cfg = m
                    .load_config()
                    .with_context(|| format!("Failed to load meta config from {}", m.config))?;
                cfg.logger
            }
            TiCISubCommands::Worker(w) => {
                let cfg = w
                    .load_config()
                    .with_context(|| format!("Failed to load worker config from {}", w.config))?;
                cfg.logger
            }
        };
        self.overwrite_log_by_cmdline(&mut log_config);
        Ok(log_config)
    }

    fn validate(&self) -> Result<()> {
        let cfg_path = match &self.command {
            TiCISubCommands::Meta(MetaArgs { config, .. }) => config,
            TiCISubCommands::Worker(WorkerArgs { config, .. }) => config,
        };
        ensure!(!cfg_path.is_empty(), "Config path cannot be empty");
        Ok(())
    }

    async fn run(&self) -> Result<()> {
        match &self.command {
            TiCISubCommands::Meta(m) => m.run().await,
            TiCISubCommands::Worker(w) => w.run().await,
        }
    }
}

#[derive(clap::Subcommand, Debug)]
enum TiCISubCommands {
    /// Start a TiCI Meta Node
    Meta(MetaArgs),
    /// Start a TiCI Worker Node
    Worker(WorkerArgs),
}

#[derive(clap::Args, Debug)]
struct MetaArgs {
    /// Path to configuration file
    #[arg(long, value_name = "FILE")]
    config: String,
    /// Server listen host
    #[arg(short = 'H', long, value_name = "HOST")]
    host: Option<String>,
    /// Server listen port
    #[arg(short = 'P', long, value_name = "PORT")]
    port: Option<String>,
    /// Status port
    #[arg(long = "status-port", value_name = "STATUS_PORT")]
    status_port: Option<String>,
    /// Advertised host for external access
    #[arg(long = "advertise-host", value_name = "ADVERTISE_HOST")]
    advertise_host: Option<String>,
    /// PD addresses, multiple separated by comma
    #[arg(long = "pd-addr", value_name = "PD_ADDR")]
    pd_addr: Option<String>,
    /// TiDB server dsn, multiple separated by comma
    #[arg(long = "tidb-addr", value_name = "TIDB_SERVER")]
    tidb_addr: Option<String>,
}

impl MetaArgs {
    fn load_config(&self) -> Result<MetaServiceConfig> {
        MetaServiceConfig::load(self.config.as_str())
    }

    fn get_config(&self) -> Result<MetaServiceConfig> {
        let mut cfg = self.load_config()?;
        debug!("before address config, config={:?}", cfg);
        cfg.server.adjust_by_cmdline(
            self.host.as_deref().unwrap_or(""),
            self.port.as_deref().unwrap_or(""),
            self.status_port.as_deref().unwrap_or(""),
            self.advertise_host.as_deref().unwrap_or(""),
            self.pd_addr.as_deref().unwrap_or(""),
        )?;
        cfg.tidb_server.overwrite_by_cmdline(&self.tidb_addr);
        info!("after address config, config={:?}", cfg);
        Ok(cfg)
    }

    async fn run(&self) -> Result<()> {
        let cfg = self.get_config()?;

        // Start the server
        let (tx, rx) = tokio::sync::oneshot::channel::<()>();
        if cfg.security.is_valid() {
            info!("use aws_lc_rs for rustls::crypto");
            use rustls::crypto::CryptoProvider;
            CryptoProvider::install_default(rustls::crypto::aws_lc_rs::default_provider())
                .expect("Failed to install crypto provider");
        }
        let mut server = MetaServiceServer::new(cfg).await.unwrap();
        // let meta service wait to become leader
        server.wait_become_leader().await;
        server.start(rx).await?;
        info!("Running...");
        tokio::signal::ctrl_c().await.ok();
        info!("Stopping...");
        server.stop(tx).await?;
        Ok(())
    }
}

#[derive(clap::Args, Debug)]
struct WorkerArgs {
    /// Path to configuration file
    #[arg(long, value_name = "FILE")]
    config: String,
    /// Server listen host
    #[arg(short = 'H', long, value_name = "HOST")]
    host: Option<String>,
    /// Server listen port
    #[arg(short = 'P', long, value_name = "PORT")]
    port: Option<String>,
    /// Status port
    #[arg(long = "status-port", value_name = "STATUS_PORT")]
    status_port: Option<String>,
    /// Advertised host for external access
    #[arg(long = "advertise-host", value_name = "ADVERTISE_HOST")]
    advertise_host: Option<String>,
    /// PD addresses, multiple separated by comma
    #[arg(long = "pd-addr", value_name = "PD_ADDR")]
    pd_addr: Option<String>,
    /// Path to the data directory
    #[arg(long, value_name = "DIR")]
    data_dir: Option<String>,
}

impl WorkerArgs {
    fn load_config(&self) -> Result<WorkerNodeConfig> {
        WorkerNodeConfig::load(self.config.as_str())
    }

    fn get_config(&self) -> Result<WorkerNodeConfig> {
        let mut cfg = self.load_config()?;
        debug!("before address config, config={:?}", cfg);
        // override the server configs by cmdline
        cfg.server.adjust_by_cmdline(
            self.host.as_deref().unwrap_or(""),
            self.port.as_deref().unwrap_or(""),
            self.status_port.as_deref().unwrap_or(""),
            self.advertise_host.as_deref().unwrap_or(""),
            self.pd_addr.as_deref().unwrap_or(""),
        )?;
        // override the storage configs by cmdline
        cfg.storage.overwrite_by_cmdline(&self.data_dir);
        info!("after address config, config={:?}", cfg);
        Ok(cfg)
    }

    async fn run(&self) -> Result<()> {
        let cfg = self.get_config()?;

        // Start the server
        let (tx, rx) = tokio::sync::oneshot::channel::<()>();
        if cfg.security.is_valid() {
            info!("use aws_lc_rs for rustls::crypto");
            use rustls::crypto::CryptoProvider;
            CryptoProvider::install_default(rustls::crypto::aws_lc_rs::default_provider())
                .expect("Failed to install crypto provider");
        }
        let mut server = tici::worker::server::WorkerServiceServer::new(cfg).await?;
        server.start(rx).await?;
        info!("Running...");
        tokio::signal::ctrl_c().await.ok();
        info!("Stopping...");
        server.stop(tx).await?;
        Ok(())
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    // Setup CLI argument parsing with clap derive
    let cmd = TiCICommand::parse();
    cmd.validate()?;
    let log_config = cmd.load_log_config()?;
    let _logger_handle = init_logger(&log_config)?;

    info!("set panic hook");
    tici_system::panic_hook::setup_panic_hook();
    info!("set panic hook end");

    info!("Welcome to TiCI Server");
    info!("running with args: {:?}", cmd);
    info!(
        "build info={}",
        tici_shard::utils::build_info().replace("\n", " ")
    );

    // Start the server
    cmd.run()
        .await
        .inspect_err(|e| error!("Failed to start server, err={:?}", e))
}