use std::collections::HashSet;
use std::io::{self, Write};
use std::thread;
use std::time::{Duration, Instant};

mod cpu;
mod gpu;
mod mem;
mod temp;

fn parse_interval_ms() -> u64 {
    let args: Vec<String> = std::env::args().collect();
    let mut i = 1;
    while i < args.len() {
        if args[i] == "--interval" {
            if let Some(ms) = args.get(i + 1).and_then(|s| s.parse().ok()) {
                return ms;
            }
        }
        i += 1;
    }
    1000
}

fn parse_types() -> HashSet<String> {
    let args: Vec<String> = std::env::args().collect();
    let mut i = 1;
    while i < args.len() {
        if args[i] == "--types" {
            if let Some(list) = args.get(i + 1) {
                return list.split(',').map(|s| s.trim().to_string()).collect();
            }
        }
        i += 1;
    }
    HashSet::new() // empty = all types
}

fn main() {
    let interval = Duration::from_millis(parse_interval_ms());
    let types = parse_types();
    let stdout = io::stdout();
    let mut out = io::BufWriter::new(stdout.lock());
    let mut prev: Vec<cpu::Sample> = vec![];
    let mut freqs: Vec<f64> = vec![];
    let mut tick = 0u64;

    let emit_cpu = types.is_empty() || types.contains("cpu");
    let emit_mem = types.is_empty() || types.contains("mem");
    let emit_temp = types.is_empty() || types.contains("temp");
    let emit_graphics = types.is_empty() || types.contains("gpu");

    let mut gpu = if emit_graphics {
        gpu::detect_gpu()
    } else {
        gpu::GpuBackend::None
    };

    loop {
        let t0 = Instant::now();

        if emit_cpu {
            let curr = cpu::read_stat();
            if tick.is_multiple_of(2) {
                freqs = cpu::read_core_freqs();
            }
            cpu::emit_cpu(&mut out, &prev, &curr, &freqs);
            prev = curr;
        }

        if emit_mem && tick.is_multiple_of(2) {
            mem::emit_mem(&mut out);
        }

        if tick.is_multiple_of(4) {
            if emit_temp {
                temp::emit_temp(&mut out);
            }
            if emit_graphics {
                let emit = match &gpu {
                    gpu::GpuBackend::Amd { .. } | gpu::GpuBackend::Intel { .. } => true,
                    gpu::GpuBackend::Nvidia => tick.is_multiple_of(8),
                    gpu::GpuBackend::None => false,
                };
                if emit {
                    gpu::emit_gpu(&mut out, &mut gpu);
                }
            }
        }

        let _ = out.flush();
        tick += 1;

        let elapsed = t0.elapsed();
        if let Some(remaining) = interval.checked_sub(elapsed) {
            thread::sleep(remaining);
        }
    }
}