use std::{
    fs, io,
    path::{Path, PathBuf},
    process::ExitCode,
};

use clap::Parser;

#[derive(Debug, Parser)]
struct Args {
    #[clap(short, long, help = "Display threads of processes")]
    threads: bool,
}

#[derive(Debug)]
struct ThreadInfo {
    tid: u32,
    name: String,
}

#[derive(Debug)]
struct ProcessInfo {
    pid: u32,
    name: String,
    threads: Vec<ThreadInfo>,
}

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, args: &Args) -> io::Result<ProcessInfo> {
    let base = PathBuf::from("/sys/proc").join(format!("{pid}"));
    let name = fs::read_to_string(base.join("name"))?;
    let threads = if args.threads {
        threads(base).unwrap_or_default()
    } else {
        vec![]
    };
    Ok(ProcessInfo { pid, name, threads })
}

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)
}

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

fn run(args: &Args) -> io::Result<()> {
    let mut processes = processes(args)?;
    processes.sort_by(|a, b| Ord::cmp(&a.pid, &b.pid));
    if args.threads {
        println!("{:>4} {:>4} NAME", "PID", "TID");
        for process in processes {
            println!("{:4} {:4} {}", process.pid, "", process.name.trim());
            let count = process.threads.len();
            for (i, thread) in process.threads.into_iter().enumerate() {
                let ch = if i == count - 1 { '└' } else { '├' };
                println!("{:>4} {:4} {}", ch, thread.tid, thread.name.trim());
            }
        }
    } else {
        println!("{:>4} NAME", "PID");
        for process in processes {
            println!("{:4} {}", process.pid, process.name.trim());
        }
    }
    Ok(())
}

fn main() -> ExitCode {
    let args = Args::parse();
    match run(&args) {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("{error}");
            ExitCode::FAILURE
        }
    }
}
