use std::{collections::HashMap, os::unix::net::UnixStream, time::Instant};
use std::{fs, io};

use procfs::process::Process;

use crate::energy::{connect_to_service, request_all_processes, Pid, ProcessInfo};

use crate::SOCKET_PATH;

pub struct ProcessData {
    pub processes: Vec<Pid>,
    pub total_energy_history: Vec<(f64, f64)>,
    pub started_at: Instant,
    pub socket: UnixStream,
    pub process_info: HashMap<Pid, ProcessInfo>,
    pub rapl_offset: f64,
}

impl ProcessData {
    pub fn new() -> Self {
        Self {
            processes: Vec::new(),
            total_energy_history: Vec::new(),
            started_at: Instant::now(),
            process_info: HashMap::new(),
            socket: connect_to_service(SOCKET_PATH).expect("Failed to connect to socket"),
            rapl_offset: read_package_energy().unwrap(),
        }
    }

    pub fn update(&mut self) {
        self.fetch_processes();
    }

    pub fn fetch_processes(&mut self) {
        let energy_data = match request_all_processes(&mut self.socket) {
            Ok(data) => data,
            Err(_) => return, // Can't get energy data
        };

        // Use process data with procfs to get process names
        let mut updated_processes = Vec::new();
        let mut total_energy = 0.0;

        let elapsed = self.started_at.elapsed().as_secs_f64();
        for (pid, (energy, tree_energy)) in energy_data {
            let process_info = self.process_info.entry(pid).or_insert_with(|| {
                let name = get_process_name(pid).unwrap_or_else(|| format!("Process {}", pid));

                ProcessInfo::new(pid, name, energy, tree_energy)
            });

            let process_delta = energy - process_info.energy;
            let delta = tree_energy - process_info.tree_energy;
            let delta = (process_info.history.last().map(|x| x.1).unwrap_or(delta) + delta) * 0.5;
            process_info.energy = energy;
            process_info.tree_energy = tree_energy;
            process_info.add_history_point(elapsed, delta);

            updated_processes.push(pid);
            total_energy += process_delta;
        }

        // Sort by energy consumption (descending)
        updated_processes.sort_by(|a, b| {
            self.process_info[b]
                .energy
                .partial_cmp(&self.process_info[a].energy)
                .unwrap()
        });
        let mut sorted_processes = updated_processes.clone();

        // Update our local process list
        self.processes = updated_processes;
        sorted_processes.sort_unstable();

        self.process_info
            .retain(|pid, _| sorted_processes.binary_search(pid).is_ok());

        // Add total energy to history
        let elapsed = self.started_at.elapsed().as_secs_f64();

        let rapl_energy = read_package_energy().unwrap();
        let total_energy = rapl_energy - self.rapl_offset;
        self.rapl_offset = rapl_energy;

        self.add_total_energy_point(elapsed, total_energy);
    }

    fn add_total_energy_point(&mut self, timestamp: f64, energy: f64) {
        self.total_energy_history.push((timestamp, energy));

        // Keep history bounded
        if self.total_energy_history.len() > 60 {
            self.total_energy_history.remove(0);
        }
    }
}

// Helper function to get process name using procfs
fn get_process_name(pid: Pid) -> Option<String> {
    match Process::new(pid) {
        Ok(process) => match process.stat() {
            Ok(stat) => Some(stat.comm),
            Err(_) => None,
        },
        Err(_) => None,
    }
}

/// Read current package energy counter value in joules
pub fn read_package_energy() -> io::Result<f64> {
    let energy_path = "/sys/class/powercap/intel-rapl/intel-rapl:0/energy_uj";
    let energy_str = fs::read_to_string(energy_path)?;
    let energy_uj = energy_str
        .trim()
        .parse::<u64>()
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;

    // Convert from microjoules to joules
    Ok(energy_uj as f64 / 1_000_000.0)
}