// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0.

mod config;
mod metrics;
mod pool;
mod report;
mod slow_task;
mod workload;

use std::{path::PathBuf, thread};

use anyhow::Result;
use clap::Parser;

use crate::{
    config::{ConfigOverride, SimulatorConfig},
    metrics::KeyspaceMetricStore,
    pool::SimulatorReadPool,
};

#[derive(Debug, Parser)]
#[clap(name = "yatp-sim")]
#[clap(about = "Generate controlled workloads on a real YATP read pool")]
struct Args {
    /// Path to the workload TOML file. Relative paths fall back to the
    /// yatp-sim crate root when they do not exist in the working directory.
    #[clap(long, value_name = "PATH")]
    config: PathBuf,

    /// Override a top-level or read-pool config field after loading the TOML
    /// file.
    #[clap(long, value_name = "FIELD=VALUE")]
    set: Vec<String>,

    /// Override a keyspace group selected by its configured name.
    #[clap(
        long,
        value_names = &["GROUP", "FIELD=VALUE"],
        number_of_values = 2,
        allow_hyphen_values = true
    )]
    set_keyspace: Vec<String>,

    /// Override a task selected by its keyspace group and configured task name.
    #[clap(
        long,
        value_names = &["GROUP", "TASK", "FIELD=VALUE"],
        number_of_values = 3,
        allow_hyphen_values = true
    )]
    set_task: Vec<String>,
}

impl Args {
    fn config_overrides(&self) -> Vec<ConfigOverride> {
        debug_assert_eq!(self.set_keyspace.len() % 2, 0);
        debug_assert_eq!(self.set_task.len() % 3, 0);
        let mut overrides = Vec::with_capacity(
            self.set.len() + self.set_keyspace.len() / 2 + self.set_task.len() / 3,
        );
        overrides.extend(self.set.iter().cloned().map(ConfigOverride::Global));
        overrides.extend(self.set_keyspace.chunks_exact(2).map(|values| {
            ConfigOverride::Keyspace {
                group: values[0].clone(),
                assignment: values[1].clone(),
            }
        }));
        overrides.extend(
            self.set_task
                .chunks_exact(3)
                .map(|values| ConfigOverride::Task {
                    group: values[0].clone(),
                    task: values[1].clone(),
                    assignment: values[2].clone(),
                }),
        );
        overrides
    }
}

fn main() -> Result<()> {
    let args = Args::parse();
    let overrides = args.config_overrides();
    let config = SimulatorConfig::load(&args.config, &overrides)?;
    let metric_store = KeyspaceMetricStore::new()?;
    let read_pool = SimulatorReadPool::new(&config.read_pool, metric_store.keyspace_metrics())?;
    let handle = read_pool.handle();
    let worker_cpu_time = read_pool.worker_cpu_time_source();

    println!(
        "starting workload: duration={:?}, workers={}, fairness={}, global_cpu_admission={}",
        config.duration.0,
        config.read_pool.workers,
        config.read_pool.keyspace_fairness.enabled,
        config.read_pool.fair_global_cpu_admission.enabled,
    );
    let result = workload::run(&config, &handle, metric_store, worker_cpu_time);
    // workload::run finalizes all report snapshots before returning. Keep the
    // pool alive during cooldown without extending any reported interval.
    if result.is_ok() && !config.cooldown_duration.0.is_zero() {
        println!(
            "workload metrics finished; waiting {:?} before shutdown",
            config.cooldown_duration.0
        );
        thread::sleep(config.cooldown_duration.0);
    }
    drop(handle);
    read_pool.shutdown();
    let report_data = result?;
    report::write_reports(&report_data)?;
    metrics::print_run_summary(&report_data);
    println!("workload finished");
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_cli_parses_repeated_config_overrides() {
        let args = Args::try_parse_from([
            "yatp-sim",
            "--config",
            "workload.toml",
            "--set",
            "duration=30s",
            "--set",
            "read-pool.workers=8",
            "--set-keyspace",
            "noisy tenant",
            "count=3",
            "--set-keyspace",
            "victim",
            "count=2",
            "--set-task",
            "noisy tenant",
            "normal/read",
            "concurrency=64",
            "--set-task",
            "victim",
            "normal",
            "qps=300",
        ])
        .unwrap();

        assert_eq!(
            args.config_overrides(),
            vec![
                ConfigOverride::Global("duration=30s".to_owned()),
                ConfigOverride::Global("read-pool.workers=8".to_owned()),
                ConfigOverride::Keyspace {
                    group: "noisy tenant".to_owned(),
                    assignment: "count=3".to_owned(),
                },
                ConfigOverride::Keyspace {
                    group: "victim".to_owned(),
                    assignment: "count=2".to_owned(),
                },
                ConfigOverride::Task {
                    group: "noisy tenant".to_owned(),
                    task: "normal/read".to_owned(),
                    assignment: "concurrency=64".to_owned(),
                },
                ConfigOverride::Task {
                    group: "victim".to_owned(),
                    task: "normal".to_owned(),
                    assignment: "qps=300".to_owned(),
                },
            ]
        );
    }

    #[test]
    fn test_cli_rejects_combined_keyspace_selector_and_assignment() {
        let result = Args::try_parse_from([
            "yatp-sim",
            "--config",
            "workload.toml",
            "--set-keyspace",
            "noisy count=3",
        ]);

        result.unwrap_err();
    }

    #[test]
    fn test_cli_rejects_incomplete_task_override() {
        let result = Args::try_parse_from([
            "yatp-sim",
            "--config",
            "workload.toml",
            "--set-task",
            "noisy",
            "concurrency=64",
        ]);

        result.unwrap_err();
    }

    #[test]
    fn test_cli_accepts_selectors_starting_with_hyphens() {
        let args = Args::try_parse_from([
            "yatp-sim",
            "--config",
            "workload.toml",
            "--set-keyspace",
            "-noisy",
            "count=3",
            "--set-task",
            "-noisy",
            "-normal",
            "concurrency=64",
        ])
        .unwrap();

        assert_eq!(
            args.config_overrides(),
            vec![
                ConfigOverride::Keyspace {
                    group: "-noisy".to_owned(),
                    assignment: "count=3".to_owned(),
                },
                ConfigOverride::Task {
                    group: "-noisy".to_owned(),
                    task: "-normal".to_owned(),
                    assignment: "concurrency=64".to_owned(),
                },
            ]
        );
    }

    #[test]
    fn test_cli_overrides_are_applied_to_loaded_config() {
        let args = Args::try_parse_from([
            "yatp-sim",
            "--config",
            "configs/five-tenants.toml",
            "--set",
            "duration=30s",
            "--set",
            "warmup-duration=0s",
            "--set-keyspace",
            "B",
            "count=2",
            "--set-task",
            "B",
            "burst",
            "start-time=0s",
        ])
        .unwrap();

        let config = SimulatorConfig::load(&args.config, &args.config_overrides()).unwrap();
        assert_eq!(config.duration.0, std::time::Duration::from_secs(30));
        assert_eq!(config.warmup_duration.0, std::time::Duration::ZERO);
        assert_eq!(config.keyspace_groups[1].count, 2);
        assert_eq!(
            config.keyspace_groups[1].tasks[1].start_time.0,
            std::time::Duration::ZERO
        );
    }
}