use std::env;
use std::fs;
use std::io::{Read, Write};
use std::time::{Duration, Instant};
use log::{info, warn, LevelFilter};

const IA32_PERF_CTL: u32 = 0x199;
const POLL_MS: u64 = 100;
const SAMPLE_WINDOW: usize = 10;
const STATE_WRITE_INTERVAL_S: u64 = 1;
const MSR_ERROR_SUPPRESS_COUNT: u32 = 1;
const THERMAL_CACHE_MS: u64 = 1000;

#[derive(Clone, Copy, PartialEq, Debug)]
enum Governor { Performance, Powersave, Ondemand, Conservative, Schedutil }

struct StderrLogger;
impl log::Log for StderrLogger {
    fn enabled(&self, m: &log::Metadata) -> bool { m.level() <= LevelFilter::Info }
    fn log(&self, r: &log::Record) { eprintln!("[{}] cpufreqd: {}", r.level(), r.args()); }
    fn flush(&self) {}
}

#[derive(Clone)]
struct PState { freq_khz: u32, power_mw: u32, latency_us: u32, ctl: u64 }

struct CpuInfo {
    id: u32, pstates: Vec<PState>, current_idx: usize,
    load_history: [f64; SAMPLE_WINDOW], load_idx: usize, throttle: bool,
    msr_errors: u32, msr_suppressed: bool,
}

fn detect_cpus() -> Vec<u32> {
    let mut v = Vec::new();
    if let Ok(e) = fs::read_dir("/dev/cpu") {
        for x in e.flatten() {
            if let Ok(n) = x.file_name().into_string() {
                if let Ok(id) = n.parse() { v.push(id); }
            }
        }
    }
    if v.is_empty() { v.push(0); }
    v
}

fn read_acpi_pss(cpu: u32) -> Vec<PState> {
    let path = format!("/scheme/acpi/processor/CPU{}/pss", cpu);
    if let Ok(d) = fs::read_to_string(&path) {
        let mut s = Vec::new();
        for l in d.lines() {
            let w: Vec<&str> = l.split_whitespace().collect();
            if w.len() >= 6 {
                if let (Ok(f), Ok(pw), Ok(la), Ok(ct)) =
                    (w[0].parse(), w[2].parse(), w[4].parse(), u64::from_str_radix(w[5], 16))
                { s.push(PState { freq_khz: f, power_mw: pw, latency_us: la, ctl: ct }); }
            }
        }
        if !s.is_empty() { return s; }
    }
    vec![
        PState { freq_khz: 2400, power_mw: 35000, latency_us: 10, ctl: 0x1a00 },
        PState { freq_khz: 2000, power_mw: 25000, latency_us: 10, ctl: 0x1600 },
        PState { freq_khz: 1600, power_mw: 18000, latency_us: 10, ctl: 0x1200 },
        PState { freq_khz: 1200, power_mw: 12000, latency_us: 10, ctl: 0x0e00 },
    ]
}

fn write_msr(cpu: u32, msr: u32, val: u64) -> bool {
    fs::OpenOptions::new().write(true).open(format!("/dev/cpu/{}/msr", cpu)).ok()
        .map(|mut f| f.write_all(&val.to_ne_bytes()).is_ok()).unwrap_or(false)
}

fn measure_load(cpu: u32, prev: &mut (u64, u64)) -> f64 {
    if let Ok(d) = fs::read_to_string(format!("/scheme/sys/cpu/{}/stat", cpu)) {
        let p: Vec<u64> = d.split_whitespace().filter_map(|s| s.parse().ok()).collect();
        if p.len() >= 4 {
            let t: u64 = p.iter().sum(); let i = p.get(3).copied().unwrap_or(0);
            let (pt, pi) = *prev; *prev = (t, i);
            if t > pt { let td = t - pt; let id = i.saturating_sub(pi); return 1.0 - (id as f64 / td as f64); }
        }
    }
    0.0
}

fn avg_load(ci: &CpuInfo) -> f64 { ci.load_history.iter().sum::<f64>() / SAMPLE_WINDOW as f64 }

fn choose_pstate(g: Governor, ci: &CpuInfo) -> usize {
    if ci.throttle { return (ci.pstates.len().saturating_sub(1)).min(ci.pstates.len()); }
    let l = avg_load(ci); let m = ci.pstates.len().saturating_sub(1); let c = ci.current_idx.min(m);
    match g {
        Governor::Performance => 0,
        Governor::Powersave => m,
        Governor::Ondemand => { if l > 0.8 && c > 0 { c - 1 } else if l < 0.3 && c < m { c + 1 } else { c } }
        Governor::Conservative => { if l > 0.9 && c > 0 { c - 1 } else if l < 0.2 && c < m { c + 1 } else { c } }
        Governor::Schedutil => { let t = ((l * (m + 1) as f64) as usize).min(m); if t < c && c < m { c + 1 } else if t > c && c > 0 { c - 1 } else { c } }
    }
}

struct ThermalCache { data: bool, last_check: Instant }
impl ThermalCache {
    fn new() -> Self { Self { data: false, last_check: Instant::now() - Duration::from_secs(10) } }
    fn get(&mut self) -> bool {
        if self.last_check.elapsed() < Duration::from_millis(THERMAL_CACHE_MS) { return self.data; }
        self.data = check_thermal_raw(); self.last_check = Instant::now(); self.data
    }
}

fn check_thermal_raw() -> bool {
    if let Ok(d) = fs::read_to_string("/scheme/thermal/summary") { d.contains("critical") || d.contains("throttle") } else { false }
}

fn write_scheme_state(governor: Governor, cpus: &[CpuInfo]) {
    let mut out = format!("governor={:?}\n", governor);
    for ci in cpus {
        if ci.pstates.is_empty() { continue; }
        let p = &ci.pstates[ci.current_idx.min(ci.pstates.len() - 1)];
        out.push_str(&format!("CPU{}: {} kHz, {} mW, load={:.1}%\n", ci.id, p.freq_khz, p.power_mw, avg_load(ci) * 100.0));
    }
    let _ = fs::write("/scheme/cpufreq/state", out);
}

fn main() {
    log::set_logger(&StderrLogger).ok();
    log::set_max_level(LevelFilter::Info);
    let governor = match env::var("CPUFREQ_GOVERNOR").unwrap_or_default().as_str() {
        "performance" => Governor::Performance, "powersave" => Governor::Powersave,
        "conservative" => Governor::Conservative, "schedutil" => Governor::Schedutil,
        _ => Governor::Ondemand,
    };
    let cpus = detect_cpus();
    info!("detected {} CPU(s), governor={:?}", cpus.len(), governor);
    let mut ci: Vec<CpuInfo> = cpus.iter().map(|&id| {
        let ps = read_acpi_pss(id);
        info!("CPU{}: {} P-states ({} - {} kHz)", id, ps.len(), ps.first().map_or(0, |p| p.freq_khz), ps.last().map_or(0, |p| p.freq_khz));
        CpuInfo { id, pstates: ps, current_idx: 0, load_history: [0.0; SAMPLE_WINDOW], load_idx: 0, throttle: false, msr_errors: 0, msr_suppressed: false }
    }).collect();
    let mut prev: Vec<(u64, u64)> = vec![(0, 0); cpus.len()];
    let mut thermal = ThermalCache::new();
    let mut last_state_write = Instant::now();
    for c in &ci { if !c.pstates.is_empty() { write_msr(c.id, IA32_PERF_CTL, c.pstates[0].ctl); } }
    loop {
        std::thread::sleep(Duration::from_millis(POLL_MS));
        let tt = thermal.get();
        for (i, c) in ci.iter_mut().enumerate() {
            if c.pstates.is_empty() { continue; }
            let l = measure_load(c.id, &mut prev[i]);
            c.load_history[c.load_idx] = l; c.load_idx = (c.load_idx + 1) % SAMPLE_WINDOW; c.throttle = tt;
            let n = choose_pstate(governor, c);
            if n != c.current_idx && n < c.pstates.len() {
                let ct = c.pstates[n].ctl;
                if write_msr(c.id, IA32_PERF_CTL, ct) {
                    info!("CPU{}: P{}→P{} ({}→{} kHz, load={:.0}%)", c.id, c.current_idx, n, c.pstates[c.current_idx].freq_khz, c.pstates[n].freq_khz, l * 100.0);
                    c.current_idx = n; c.msr_errors = 0; c.msr_suppressed = false;
                } else {
                    c.msr_errors += 1;
                    if !c.msr_suppressed { warn!("CPU{}: MSR write failed ({}/{})", c.id, c.msr_errors, MSR_ERROR_SUPPRESS_COUNT); if c.msr_errors >= MSR_ERROR_SUPPRESS_COUNT { c.msr_suppressed = true; } }
                }
            }
        }
        if last_state_write.elapsed() >= Duration::from_secs(STATE_WRITE_INTERVAL_S) { write_scheme_state(governor, &ci); last_state_write = Instant::now(); }
    }
}