// Copyright 2023 PingCAP, Inc.

use std::{
    ffi::OsStr,
    fs,
    ops::Deref,
    path::{Path, PathBuf},
    str::FromStr,
    sync::Arc,
};

use clap::Args;
use futures::future::join_all;
use kvengine::{
    CF_NAMES, EXTRA_CF, EncryptionState, IoContext, LOCK_CF, NUM_CFS, WRITE_CF,
    dfs::{DFSConfig, FileType},
    table::{
        Iterator, Value,
        blobtable::{blobtable::BlobTable, builder::BlobFooter},
        file::{File, InMemFile},
        sstable::{BlockCache, Footer, L0Footer, L0Table, SsTable},
    },
};
use tidb_query_datatype::codec::table::{
    RECORD_ROW_KEY_LEN, check_record_key, decode_common_handle, decode_int_handle, decode_table_id,
};

use crate::common::get_file_data_from_dfs;

#[derive(Args)]
pub struct ShowSstArgs {
    /// The path of the config file.
    #[clap(long, default_value = "")]
    pub config: PathBuf,
    /// The id of the SST file.
    #[clap(long, default_value_t = 0)]
    pub id: u64,
    /// The path of local SST file.
    ///
    /// If specified, the SST file will be retrieved from the path instead of
    /// DFS.
    #[clap(long)]
    pub local: Option<PathBuf>,
    /// Dump entries in the specified CF. `--dump-cf` without a value defaults
    /// to `0` (WRITE_CF). Supported values: 0=write, 1=lock, 2=extra.
    #[clap(long, value_name = "CF")]
    pub dump_cf: Option<Option<DumpCf>>,
}

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

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)
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct DumpCf(usize);

impl FromStr for DumpCf {
    type Err = String;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        Ok(Self(parse_dump_cf(value)?))
    }
}

fn parse_dump_cf(value: &str) -> Result<usize, String> {
    let cf = value.parse::<usize>().map_err(|_| {
        format!(
            "invalid CF `{value}`, expected 0 ({}) / 1 ({}) / 2 ({})",
            CF_NAMES[WRITE_CF], CF_NAMES[LOCK_CF], CF_NAMES[EXTRA_CF]
        )
    })?;
    if cf >= NUM_CFS {
        return Err(format!(
            "invalid CF `{value}`, expected 0 ({}) / 1 ({}) / 2 ({})",
            CF_NAMES[WRITE_CF], CF_NAMES[LOCK_CF], CF_NAMES[EXTRA_CF]
        ));
    }
    Ok(cf)
}

fn resolve_dump_cf(dump_cf: Option<Option<DumpCf>>) -> Option<usize> {
    dump_cf.map(|cf| cf.map_or(WRITE_CF, |cf| cf.0))
}

#[derive(Debug)]
enum SstType {
    Unknown,
    L0,
    L1Plus,
    Blob,
}

fn detect_sst_type(file: &dyn File) -> SstType {
    // L1+ tables should be checked before L0 tables.
    if file.size() >= SsTable::footer_size() as u64 {
        let mut footer = Footer::default();
        footer.unmarshal(&file.read_footer(SsTable::footer_size()).unwrap());
        if footer.is_match() {
            return SstType::L1Plus;
        }
    }

    if file.size() >= L0Table::footer_size() as u64 {
        let mut footer = L0Footer::default();
        footer.unmarshal(&file.read_footer(L0Table::footer_size()).unwrap());
        if footer.is_match() {
            return SstType::L0;
        }
    }

    if file.size() >= BlobTable::footer_size() as u64 {
        let mut footer = BlobFooter::default();
        footer.unmarshal(&file.read_footer(BlobTable::footer_size()).unwrap());
        if footer.is_match() {
            return SstType::Blob;
        }
    }
    SstType::Unknown
}

pub fn execute_show_sst(args: ShowSstArgs) {
    let dump_cf = resolve_dump_cf(args.dump_cf);
    let mut config = ShowSstConfig::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();
    }
    config.dfs.override_from_env();

    let data = match args.local {
        Some(local) => get_file_data_from_local(&local),
        None => get_file_data_from_dfs(args.id, FileType::Sst, &config.dfs),
    };
    let file = Arc::new(InMemFile::new(args.id, data));
    let sst_type = detect_sst_type(file.as_ref());
    match sst_type {
        SstType::L0 => {
            let l0 = L0Table::new(file, BlockCache::None, false, EncryptionState::Disabled)
                .unwrap()
                .unwrap();
            print_l0_table(&l0);
            if let Some(dump_cf) = dump_cf {
                print_l0_cf_dump(&l0, dump_cf, 2);
            }
        }
        SstType::Blob => {
            let blob = BlobTable::new(file).unwrap();
            print_blob_table(&blob);
            if dump_cf.is_some() {
                eprintln!("--dump-cf only supports SST tables");
            }
        }
        SstType::L1Plus => {
            let ln = SsTable::new(file, BlockCache::None, EncryptionState::Disabled).unwrap();
            println!("[SST {}, level 1+]", ln.id());
            print_sstable(&ln, 2);
            if let Some(dump_cf) = dump_cf {
                print_sstable_cf_dump(&ln, dump_cf, 2);
            }
        }
        SstType::Unknown => {
            eprintln!("Not a valid SST file");
        }
    }
}

pub(crate) fn print_l0_table(l0: &L0Table) {
    println!("[SST {}, level 0]", l0.id());
    println!("  size: {}", l0.size());
    println!("  max_ts: {}", l0.max_ts());
    println!("  entries: {}", l0.entries());
    println!("  tombs: {}", l0.tombs());
    println!("  entries_write_cf: {}", l0.entries_write_cf());
    println!("  kv_size: {}", l0.kv_size());
    println!("  version: {}", l0.snap_version());
    println!(
        "  smallest: {}",
        log_wrappers::hex_encode_upper(l0.smallest().deref())
    );
    println!(
        "  biggest: {}",
        log_wrappers::hex_encode_upper(l0.biggest().deref())
    );
    println!("  total_blob_size: {}", l0.total_blob_size());

    for cf in 0..kvengine::NUM_CFS {
        println!("  [CF {}]", cf);
        if let Some(tbl) = l0.get_cf(cf) {
            print_sstable(tbl, 4);
        } else {
            println!("    None");
        }
    }
}

pub(crate) fn print_blob_table(blob: &BlobTable) {
    println!("[BLOB {}]", blob.id());
    println!("  version: {}", blob.version());
    println!("  size: {}", blob.size());
    println!(
        "  smallest: {}",
        log_wrappers::hex_encode_upper(blob.smallest_key().deref())
    );
    println!(
        "  biggest: {}",
        log_wrappers::hex_encode_upper(blob.biggest_key().deref())
    );
    println!("  total_blob_size: {}", blob.total_blob_size());
    println!("  compression_tp: {}", blob.compression_tp());
    println!("  min_blob_size: {}", blob.min_blob_size());
}

#[derive(Debug, PartialEq, Eq)]
struct WriteCfDumpEntry {
    table_id: i64,
    handle: String,
    handle_raw: String,
    version: u64,
    deleted: bool,
    tombstoned: bool,
}

#[derive(Debug, Default, PartialEq, Eq)]
struct WriteCfDumpResult {
    entries: Vec<WriteCfDumpEntry>,
    skipped_non_record_keys: usize,
}

#[derive(Debug, PartialEq, Eq)]
struct RawCfDumpEntry {
    key: String,
    version: u64,
    deleted: bool,
    value_empty: bool,
    meta: u8,
    user_meta: String,
    value: String,
}

fn print_l0_cf_dump(l0: &L0Table, cf: usize, indent: usize) {
    let indent = " ".repeat(indent);
    println!("{}[CF {} {} dump]", indent, cf, CF_NAMES[cf]);
    let Some(tbl) = l0.get_cf(cf).as_ref() else {
        println!("{}  None", indent);
        return;
    };
    print_cf_dump_body(tbl, cf, indent.len() + 2);
}

fn print_sstable_cf_dump(tbl: &SsTable, cf: usize, indent: usize) {
    let indent = " ".repeat(indent);
    println!("{}[CF {} {} dump]", indent, cf, CF_NAMES[cf]);
    print_cf_dump_body(tbl, cf, indent.len() + 2);
}

fn print_cf_dump_body(tbl: &SsTable, cf: usize, indent: usize) {
    match cf {
        WRITE_CF => print_write_cf_dump_body(tbl, indent),
        LOCK_CF | EXTRA_CF => print_raw_cf_dump_body(tbl, indent),
        _ => unreachable!("validated by clap"),
    }
}

fn print_write_cf_dump_body(tbl: &SsTable, indent: usize) {
    let result = collect_write_cf_dump(tbl);
    let indent = " ".repeat(indent);
    if result.entries.is_empty() {
        println!("{}No row-record keys found.", indent);
    } else {
        for (idx, entry) in result.entries.iter().enumerate() {
            println!(
                "{}{}: table_id={}, handle={}, handle_raw={}, version={}, deleted={}, tombstoned={}",
                indent,
                idx,
                entry.table_id,
                entry.handle,
                entry.handle_raw,
                entry.version,
                entry.deleted,
                entry.tombstoned
            );
        }
    }
    if result.skipped_non_record_keys > 0 {
        println!(
            "{}skipped_non_record_keys: {}",
            indent, result.skipped_non_record_keys
        );
    }
}

fn print_raw_cf_dump_body(tbl: &SsTable, indent: usize) {
    let entries = collect_raw_cf_dump(tbl);
    let indent = " ".repeat(indent);
    if entries.is_empty() {
        println!("{}No entries found.", indent);
        return;
    }
    for (idx, entry) in entries.iter().enumerate() {
        println!(
            "{}{}: key={}, version={}, deleted={}, value_empty={}, meta=0x{:02X}, user_meta={}, value={}",
            indent,
            idx,
            entry.key,
            entry.version,
            entry.deleted,
            entry.value_empty,
            entry.meta,
            entry.user_meta,
            entry.value
        );
    }
}

fn collect_write_cf_dump(tbl: &SsTable) -> WriteCfDumpResult {
    let mut iter = tbl.new_iterator(false, true);
    let mut result = WriteCfDumpResult::default();
    iter.rewind();
    while iter.valid() {
        collect_current_write_cf_entry(&mut result, iter.key(), iter.value());
        while iter.next_version() {
            collect_current_write_cf_entry(&mut result, iter.key(), iter.value());
        }
        iter.next();
    }
    result
}

fn collect_raw_cf_dump(tbl: &SsTable) -> Vec<RawCfDumpEntry> {
    let mut iter = tbl.new_iterator(false, true);
    let mut entries = Vec::new();
    iter.rewind();
    while iter.valid() {
        entries.push(raw_cf_dump_entry(iter.key(), iter.value()));
        while iter.next_version() {
            entries.push(raw_cf_dump_entry(iter.key(), iter.value()));
        }
        iter.next();
    }
    entries
}

fn collect_current_write_cf_entry(
    result: &mut WriteCfDumpResult,
    key: kvengine::table::InnerKey<'_>,
    value: Value,
) {
    match parse_write_cf_dump_entry(key, value) {
        Some(entry) => result.entries.push(entry),
        None => result.skipped_non_record_keys += 1,
    }
}

fn parse_write_cf_dump_entry(
    key: kvengine::table::InnerKey<'_>,
    value: Value,
) -> Option<WriteCfDumpEntry> {
    let key = key.deref();
    if check_record_key(key).is_err() {
        return None;
    }

    let table_id = decode_table_id(key).ok()?;
    let handle_buf = decode_common_handle(key).ok()?;
    let handle_raw = log_wrappers::hex_encode_upper(handle_buf);
    let handle = if key.len() == RECORD_ROW_KEY_LEN {
        match decode_int_handle(key) {
            Ok(int_handle) => int_handle.to_string(),
            Err(_) => format!("raw {}", handle_raw),
        }
    } else {
        format!("raw {}", handle_raw)
    };

    Some(WriteCfDumpEntry {
        table_id,
        handle,
        handle_raw,
        version: value.version,
        deleted: value.is_value_empty(),
        tombstoned: value.is_deleted(),
    })
}

fn raw_cf_dump_entry(key: kvengine::table::InnerKey<'_>, value: Value) -> RawCfDumpEntry {
    RawCfDumpEntry {
        key: log_wrappers::hex_encode_upper(key.deref()),
        version: value.version,
        deleted: value.is_deleted(),
        value_empty: value.is_value_empty(),
        meta: value.meta,
        user_meta: log_wrappers::hex_encode_upper(value.user_meta()),
        value: log_wrappers::hex_encode_upper(value.get_value()),
    }
}

pub(crate) fn print_sstable(tbl: &SsTable, indent: usize) {
    let indent = " ".repeat(indent);
    println!("{}size: {}", indent, tbl.size());
    println!("{}index_size: {}", indent, tbl.index_size());
    println!("{}filter_size: {}", indent, tbl.filter_size());
    println!("{}max_ts: {}", indent, tbl.max_ts);
    println!("{}entries: {}", indent, tbl.entries);
    println!("{}old_entries: {}", indent, tbl.old_entries);
    println!("{}tombs: {}", indent, tbl.tombs);
    println!("{}kv_size: {}", indent, tbl.kv_size);
    println!(
        "{}smallest: {}",
        indent,
        log_wrappers::hex_encode_upper(tbl.smallest().deref())
    );
    println!(
        "{}biggest: {}",
        indent,
        log_wrappers::hex_encode_upper(tbl.biggest().deref())
    );
    println!("{}compression_type: {}", indent, tbl.compression_type());
    println!("{}total_blob_size: {}", indent, tbl.total_blob_size());
    println!("{}encryption_ver: {}", indent, tbl.encryption_ver());

    let idx = tbl.load_index(true);
    println!("{}num_blocks: {}", indent, idx.num_blocks());
}

#[derive(Args)]
pub struct ScanBadTableFileArgs {
    /// The local path to scan SST files.
    #[clap(long)]
    pub path: PathBuf,
    #[clap(long, default_value_t = 4)]
    pub concurrency: usize,
}

pub fn execute_scan_bad_table(args: ScanBadTableFileArgs) {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .max_blocking_threads(args.concurrency + 1)
        .enable_all()
        .build()
        .unwrap();
    if args.path.is_dir() {
        let handles = iterate_dir(&args.path, &rt);
        rt.block_on(join_all(handles));
    } else {
        eprintln!("{} is not a directory", args.path.display());
    }
}

fn iterate_dir(path: &Path, rt: &tokio::runtime::Runtime) -> Vec<tokio::task::JoinHandle<()>> {
    let mut handles = vec![];
    let read_dir = fs::read_dir(path).ctx("read_dir").unwrap();
    for entry in read_dir.flatten() {
        let sub_path = entry.path();
        if sub_path.is_dir() {
            handles.extend(iterate_dir(&sub_path, rt));
        } else {
            let h = rt.spawn_blocking(move || handle_table_file(&sub_path));
            handles.push(h)
        }
    }
    handles
}

fn handle_table_file(path: &Path) {
    if path.extension().is_none_or(|ext| ext != OsStr::new("sst")) {
        // TODO: support other file types
        return;
    }
    let filename = path.file_name().unwrap().to_string_lossy();
    let data = get_file_data_from_local(path);
    let file = Arc::new(InMemFile::new(0, data));
    let sst_type = detect_sst_type(file.as_ref());
    match sst_type {
        SstType::L0 => {
            let _ = L0Table::new(file, BlockCache::None, false, EncryptionState::Disabled)
                .unwrap_or_else(|err| panic!("{}({:?}) is invalid: {:?}", filename, sst_type, err))
                .unwrap();
        }
        SstType::Blob => {
            let _ = BlobTable::new(file)
                .unwrap_or_else(|err| panic!("{}({:?}) is invalid {:?}", filename, sst_type, err));
        }
        SstType::L1Plus => {
            let _ = SsTable::new(file, BlockCache::None, EncryptionState::Disabled)
                .unwrap_or_else(|err| panic!("{}({:?}) is invalid: {:?}", filename, sst_type, err));
        }
        SstType::Unknown => {
            panic!("{} is not a SST file", filename);
        }
    }
    println!("{filename}({sst_type:?}) is OK");
}