// Copyright 2025 PingCAP, Inc.

use std::{
    path::{Path, PathBuf},
    sync::Arc,
};

use async_trait::async_trait;
use clap::Args;
use kvengine::{
    dfs::{DFSConfig, FileType},
    table::{file::InMemFile, schema_file::SchemaFile},
};
use schema::{DEFAULT_SCAN_BATCH_SIZE, KvGetter, KvScanner};
use security::SecurityConfig;
use tikv_client::{
    BoundRange, Key, KvPair, TimestampExt, TransactionClient as TxnClient, TransactionOptions,
    Value,
};
use txn_types::TimeStamp;

use crate::common::get_file_data_from_dfs;

#[derive(Args)]
pub struct ShowSchemaArgs {
    /// The path of the config file.
    #[clap(long, default_value = "")]
    pub config: PathBuf,
    /// The id of the Schema file.
    #[clap(long, default_value_t = 0)]
    pub id: u64,
    /// The path of local Schema file.
    ///
    /// If specified, the Schema file will be retrieved from the path instead of
    /// DFS.
    #[clap(long)]
    pub local: Option<PathBuf>,
    /// The table id to show.
    #[clap(long)]
    pub table_id: Option<i64>,

    /// Read schema meta directly from TiKV and print JSON.
    ///
    /// This is useful for debugging schema manager behavior (e.g. column
    /// defaults) without relying on schema files.
    #[clap(long)]
    pub from_tikv: bool,
    /// PD endpoints, use `,` to separate multiple PDs (only for --from-tikv).
    #[clap(long, default_value_t = String::new())]
    pub pd: String,
    /// Keyspace id to read schema meta from (only for --from-tikv).
    #[clap(long, default_value_t = 0)]
    pub keyspace_id: u32,
}

#[derive(Clone, Serialize, Deserialize, PartialEq, Debug, Default)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct ShowSchemaConfig {
    pub dfs: DFSConfig,
}

impl ShowSchemaConfig {
    pub fn from_args(config_path: &Path) -> Self {
        let mut config = Self::default();
        if config_path.exists() {
            let data = std::fs::read(config_path).expect("failed to read config file");
            config = toml::from_slice(&data).unwrap();
        } else if !config_path.as_os_str().is_empty() {
            panic!("config file not found: {:?}", config_path);
        }
        config.dfs.override_from_env();
        config
    }
}

#[derive(Clone, Serialize, Deserialize, PartialEq, Debug, Default)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
struct ShowSchemaMetaConfig {
    pub pd: pd_client::Config,
    pub security: SecurityConfig,
}

fn get_show_schema_meta_config_from_args(args: &ShowSchemaArgs) -> ShowSchemaMetaConfig {
    let mut config = ShowSchemaMetaConfig::default();
    if args.config.exists() {
        let data = std::fs::read(args.config.clone()).expect("failed to read config file");
        config = toml::from_slice(&data).unwrap();
    } else if !args.config.as_os_str().is_empty() {
        panic!("config file not found: {:?}", args.config);
    }
    config.security.override_from_env();
    if !args.pd.is_empty() {
        config.pd.endpoints = args.pd.split(',').map(|x| x.to_owned()).collect();
    }
    config
}

pub(crate) fn get_file_data_from_local(local: &Path) -> bytes::Bytes {
    let data = std::fs::read(local).unwrap_or_else(|err| {
        panic!("failed to read local file from {:?}: {:?}", local, err);
    });
    bytes::Bytes::from(data)
}

fn parse_file_id(path: &Path) -> u64 {
    let file_prefix = path.file_prefix().unwrap().to_str().unwrap();
    u64::from_str_radix(file_prefix, 16).unwrap()
}

pub(crate) fn execute_show_schema(args: ShowSchemaArgs) {
    if args.from_tikv {
        execute_show_schema_meta(args);
        return;
    }
    execute_show_schema_file(args);
}

fn execute_show_schema_file(args: ShowSchemaArgs) {
    let config = ShowSchemaConfig::from_args(&args.config);
    let (file_id, schema_file_data) = if let Some(local) = args.local {
        (parse_file_id(&local), get_file_data_from_local(&local))
    } else {
        (
            args.id,
            get_file_data_from_dfs(args.id, FileType::Schema, &config.dfs),
        )
    };
    let file = InMemFile::new(file_id, schema_file_data);
    let schema_file = SchemaFile::open(Arc::new(file)).unwrap();
    println!("[Schema file]");
    println!("  id: {}", schema_file.get_file_id());
    println!("  keyspace_id: {}", schema_file.get_keyspace_id());
    println!("  version: {}", schema_file.get_version());
    println!("  restore_version: {}", schema_file.get_restore_version());
    let tables = schema_file.export_schemas();
    if let Some(table_id) = args.table_id {
        println!("  table: {}", table_id);
        println!("    schema: {:?}", tables.get(&table_id).unwrap());
        return;
    }
    for (table_id, schema) in tables {
        println!("    table: {}", table_id);
        println!("    schema: {:?}", schema);
    }
}

const DEFAULT_GRPC_MAX_DECODING_MESSAGE_SIZE: usize = 128 * 1024 * 1024; // 128MB

struct SchemaMetaKvClient {
    txn_client: TxnClient,
    scan_batch_size: u32,
}

#[async_trait]
impl KvScanner for SchemaMetaKvClient {
    async fn scan(
        &self,
        start: &[u8],
        end: &[u8],
        start_ts: TimeStamp,
    ) -> std::result::Result<Vec<(Vec<u8>, Vec<u8>)>, String> {
        let txn_start_ts = tikv_client::Timestamp::from_version(start_ts.into_inner());
        let mut snapshot = self
            .txn_client
            .snapshot(txn_start_ts, TransactionOptions::new_pessimistic());

        let scan_range: BoundRange = (start.to_vec()..end.to_vec()).into();
        let batch: Vec<KvPair> = snapshot
            .scan(scan_range, self.scan_batch_size)
            .await
            .map_err(|e| e.to_string())?
            .collect();
        Ok(batch
            .into_iter()
            .map(|KvPair(key, val)| (key.into(), val))
            .collect())
    }
}

#[async_trait]
impl KvGetter for SchemaMetaKvClient {
    async fn get(
        &self,
        key: &[u8],
        start_ts: TimeStamp,
    ) -> std::result::Result<Option<Vec<u8>>, String> {
        let txn_start_ts = tikv_client::Timestamp::from_version(start_ts.into_inner());
        let mut snapshot = self
            .txn_client
            .snapshot(txn_start_ts, TransactionOptions::new_pessimistic());
        snapshot.get(key.to_vec()).await.map_err(|e| {
            format!(
                "cse-ctl show schema: kv get failed: {}: {e:?})",
                log_wrappers::Value::key(key)
            )
        })
    }

    async fn batch_get(
        &self,
        keys: &[Vec<u8>],
        start_ts: TimeStamp,
    ) -> std::result::Result<Vec<Option<Vec<u8>>>, String> {
        let txn_start_ts = tikv_client::Timestamp::from_version(start_ts.into_inner());
        let mut snapshot = self
            .txn_client
            .snapshot(txn_start_ts, TransactionOptions::new_pessimistic());
        let pairs: std::collections::HashMap<Key, Value> = snapshot
            .batch_get(keys.to_vec())
            .await
            .map_err(|e| format!("cse-ctl show schema: kv batch get failed: {e:?})"))?
            .map(|pair| (pair.0, pair.1))
            .collect();

        let mut vals = Vec::with_capacity(keys.len());
        for key in keys {
            vals.push(pairs.get(&Key::from(key.to_vec())).cloned());
        }
        Ok(vals)
    }
}

#[derive(Serialize)]
struct SchemaMetaDump {
    keyspace_id: u32,
    start_ts: u64,
    meta_schema_version: i64,
    synced_schema_version: i64,
    tables: Vec<schema::schema::TableInfo>,
}

fn execute_show_schema_meta(args: ShowSchemaArgs) {
    if args.local.is_some() || args.id != 0 {
        panic!("--from-tikv does not support --local/--id");
    }

    let config = get_show_schema_meta_config_from_args(&args);
    if config.pd.endpoints.is_empty() {
        panic!("PD endpoints not set, please specify --pd or config.pd.endpoints");
    }

    let mut tikv_client_config = tikv_client::Config::default()
        .with_grpc_max_decoding_message_size(DEFAULT_GRPC_MAX_DECODING_MESSAGE_SIZE);
    if !config.security.ca_path.is_empty()
        || !config.security.cert_path.is_empty()
        || !config.security.key_path.is_empty()
    {
        tikv_client_config = tikv_client_config.with_security(
            config.security.ca_path.clone(),
            config.security.cert_path.clone(),
            config.security.key_path.clone(),
        );
    }

    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();

    runtime
        .block_on(async move {
            let txn_client = TxnClient::new_with_config_api_v2_no_prefix(
                config.pd.endpoints,
                tikv_client_config,
            )
            .await
            .map_err(|e| format!("create txn client failed: {e:?}"))?;
            let now = txn_client
                .current_timestamp()
                .await
                .map_err(|e| format!("get timestamp failed: {e:?}"))?;
            let start_ts = TimeStamp::new(now.version());

            let kv = Arc::new(SchemaMetaKvClient {
                txn_client,
                scan_batch_size: DEFAULT_SCAN_BATCH_SIZE,
            });
            let kv_getter: Arc<dyn KvGetter> = kv.clone();
            let kv_scanner: Arc<dyn KvScanner> = kv.clone();

            let meta_schema_version =
                schema::get_schema_version(kv_getter.clone(), args.keyspace_id, start_ts).await?;
            let sync_result =
                schema::sync_schema(kv_getter, kv_scanner, args.keyspace_id, None, start_ts)
                    .await?;
            let synced_schema_version = sync_result.schema_version;
            let mut tables = sync_result.table_infos;
            if let Some(table_id) = args.table_id {
                tables.retain(|t| t.id == table_id);
            }
            tables.sort_by_key(|t| t.id);

            let dump = SchemaMetaDump {
                keyspace_id: args.keyspace_id,
                start_ts: start_ts.into_inner(),
                meta_schema_version,
                synced_schema_version,
                tables,
            };
            println!("{}", serde_json::to_string_pretty(&dump).unwrap());
            Ok::<(), String>(())
        })
        .unwrap_or_else(|e| panic!("{e}"));
}