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

use std::{
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, AtomicU64, AtomicUsize},
    },
    thread,
    time::{Duration, Instant},
};

use anyhow::{Context, Result, ensure};
use prometheus::IntGauge;
use resource_control::{CpuThrottleConfig, CpuThrottleManager, ReadLimiter};
use tikv::read_pool::{FairGlobalCpuAdmissionController, ReadPool, ReadPoolHandle};
use tikv_util::{
    task_reschedule::DEFAULT_MAX_TIME_SLICE_US,
    yatp_pool::{DefaultTicker, FairQueueAdmissionController, KeyspaceMetrics, YatpPoolBuilder},
};

use crate::{
    config::{FairGlobalCpuAdmissionConfig, ReadPoolConfig},
    metrics::WorkerCpuTimeSource,
};

pub struct SimulatorReadPool {
    read_pool: ReadPool,
    worker_cpu_time: Arc<WorkerCpuTimeSampler>,
}

impl SimulatorReadPool {
    pub fn new(config: &ReadPoolConfig, keyspace_metrics: Arc<KeyspaceMetrics>) -> Result<Self> {
        let worker_cpu_time = Arc::new(WorkerCpuTimeSampler::new());
        let register_worker_cpu_time = worker_cpu_time.clone();
        let mut builder = YatpPoolBuilder::new(DefaultTicker);
        builder
            .name_prefix("yatp-sim")
            .thread_count(config.workers, config.workers, config.workers)
            .keyspace_metrics(Some(keyspace_metrics.clone()))
            .after_start(move || register_worker_cpu_time.register_current_thread());

        let fair_global_cpu_admission_controller =
            build_fair_global_cpu_admission_controller(&config.fair_global_cpu_admission)?;
        let (pool, keyspace_fair_queue) = if config.keyspace_fairness.enabled {
            let controller = fair_global_cpu_admission_controller
                .as_ref()
                .map(|controller| controller.clone() as Arc<dyn FairQueueAdmissionController>);
            let (pool, queue) = builder.build_keyspace_fair_pool(controller);
            (pool, Some(queue))
        } else {
            ensure!(
                fair_global_cpu_admission_controller.is_none(),
                "fair global CPU admission requires keyspace fairness"
            );
            (builder.build_multi_level_pool(), None)
        };

        // ReadPool requires these gauges for internal bookkeeping. They are
        // deliberately unregistered because they are not part of the
        // simulator report.
        let running_tasks = IntGauge::new(
            "yatp_sim_internal_running_tasks",
            "Internal YATP simulator running task count",
        )?;
        let running_threads = IntGauge::new(
            "yatp_sim_internal_running_threads",
            "Internal YATP simulator running thread count",
        )?;
        let read_pool = ReadPool::Yatp {
            pool,
            keyspace_fair_queue,
            fairness_use_wall_time: Arc::new(AtomicBool::new(
                config.keyspace_fairness.use_wall_time,
            )),
            fairness_slice_precharge_us: Arc::new(AtomicU64::new(
                config.keyspace_fairness.slice_precharge_us,
            )),
            // The simulated workload never reaches a wall-time checkpoint, so
            // this slice is inert here until the workload calls
            // `tikv_util::cooperative_checkpoint`.
            fairness_max_time_slice_us: Arc::new(AtomicU64::new(DEFAULT_MAX_TIME_SLICE_US)),
            keyspace_metrics: Some(keyspace_metrics),
            running_tasks,
            running_threads,
            max_tasks: Arc::new(AtomicUsize::new(usize::MAX)),
            keyspace_task_limiter: None,
            pool_size: Arc::new(AtomicUsize::new(config.workers)),
            fair_global_cpu_admission_controller,
        };
        worker_cpu_time.wait_for_threads(config.workers)?;
        Ok(Self {
            read_pool,
            worker_cpu_time,
        })
    }

    pub fn handle(&self) -> ReadPoolHandle {
        self.read_pool.handle()
    }

    pub(crate) fn worker_cpu_time_source(&self) -> Arc<dyn WorkerCpuTimeSource> {
        self.worker_cpu_time.clone()
    }

    pub fn shutdown(self) {
        self.read_pool.shutdown();
    }
}

struct WorkerCpuTimeSampler {
    thread_handles: Mutex<Vec<usize>>,
}

impl WorkerCpuTimeSampler {
    fn new() -> Self {
        Self {
            thread_handles: Mutex::new(Vec::new()),
        }
    }

    fn register_current_thread(&self) {
        if let Some(handle) = current_native_thread_handle() {
            self.thread_handles
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner())
                .push(handle);
        }
    }

    fn wait_for_threads(&self, expected_threads: usize) -> Result<()> {
        ensure!(
            native_thread_cpu_supported(),
            "yatp-sim worker thread CPU sampling is unsupported on this platform"
        );
        let deadline = Instant::now() + Duration::from_secs(5);
        loop {
            let registered_threads = self
                .thread_handles
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner())
                .len();
            if registered_threads >= expected_threads {
                return Ok(());
            }
            ensure!(
                Instant::now() < deadline,
                "timed out waiting for {expected_threads} YATP workers to register for CPU sampling; registered {registered_threads}"
            );
            thread::sleep(Duration::from_millis(1));
        }
    }
}

impl WorkerCpuTimeSource for WorkerCpuTimeSampler {
    fn snapshot_cpu_nanos(&self) -> Result<u64> {
        let thread_handles = self
            .thread_handles
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
            .clone();
        ensure!(
            !thread_handles.is_empty(),
            "no YATP worker threads are registered for CPU sampling"
        );
        thread_handles.into_iter().try_fold(0u64, |total, handle| {
            let cpu_nanos = read_native_thread_cpu_nanos(handle)
                .context("failed to read CPU time from a YATP worker thread")?;
            total
                .checked_add(cpu_nanos)
                .context("YATP worker thread CPU time overflow")
        })
    }
}

#[cfg(any(target_os = "linux", target_os = "macos"))]
fn native_thread_cpu_supported() -> bool {
    true
}

#[cfg(not(any(target_os = "linux", target_os = "macos")))]
fn native_thread_cpu_supported() -> bool {
    false
}

#[cfg(any(target_os = "linux", target_os = "macos"))]
fn current_native_thread_handle() -> Option<usize> {
    Some(unsafe { libc::pthread_self() as usize })
}

#[cfg(not(any(target_os = "linux", target_os = "macos")))]
fn current_native_thread_handle() -> Option<usize> {
    None
}

#[cfg(target_os = "linux")]
fn read_native_thread_cpu_nanos(thread_handle: usize) -> Option<u64> {
    let mut clock_id = 0;
    let result =
        unsafe { libc::pthread_getcpuclockid(thread_handle as libc::pthread_t, &mut clock_id) };
    if result != 0 {
        return None;
    }
    let mut time = libc::timespec {
        tv_sec: 0,
        tv_nsec: 0,
    };
    if unsafe { libc::clock_gettime(clock_id, &mut time) } != 0 {
        return None;
    }
    let seconds = u64::try_from(time.tv_sec).ok()?;
    let nanoseconds = u64::try_from(time.tv_nsec).ok()?;
    seconds.checked_mul(1_000_000_000)?.checked_add(nanoseconds)
}

#[cfg(target_os = "macos")]
fn read_native_thread_cpu_nanos(thread_handle: usize) -> Option<u64> {
    let mach_thread = unsafe { libc::pthread_mach_thread_np(thread_handle as libc::pthread_t) };
    if mach_thread == 0 {
        return None;
    }
    let mut info: libc::thread_basic_info = unsafe { std::mem::zeroed() };
    let mut count = libc::THREAD_BASIC_INFO_COUNT;
    let result = unsafe {
        libc::thread_info(
            mach_thread,
            libc::THREAD_BASIC_INFO as libc::thread_flavor_t,
            &mut info as *mut libc::thread_basic_info as libc::thread_info_t,
            &mut count,
        )
    };
    if result != libc::KERN_SUCCESS {
        return None;
    }
    time_value_to_nanos(info.user_time)?.checked_add(time_value_to_nanos(info.system_time)?)
}

#[cfg(target_os = "macos")]
fn time_value_to_nanos(time: libc::time_value_t) -> Option<u64> {
    let seconds = u64::try_from(time.seconds).ok()?;
    let microseconds = u64::try_from(time.microseconds).ok()?;
    seconds
        .checked_mul(1_000_000_000)?
        .checked_add(microseconds.checked_mul(1_000)?)
}

#[cfg(not(any(target_os = "linux", target_os = "macos")))]
fn read_native_thread_cpu_nanos(_: usize) -> Option<u64> {
    None
}

fn build_fair_global_cpu_admission_controller(
    config: &FairGlobalCpuAdmissionConfig,
) -> Result<Option<Arc<FairGlobalCpuAdmissionController>>> {
    if !config.enabled {
        return Ok(None);
    }

    let refill_interval_ms = u64::try_from(config.refill_interval.0.as_millis())
        .context("fair global CPU refill interval is too large")?;
    let mut cpu_config = CpuThrottleConfig::default();
    cpu_config.enabled = true;
    cpu_config.refill_interval_ms = refill_interval_ms;
    let manager = Arc::new(CpuThrottleManager::new(
        cpu_config,
        config.global_cpu_quota_us()?,
        std::collections::HashMap::new(),
    ));

    let mut resource_config = resource_control::Config::default();
    resource_config.enabled = true;
    resource_config.read_enabled = true;
    resource_config.cpu_throttle.enabled = true;
    resource_config.cpu_throttle.refill_interval = config.refill_interval;
    let read_limiter = ReadLimiter::new(resource_config);
    read_limiter.set_cpu_throttle_manager(manager.clone());

    Ok(Some(Arc::new(FairGlobalCpuAdmissionController::new(
        manager,
        read_limiter,
    ))))
}

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

    #[test]
    fn test_pool_installs_fair_global_cpu_admission_controller() {
        let mut config = ReadPoolConfig::default();
        config.workers = 1;
        config.keyspace_fairness.enabled = true;
        config.fair_global_cpu_admission.enabled = true;
        config.fair_global_cpu_admission.max_cpu_cores = 0.5;
        config.fair_global_cpu_admission.refill_interval =
            tikv_util::config::ReadableDuration::millis(10);

        let metric_store = KeyspaceMetricStore::new().unwrap();
        let pool = SimulatorReadPool::new(&config, metric_store.keyspace_metrics()).unwrap();
        assert!(pool.handle().uses_fair_global_cpu_gate());
        pool.shutdown();
    }

    #[test]
    fn test_pool_skips_fair_global_cpu_admission_controller_when_disabled() {
        let mut config = ReadPoolConfig::default();
        config.workers = 1;
        config.keyspace_fairness.enabled = true;

        let metric_store = KeyspaceMetricStore::new().unwrap();
        let pool = SimulatorReadPool::new(&config, metric_store.keyspace_metrics()).unwrap();
        assert!(!pool.handle().uses_fair_global_cpu_gate());
        pool.shutdown();
    }
}