#![feature(yggdrasil_os, rustc_private)]

use std::{
    collections::{HashMap, HashSet},
    fs, io,
    os::fd::AsRawFd,
    path::{Path, PathBuf},
    time::{Duration, Instant},
};

use cross::io::{Poll, PollEvent, TimerFd};
use libterm::{Term, TermKey, tui::PrettyGauge};
use libutil::fmt::FormatSize;
use ratatui::{
    Frame, Terminal,
    layout::{Constraint, Layout, Rect},
    style::{Color, Style},
    text::{Line, Span},
    widgets::{Cell, Paragraph, Row, Table},
};
use runtime::{abi::system, rt::system::SystemMemoryStats};

struct ProcessInfo {
    pid: u32,
    name: String,
    run_time: u64,
    threads: Vec<ThreadInfo>,
}

struct ThreadInfo {
    tid: u32,
    name: String,
}

struct State {
    processes: Vec<ProcessInfo>,
    memory: Option<SystemMemoryStats>,
    expanded: HashSet<u32>,
    cpu_load: HashMap<usize, CpuStats>,
    cpu_stats: HashMap<usize, CpuStats>,
    selection: usize,
}

struct CpuStats {
    total_ns: u64,
    idle_ns: u64,
}

impl CpuStats {
    fn read_ns<P: AsRef<Path>>(path: P) -> Option<u64> {
        fs::read_to_string(path).ok()?.trim().parse().ok()
    }

    fn read(index: usize) -> Self {
        let path = PathBuf::from("/sys/kernel/sched").join(format!("{index}"));
        let total_ns = Self::read_ns(path.join("total")).unwrap_or_default();
        let idle_ns = Self::read_ns(path.join("idle")).unwrap_or_default();
        Self { total_ns, idle_ns }
    }
}

fn memory_info() -> io::Result<SystemMemoryStats> {
    runtime::rt::system::get_system_info!(system::MemoryUsage).map_err(io::Error::from)
}

fn cpu_stats(stats: &mut HashMap<usize, CpuStats>, load: &mut HashMap<usize, CpuStats>) {
    if stats.is_empty() {
        let Ok(dir) = fs::read_dir("/sys/kernel/sched") else {
            return;
        };
        for entry in dir {
            let Ok(entry) = entry else { continue };
            let filename = entry.file_name();
            let Some(filename) = filename.to_str() else {
                continue;
            };
            let Ok(index) = filename.parse() else {
                continue;
            };
            let data = CpuStats::read(index);
            stats.insert(index, data);
        }
        return;
    }

    for (cpu, stats) in stats.iter_mut() {
        let new_stats = CpuStats::read(*cpu);
        let idle_ns = new_stats.idle_ns.wrapping_sub(stats.idle_ns);
        let total_ns = new_stats.total_ns.wrapping_sub(stats.total_ns).max(idle_ns);
        load.insert(*cpu, CpuStats { total_ns, idle_ns });
        *stats = new_stats;
    }
}

fn list_ids<P: AsRef<Path>>(path: P) -> io::Result<Vec<u32>> {
    let mut list = vec![];
    for entry in fs::read_dir(path)? {
        let Ok(entry) = entry else {
            continue;
        };
        let filename = entry.file_name();
        let Some(name) = filename.to_str() else {
            continue;
        };
        let Ok(pid) = name.parse() else {
            continue;
        };
        list.push(pid);
    }
    Ok(list)
}

fn process_info(pid: u32) -> io::Result<ProcessInfo> {
    let base = PathBuf::from("/sys/proc").join(format!("{pid}"));
    let name = fs::read_to_string(base.join("name"))?;
    let run_time = fs::read_to_string(base.join("runtime"))
        .ok()
        .and_then(|p| p.trim().parse::<u64>().ok())
        .unwrap_or_default();
    let threads = threads(base).unwrap_or_default();
    Ok(ProcessInfo {
        pid,
        name,
        threads,
        run_time,
    })
}

fn processes() -> io::Result<Vec<ProcessInfo>> {
    let pids = list_ids("/sys/proc")?;
    let processes = pids
        .into_iter()
        .filter_map(|pid| process_info(pid).ok())
        .collect();
    Ok(processes)
}

fn thread_info<P: AsRef<Path>>(pid_path: P, tid: u32) -> io::Result<ThreadInfo> {
    let base = pid_path.as_ref().join(format!("{tid}"));
    let name = fs::read_to_string(base.join("name"))?;
    Ok(ThreadInfo { tid, name })
}

fn threads<P: AsRef<Path>>(pid_path: P) -> io::Result<Vec<ThreadInfo>> {
    let pid_path = pid_path.as_ref();
    let tids = list_ids(pid_path)?;
    let threads = tids
        .into_iter()
        .filter_map(|tid| thread_info(pid_path, tid).ok())
        .collect();
    Ok(threads)
}

impl State {
    fn new() -> io::Result<Self> {
        let processes = processes()?;
        let mut this = Self {
            processes,
            memory: None,
            cpu_load: HashMap::new(),
            cpu_stats: HashMap::new(),
            expanded: HashSet::new(),
            selection: 0,
        };
        this.refresh()?;
        Ok(this)
    }

    fn next(&mut self) {
        if self.processes.is_empty() {
            self.selection = 0;
            return;
        }
        if self.selection < self.processes.len() - 1 {
            self.selection += 1;
        }
    }

    fn prev(&mut self) {
        if self.selection > 0 {
            self.selection -= 1;
        }
    }

    fn toggle_selected(&mut self) {
        if let Some(selected) = self.processes.get(self.selection) {
            let pid = selected.pid;
            if self.expanded.contains(&pid) {
                self.expanded.remove(&pid);
            } else {
                self.expanded.insert(pid);
            }
        }
    }

    fn kill_selected(&self) -> io::Result<()> {
        if let Some(selected) = self.processes.get(self.selection) {
            cross::signal::send_kill(selected.pid)?;
        }
        Ok(())
    }

    fn refresh(&mut self) -> io::Result<()> {
        let selected_pid = self
            .processes
            .get(self.selection)
            .map(|process| process.pid);
        cpu_stats(&mut self.cpu_stats, &mut self.cpu_load);
        self.memory = memory_info().ok();
        self.processes = processes()?;
        let new_selection = selected_pid
            .and_then(|pid| self.processes.iter().position(|process| process.pid == pid));
        self.selection = new_selection.unwrap_or(0);

        Ok(())
    }

    fn render_processes(&self, frame: &mut Frame, area: Rect) {
        let normal_style = Style::default();
        let selected_style = Style::default().bg(Color::Cyan).fg(Color::Black);
        let hint_style = Style::default().bg(Color::Blue).fg(Color::White);

        let rows = self.processes.iter().enumerate().flat_map(|(i, process)| {
            let style = if i == self.selection {
                selected_style
            } else {
                normal_style
            };

            let run_hours = process.run_time / 3600;
            let run_minutes = (process.run_time / 60) % 60;
            let run_seconds = process.run_time % 60;

            let mut rows = vec![
                Row::new(vec![
                    Cell::from(format!("{}", process.pid)),
                    Cell::from(""),
                    Cell::from(process.name.as_str()),
                    Cell::from(format!("{run_hours:02}:{run_minutes:02}:{run_seconds:02}")),
                ])
                .height(1)
                .style(style),
            ];

            if self.expanded.contains(&process.pid) {
                let count = process.threads.len();
                for (i, thread) in process.threads.iter().enumerate() {
                    let ch = if i == count - 1 { "└" } else { "├" };
                    rows.push(
                        Row::new(vec![
                            Cell::from(ch),
                            Cell::from(format!("{}", thread.tid)),
                            Cell::from(thread.name.as_str()),
                            Cell::from(""),
                        ])
                        .style(style),
                    );
                }
            }

            rows
        });

        let header = Row::new(vec![
            Cell::from("PID"),
            Cell::from("TID"),
            Cell::from("NAME"),
            Cell::from("Time"),
        ])
        .height(1)
        .style(hint_style);

        let table = Table::new(
            rows,
            &[
                Constraint::Min(6),
                Constraint::Min(6),
                Constraint::Percentage(100),
                Constraint::Min(9),
            ],
        )
        .highlight_symbol(">>")
        .header(header);

        frame.render_widget(table, area);
    }

    fn render_keybar(&self, frame: &mut Frame, area: Rect) {
        const KEYS: &[(&str, &str)] = &[
            (" j ", "Down"),
            (" k ", "Up"),
            (" Enter ", "Toggle"),
            (" K ", "Kill"),
        ];

        let hint_style = Style::default().bg(Color::Blue).fg(Color::White);
        let key_style = Style::default().bg(Color::Cyan).fg(Color::Black);

        let mut status_spans = vec![];
        for (i, &(key, hint)) in KEYS.iter().enumerate() {
            if i != 0 {
                status_spans.push(Span::from(" "));
            }
            status_spans.push(Span::styled(key, key_style));
            status_spans.push(Span::styled(hint, hint_style));
        }

        let keybar = Paragraph::new(vec![Line::from(status_spans)]);

        frame.render_widget(keybar, area);
    }

    fn render_stats_header(&self, frame: &mut Frame, area: Rect) {
        let cpu_count = self.cpu_stats.len().max(self.cpu_load.len());
        let cpu_col_count = (frame.area().width as usize / 32).max(1);
        let cpu_row_count = cpu_count.div_ceil(cpu_col_count);

        let row_layout = vec![Constraint::Min(1); 1 + cpu_row_count];
        let row_rects = Layout::vertical(&row_layout).split(area);
        let mut cpu_constraints = vec![];
        let mut cpu_row_rects = vec![];
        for _ in 0..cpu_col_count {
            cpu_constraints.push(Constraint::Fill(1));
        }
        for row_rect in &row_rects[1..] {
            cpu_row_rects.push(
                Layout::horizontal(&cpu_constraints)
                    .spacing(1)
                    .split(*row_rect),
            );
        }

        // Memory gauge
        let (memory_total, memory_used) = match self.memory.as_ref() {
            Some(memory) => (memory.total_usable_pages, memory.allocated_pages),
            None => (1, 0),
        };
        let memory_total = memory_total.max(1);
        let usage = memory_used as f64 / memory_total as f64;
        let fill_color = match usage {
            u if u >= 0.85 => Color::from_u32(0x00FF0000),
            u if u >= 0.65 => Color::from_u32(0x00FF8800),
            u if u >= 0.45 => Color::from_u32(0x00FFFF00),
            _ => Color::from_u32(0x0000FF00),
        };
        let memory_used = memory_used * 4096;
        let memory_total = memory_total * 4096;
        // let label = Paragraph::new();
        let label = format!(
            "{}/{}",
            FormatSize::default(memory_used as u64),
            FormatSize::default(memory_total as u64)
        );
        let gauge = PrettyGauge::default()
            .filled_color(fill_color)
            .unfilled_color(Color::DarkGray)
            .prefix("Mem [")
            .inner_label(label)
            .ratio(usage);
        frame.render_widget(gauge, row_rects[0]);

        // CPU gauges
        for cpu in 0..cpu_count {
            let row = cpu / cpu_col_count;
            let col = cpu % cpu_col_count;
            let gauge_rect = cpu_row_rects[row][col];
            let load = match self.cpu_load.get(&cpu) {
                Some(load) if load.total_ns > 0 => 1.0 - load.idle_ns as f64 / load.total_ns as f64,
                _ => 0.0,
            };
            let label = format!("{:<3.1}%", load * 100.0);
            let gauge = PrettyGauge::default()
                .prefix(format!("{cpu:>2} ["))
                .filled_color(Color::Yellow)
                .unfilled_color(Color::DarkGray)
                .inner_label(label)
                .ratio(load);

            frame.render_widget(gauge, gauge_rect);
        }
    }

    fn render(&self, f: &mut Frame) {
        let cpus_per_row = f.area().width as usize / 32;
        let cpu_row_count = self
            .cpu_load
            .len()
            .max(self.cpu_stats.len())
            .div_ceil(cpus_per_row);

        let rects = Layout::default()
            .constraints(
                [
                    Constraint::Min(1 + cpu_row_count as u16),
                    Constraint::Percentage(99),
                    Constraint::Min(1),
                ]
                .as_ref(),
            )
            .split(f.area());

        self.render_stats_header(f, rects[0]);
        self.render_processes(f, rects[1]);
        self.render_keybar(f, rects[2]);
    }
}

fn run() -> io::Result<()> {
    const INTERVAL: Duration = Duration::from_millis(250);

    let mut state = State::new()?;

    let mut poll = Poll::new()?;
    let mut timer = TimerFd::new()?;

    let mut running = true;
    let term = Term::open()?;
    let term_fd = term.as_raw_fd();

    poll.add(&timer)?;
    poll.add(&term_fd)?;

    let mut term = Terminal::new(term)?;

    timer.start(INTERVAL)?;

    let mut last_refresh = Instant::now();
    state.refresh()?;

    while running {
        if last_refresh.elapsed() >= INTERVAL {
            last_refresh = Instant::now();
            state.refresh()?;
        }

        term.draw(|f| state.render(f))?;

        match poll.wait(Some(Duration::from_millis(100)))? {
            PollEvent::Ready(fd) if fd == timer.as_raw_fd() => {
                timer.start(INTERVAL)?;
            }
            PollEvent::Ready(fd) if fd == term_fd => {
                let key = term.backend_mut().read_key()?;
                match key {
                    TermKey::Escape | TermKey::Char('q') => {
                        running = false;
                    }
                    TermKey::Char('\n') => {
                        state.toggle_selected();
                    }
                    TermKey::Char('K') => {
                        if let Err(error) = state.kill_selected() {
                            log::error!("Kill: {error}");
                        }
                    }
                    TermKey::Char('k') | TermKey::Up => {
                        state.prev();
                    }
                    TermKey::Char('j') | TermKey::Down => {
                        state.next();
                    }
                    _ => (),
                }
            }
            PollEvent::None => {
                timer.start(INTERVAL)?;
            }
            _ => unreachable!(),
        }
    }

    Ok(())
}

fn main() {
    logsink::setup_logging(false);
    match run() {
        Ok(()) => (),
        Err(error) => {
            eprintln!("{error}");
        }
    }
}
