// Copyright 2026 PingCAP, Inc.

use std::{
    fs,
    io::{Error, ErrorKind, Result},
    path::PathBuf,
    sync::Mutex,
};

pub const PSI_PATH: &str = "/proc/pressure";

fn cgroup_v2_pressure_path(resource: &str) -> Option<PathBuf> {
    // cgroup v2 exposes per-cgroup PSI as <mount>/<cgroup>/<resource>.pressure
    // We assume the unified mount is /sys/fs/cgroup.
    let cgroup_root = PathBuf::from("/sys/fs/cgroup");
    if !cgroup_root.exists() {
        return None;
    }

    // Parse /proc/self/cgroup to get the relative cgroup path (field after the
    // second colon).
    let content = fs::read_to_string("/proc/self/cgroup").ok()?;
    let mut rel_path = None;
    for line in content.lines() {
        // cgroup v2 line format: 0::<path>
        let mut parts = line.splitn(3, ':');
        let _ = parts.next();
        let _controllers = parts.next();
        if let Some(path) = parts.next() {
            rel_path = Some(path);
            break;
        }
    }

    let rel = rel_path?;
    let mut candidate = cgroup_root;
    candidate.push(rel.trim_start_matches('/'));
    candidate.push(format!("{}.pressure", resource));
    if candidate.exists() {
        Some(candidate)
    } else {
        None
    }
}

fn has_cgroup_v2_pressure() -> bool {
    // Quick check: controllers file exists and cpu.pressure exists at least at
    // root.
    let controllers = PathBuf::from("/sys/fs/cgroup/cgroup.controllers");
    if !controllers.exists() {
        return false;
    }
    let root_cpu = PathBuf::from("/sys/fs/cgroup/cpu.pressure");
    root_cpu.exists()
}

use prometheus::{
    self, GaugeVec, Opts,
    core::{Collector, Desc},
    proto,
};

pub fn monitor_system_psi<S: Into<String>>(namespace: S) -> Result<()> {
    // Require cgroup v2 PSI; if unavailable, register nothing.
    if !has_cgroup_v2_pressure() {
        warn!("cgroup v2 PSI not available; pressure_stall metrics disabled");
        return Ok(());
    }
    let collector = PsiCollector::new(namespace.into());
    prometheus::register(Box::new(collector)).map_err(|e| Error::other(format!("{:?}", e)))
}

struct PsiMetrics {
    totals: GaugeVec,
    averages: GaugeVec,
}

impl PsiMetrics {
    fn new(namespace: String) -> PsiMetrics {
        let totals = GaugeVec::new(
            Opts::new(
                "pressure_stall_seconds_total",
                "Cumulative stall time from PSI files in seconds.",
            )
            .namespace(namespace.clone()),
            &["resource", "level"],
        )
        .unwrap();

        let averages = GaugeVec::new(
            Opts::new(
                "pressure_stall_average",
                "PSI average stall percentage reported over different windows.",
            )
            .namespace(namespace),
            &["resource", "level", "window"],
        )
        .unwrap();

        PsiMetrics { totals, averages }
    }

    fn descs(&self) -> Vec<Desc> {
        let mut descs = Vec::with_capacity(2);
        descs.extend(self.totals.desc().into_iter().cloned());
        descs.extend(self.averages.desc().into_iter().cloned());
        descs
    }

    fn reset(&mut self) {
        self.totals.reset();
        self.averages.reset();
    }
}

#[derive(Default)]
struct PsiEntry {
    resource: String,
    level: String,
    avg10: f64,
    avg60: f64,
    avg300: f64,
    total_seconds: f64,
}

fn read_psi(resource: &'static str) -> Result<Vec<PsiEntry>> {
    // Prefer cgroup v2 scoped PSI if available; fall back to host PSI.
    let path = cgroup_v2_pressure_path(resource)
        .unwrap_or_else(|| PathBuf::from(format!("/proc/pressure/{}", resource)));

    let content = match fs::read_to_string(path) {
        Ok(c) => c,
        Err(e) if e.kind() == ErrorKind::NotFound => return Ok(vec![]),
        Err(e) => return Err(e),
    };

    let mut entries = Vec::new();
    for line in content.lines() {
        let mut parts = line.split_whitespace();
        let level = match parts.next() {
            Some(lvl) => lvl,
            None => continue,
        };

        let mut entry = PsiEntry {
            resource: resource.to_string(),
            level: level.to_string(),
            ..Default::default()
        };

        for part in parts {
            if let Some((k, v)) = part.split_once('=') {
                if let Ok(val) = v.parse::<f64>() {
                    match k {
                        "avg10" => entry.avg10 = val,
                        "avg60" => entry.avg60 = val,
                        "avg300" => entry.avg300 = val,
                        "total" => entry.total_seconds = val / 1_000_000.0, // microseconds ->
                        // seconds
                        _ => {}
                    }
                }
            }
        }
        entries.push(entry);
    }

    Ok(entries)
}

struct PsiCollector {
    descs: Vec<Desc>,
    metrics: Mutex<PsiMetrics>,
}

impl PsiCollector {
    fn new(namespace: String) -> PsiCollector {
        let metrics = PsiMetrics::new(namespace);
        let descs = metrics.descs();
        PsiCollector {
            descs,
            metrics: Mutex::new(metrics),
        }
    }
}

impl Collector for PsiCollector {
    fn desc(&self) -> Vec<&Desc> {
        self.descs.iter().collect()
    }

    fn collect(&self) -> Vec<proto::MetricFamily> {
        let mut metrics = self.metrics.lock().unwrap();
        metrics.reset();

        for resource in ["cpu", "io", "memory"] {
            if let Ok(entries) = read_psi(resource) {
                for entry in entries {
                    metrics
                        .totals
                        .get_metric_with_label_values(&[
                            entry.resource.as_str(),
                            entry.level.as_str(),
                        ])
                        .unwrap()
                        .set(entry.total_seconds);

                    for (window, value) in [
                        ("avg10", entry.avg10),
                        ("avg60", entry.avg60),
                        ("avg300", entry.avg300),
                    ] {
                        metrics
                            .averages
                            .get_metric_with_label_values(&[
                                entry.resource.as_str(),
                                entry.level.as_str(),
                                window,
                            ])
                            .unwrap()
                            .set(value);
                    }
                }
            }
        }

        let mut mfs = metrics.totals.collect();
        mfs.extend(metrics.averages.collect());
        mfs
    }
}