// Copyright 2022 PingCAP, Inc.

#[macro_use]
extern crate serde_derive;

mod archive;
mod backup;
mod check_columnar;
mod check_table;
mod columnar;
mod common;
mod dfsgc;
mod encryption;
mod fts;
mod http;
mod materialize_config;
mod mvcc;
mod pack_backup;
mod recovery;
mod region;
mod resolve_lock;
mod restore;
mod schema_file;
mod sst;
mod stats;
mod test;
mod txn_file;
mod txn_log;
mod unpack_backup;
mod unsafe_recover;

use std::{env, fs::OpenOptions, io, sync::LazyLock};

use backup::{ShowBackupListArgs, execute_show_backup_list};
use clap::{Args, Parser, Subcommand};
use encryption::{EncryptionCommand, execute_encryption_command};
use native_br::common::step_to_stdout;
use region::{RegionCommand, execute_region_command};
use slog::Drain;

use crate::{
    Commands::*,
    archive::{ArchiveArgs, ShowArchiveArgs, execute_archive, execute_show_archive},
    backup::{BackupArgs, ShowBackupArgs, execute_backup, execute_show_backup},
    check_columnar::{CheckColumnarArgs, execute_check_columnar},
    check_table::{CheckTableArgs, execute_check_table},
    columnar::{ShowColumnarArgs, execute_show_columnar},
    dfsgc::{DfsGcArgs, execute_dfsgc},
    fts::{ShowFtsArgs, execute_show_fts},
    http::HttpArgs,
    materialize_config::{MaterializeConfigArgs, execute_materialize_config},
    mvcc::{MvccArgs, execute_mvcc},
    pack_backup::{PackBackupArgs, execute_pack_backup},
    recovery::{RecoveryArgs, execute_recovery},
    resolve_lock::{ResolveLockArgs, execute_resolve_lock},
    restore::{RestoreCommand, execute_restore_command},
    schema_file::{ShowSchemaArgs, execute_show_schema},
    sst::{ScanBadTableFileArgs, ShowSstArgs, execute_scan_bad_table, execute_show_sst},
    stats::{StatsArgs, execute_stats},
    test::{TestArgs, execute_test},
    txn_file::{ShowTxnChunkArgs, execute_show_txn_chunk},
    txn_log::{ShowTxnLogArgs, execute_show_txn_log},
    unpack_backup::{UnpackBackupArgs, execute_unpack_backup},
    unsafe_recover::{UnsafeRecoverArgs, execute_unsafe_recover},
};

fn main() {
    // Suppress logs for encryption command if LOG_LEVEL is not explicitly set
    if std::env::args().any(|arg| arg == "encryption") && std::env::var("LOG_LEVEL").is_err() {
        std::env::set_var("LOG_LEVEL", "error");
    }
    init_logger();
    step_to_stdout();
    let cli: Cli = Cli::parse();
    cli.apply_tls_env_overrides();
    match cli.command {
        DfsGc(dfsgc_arg) => {
            execute_dfsgc(dfsgc_arg);
        }
        UnsafeRecover(unsafe_recover) => {
            execute_unsafe_recover(unsafe_recover);
        }
        Backup(backup_args) => {
            execute_backup(backup_args);
        }
        Restore(restore_cmd) => {
            execute_restore_command(restore_cmd);
        }
        Archive(archive_args) => {
            execute_archive(archive_args);
        }
        PackBackup(pack_backup_args) => {
            execute_pack_backup(pack_backup_args);
        }
        UnpackBackup(unpack_backup_args) => {
            execute_unpack_backup(unpack_backup_args);
        }
        Stats(stats_arg) => {
            execute_stats(stats_arg);
        }
        ResolveLock(arg) => {
            execute_resolve_lock(arg);
        }
        CheckTable(args) => {
            execute_check_table(args);
        }
        Mvcc(args) => execute_mvcc(args),
        Show(args) => {
            execute_show(args);
        }
        Http(args) => {
            http::execute_http(args);
        }
        Recovery(args) => {
            execute_recovery(args);
        }
        Test(args) => {
            execute_test(args);
        }
        Region(region_cmd) => {
            execute_region_command(region_cmd);
        }
        CheckColumnar(args) => {
            execute_check_columnar(args);
        }
        MaterializeConfig(args) => {
            execute_materialize_config(args);
        }
        Encryption(encryption_cmd) => {
            execute_encryption_command(encryption_cmd);
        }
    }
}

impl Cli {
    fn tls_override_values(&self) -> Result<Option<(&str, &str, &str)>, &'static str> {
        match (
            self.ca_path.as_deref(),
            self.cert_path.as_deref(),
            self.key_path.as_deref(),
        ) {
            (Some(ca), Some(cert), Some(key)) => Ok(Some((ca, cert, key))),
            (None, None, None) => Ok(None),
            _ => Err("--ca-path, --cert-path, and --key-path must be set together"),
        }
    }

    fn apply_tls_env_overrides(&self) {
        let Some((ca, cert, key)) = self.tls_override_values().unwrap_or_else(|err| {
            eprintln!("{err}");
            std::process::exit(1);
        }) else {
            return;
        };

        env::set_var("CA_PATH", ca);
        env::set_var("CERT_PATH", cert);
        env::set_var("KEY_PATH", key);
    }
}

fn init_logger() {
    let output = env::var("LOG_FILE").ok();
    let level = tikv_util::logger::get_level_by_string(
        &env::var("LOG_LEVEL").unwrap_or_else(|_| "info".to_owned()),
    )
    .unwrap();
    let append_instead_truncate = env::var("LOG_APPEND").is_ok();

    match output {
        Some(log_file) => {
            let f = OpenOptions::new()
                .create(true)
                .write(!append_instead_truncate)
                .truncate(!append_instead_truncate)
                .append(append_instead_truncate)
                .open(log_file)
                .unwrap();
            init_logger_impl(f, level);
        }
        None => init_logger_impl(io::stdout(), level),
    };
}

fn init_logger_impl<W: 'static + io::Write + Send>(writer: W, level: slog::Level) {
    let drain: Box<dyn slog::Drain<Ok = (), Err = std::io::Error> + Send> =
        if atty::is(atty::Stream::Stdout) {
            let dec = slog_term::TermDecorator::new().build();
            Box::new(slog_term::CompactFormat::new(dec).build())
        } else {
            let dec = slog_term::PlainDecorator::new(writer);
            Box::new(slog_term::CompactFormat::new(dec).build())
        };
    let drain = std::sync::Mutex::new(drain).filter_level(level).fuse();
    let logger = slog::Logger::root(drain, slog::o!());
    slog_global::set_global(logger);
}

static VERSION_INFO: LazyLock<String> = LazyLock::new(|| {
    let build_timestamp = option_env!("TIKV_BUILD_TIME");
    tikv::tikv_version_info(build_timestamp)
});

#[derive(Parser)]
#[clap(author, version = &**VERSION_INFO, about, long_about = None)]
pub struct Cli {
    #[clap(long, global = true)]
    /// Set the CA certificate path (overrides env `CA_PATH`).
    pub ca_path: Option<String>,

    #[clap(long, global = true)]
    /// Set the certificate path (overrides env `CERT_PATH`).
    pub cert_path: Option<String>,

    #[clap(long, global = true)]
    /// Set the private key path (overrides env `KEY_PATH`).
    pub key_path: Option<String>,

    #[clap(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
pub enum Commands {
    /// Scan and mark the unused DFS files as deleted.
    DfsGc(DfsGcArgs),
    /// Unsafely recover the cluster by directly modifying the data on the raft
    /// engine.
    UnsafeRecover(UnsafeRecoverArgs),
    /// Backup backups the cluster.
    Backup(BackupArgs),
    /// Restore a backup.
    Restore(RestoreCommand),
    /// Archive old backups.
    Archive(ArchiveArgs),
    /// Pack backup files into a compact representation.
    PackBackup(PackBackupArgs),
    /// Unpack backup files from another cluster to the current cluster.
    UnpackBackup(UnpackBackupArgs),
    /// Stats s3 objects.
    Stats(StatsArgs),
    /// Resolve lock
    ResolveLock(ResolveLockArgs),
    /// CheckTable check data consistency on each table.
    CheckTable(CheckTableArgs),
    /// CheckColumnar check columnar data consistency.
    CheckColumnar(CheckColumnarArgs),
    /// Mvcc gets the mvcc information of given keys.
    Mvcc(MvccArgs),
    /// Show some information.
    Show(ShowArgs),
    /// Execute recovery mode commands.
    Recovery(RecoveryArgs),
    /// Run HTTP requests to stores.
    Http(HttpArgs),
    /// Run test tools.
    Test(TestArgs),
    /// Region related commands.
    Region(RegionCommand),
    /// Materialize a TiKV config with defaults, compatibility adjustments, and
    /// full validation applied.
    MaterializeConfig(MaterializeConfigArgs),
    /// Encryption key management commands.
    Encryption(EncryptionCommand),
}

#[derive(Args)]
pub struct ShowArgs {
    #[clap(subcommand)]
    command: ShowCommands,
}

#[derive(Subcommand)]
enum ShowCommands {
    /// Show the backup meta data.
    Backup(ShowBackupArgs),
    BackupList(ShowBackupListArgs),
    Sst(ShowSstArgs),
    Columnar(ShowColumnarArgs),
    ScanBadTableFile(ScanBadTableFileArgs),
    Archive(ShowArchiveArgs),
    Fts(ShowFtsArgs),
    TxnChunk(ShowTxnChunkArgs),
    TxnLog(ShowTxnLogArgs),
    Schema(ShowSchemaArgs),
}

fn execute_show(args: ShowArgs) {
    match args.command {
        ShowCommands::Backup(args) => {
            execute_show_backup(args);
        }
        ShowCommands::BackupList(args) => {
            execute_show_backup_list(args);
        }
        ShowCommands::Sst(args) => {
            execute_show_sst(args);
        }
        ShowCommands::Columnar(args) => {
            execute_show_columnar(args);
        }
        ShowCommands::Fts(args) => {
            execute_show_fts(args);
        }
        ShowCommands::ScanBadTableFile(args) => {
            execute_scan_bad_table(args);
        }
        ShowCommands::Archive(args) => {
            execute_show_archive(args);
        }
        ShowCommands::TxnChunk(args) => {
            execute_show_txn_chunk(args);
        }
        ShowCommands::TxnLog(args) => execute_show_txn_log(args),
        ShowCommands::Schema(args) => {
            execute_show_schema(args);
        }
    }
}

#[cfg(test)]
mod tests {
    use clap::Parser;

    use super::{Cli, Commands};

    #[test]
    fn test_cli_parses_global_tls_args() {
        let cli = Cli::parse_from([
            "cse-ctl",
            "--ca-path",
            "/tmp/ca.pem",
            "--cert-path",
            "/tmp/cert.pem",
            "--key-path",
            "/tmp/key.pem",
            "encryption",
            "status",
            "--keyspace-name",
            "test_keyspace",
        ]);

        assert_eq!(cli.ca_path.as_deref(), Some("/tmp/ca.pem"));
        assert_eq!(cli.cert_path.as_deref(), Some("/tmp/cert.pem"));
        assert_eq!(cli.key_path.as_deref(), Some("/tmp/key.pem"));
        assert!(matches!(cli.command, Commands::Encryption(_)));
    }

    #[test]
    fn test_tls_override_values_require_all_paths() {
        let cli = Cli::parse_from([
            "cse-ctl",
            "--ca-path",
            "/tmp/ca.pem",
            "encryption",
            "status",
            "--keyspace-name",
            "test_keyspace",
        ]);

        assert_eq!(
            cli.tls_override_values().unwrap_err(),
            "--ca-path, --cert-path, and --key-path must be set together"
        );
    }

    #[test]
    fn test_cli_parses_global_tls_args_with_keyspace_id() {
        let cli = Cli::parse_from([
            "cse-ctl",
            "--ca-path",
            "/tmp/ca.pem",
            "--cert-path",
            "/tmp/cert.pem",
            "--key-path",
            "/tmp/key.pem",
            "encryption",
            "status",
            "--keyspace-id",
            "42",
        ]);

        assert_eq!(cli.ca_path.as_deref(), Some("/tmp/ca.pem"));
        assert_eq!(cli.cert_path.as_deref(), Some("/tmp/cert.pem"));
        assert_eq!(cli.key_path.as_deref(), Some("/tmp/key.pem"));
        assert!(matches!(cli.command, Commands::Encryption(_)));
    }
}