#![feature(yggdrasil_os)]
use std::{
    collections::HashSet,
    fs, io,
    os::fd::AsRawFd,
    path::{Path, PathBuf},
    time::Duration,
};

use cross::io::{Poll, TimerFd};
use libterm::{Term, TermKey};
use tui::{
    layout::{Constraint, Layout},
    style::{Color, Style},
    text::{Span, Spans},
    widgets::{Cell, Paragraph, Row, Table},
    Frame, Terminal,
};

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

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

struct State {
    processes: Vec<ProcessInfo>,
    expanded: HashSet<u32>,
    selection: usize,
}

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 threads = threads(base).unwrap_or_default();
    Ok(ProcessInfo { pid, name, threads })
}

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()?;
        Ok(Self {
            processes,
            expanded: HashSet::new(),
            selection: 0,
        })
    }

    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);
        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(&self, f: &mut Frame<Term>) {
        const KEYS: &[(&str, &str)] = &[
            (" j ", "Down"),
            (" k ", "Up"),
            (" Enter ", "Toggle"),
            (" K ", "Kill"),
        ];

        let rects = Layout::default()
            .constraints([Constraint::Percentage(99), Constraint::Min(1)].as_ref())
            .split(f.size());

        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 key_style = Style::default().bg(Color::Cyan).fg(Color::Black);

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

            let mut rows = vec![Row::new(vec![
                Cell::from(format!("{}", process.pid)),
                Cell::from(""),
                Cell::from(process.name.as_str()),
            ])
            .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()),
                        ])
                        .style(style),
                    );
                }
            }

            rows
        });

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

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

        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 status = Paragraph::new(vec![Spans::from(status_spans)]);

        f.render_widget(table, rects[0]);
        f.render_widget(status, rects[1]);
    }
}

fn run() -> io::Result<()> {
    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(Duration::from_secs(1))?;

    while running {
        term.draw(|f| state.render(f))?;

        match poll.wait(None)?.unwrap() {
            fd if fd == timer.as_raw_fd() => {
                timer.start(Duration::from_secs(1))?;
                state.refresh()?;
            }
            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();
                    }
                    _ => (),
                }
            }
            _ => unreachable!(),
        }
    }

    Ok(())
}

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