#![feature(async_await)]

use futures::executor::block_on;
use futures::stream::StreamExt;
use heim::{disk, memory};
use indexmap::IndexMap;
use nu::{
    serve_plugin, CallInfo, CommandConfig, Plugin, Primitive, ReturnSuccess, ReturnValue,
    ShellError, Span, Spanned, SpannedDictBuilder, Value,
};
use std::ffi::OsStr;

struct Sys;
impl Sys {
    fn new() -> Sys {
        Sys
    }
}

//TODO: add more error checking

async fn os(span: Span) -> Option<Spanned<Value>> {
    if let (Ok(name), Ok(version)) = (sys_info::os_type(), sys_info::os_release()) {
        let mut os_idx = SpannedDictBuilder::new(span);
        os_idx.insert("name", Primitive::String(name));
        os_idx.insert("version", Primitive::String(version));

        Some(os_idx.into_spanned_value())
    } else {
        None
    }
}

async fn cpu(span: Span) -> Option<Spanned<Value>> {
    if let (Ok(num_cpu), Ok(cpu_speed)) = (sys_info::cpu_num(), sys_info::cpu_speed()) {
        let mut cpu_idx = SpannedDictBuilder::new(span);
        cpu_idx.insert("cores", Primitive::Int(num_cpu as i64));
        cpu_idx.insert("speed", Primitive::Int(cpu_speed as i64));
        Some(cpu_idx.into_spanned_value())
    } else {
        None
    }
}

async fn mem(span: Span) -> Spanned<Value> {
    let memory = memory::memory().await.unwrap();
    let swap = memory::swap().await.unwrap();

    let mut dict = SpannedDictBuilder::new(span);

    dict.insert("total", Value::bytes(memory.total().get()));
    dict.insert("free", Value::bytes(memory.free().get()));

    dict.insert("swap total", Value::bytes(swap.total().get()));
    dict.insert("swap free", Value::bytes(swap.free().get()));

    dict.into_spanned_value()
}

async fn host(span: Span) -> Spanned<Value> {
    let mut dict = SpannedDictBuilder::new(span);

    // OS
    if let Ok(platform) = heim::host::platform().await {
        dict.insert("name", Value::string(platform.system()));
        dict.insert("release", Value::string(platform.release()));
        dict.insert("hostname", Value::string(platform.hostname()));
        dict.insert("arch", Value::string(platform.architecture().as_str()));
    }

    // Uptime
    if let Ok(uptime) = heim::host::uptime().await {
        let mut uptime_dict = SpannedDictBuilder::new(span);

        let uptime = uptime.get().round() as i64;
        let days = uptime / (60 * 60 * 24);
        let hours = (uptime - days * 60 * 60 * 24) / (60 * 60);
        let minutes = (uptime - days * 60 * 60 * 24 - hours * 60 * 60) / 60;
        let seconds = uptime % 60;

        uptime_dict.insert("days", Value::int(days));
        uptime_dict.insert("hours", Value::int(hours));
        uptime_dict.insert("mins", Value::int(minutes));
        uptime_dict.insert("secs", Value::int(seconds));

        dict.insert_spanned("uptime", uptime_dict.into_spanned_value());
    }

    // Users
    let mut users = heim::host::users();
    let mut user_vec = vec![];
    while let Some(user) = users.next().await {
        let user = user.unwrap();

        user_vec.push(Spanned {
            item: Value::string(user.username()),
            span,
        });
    }
    let user_list = Value::List(user_vec);
    dict.insert("users", user_list);

    dict.into_spanned_value()
}

async fn disks(span: Span) -> Value {
    let mut output = vec![];
    let mut partitions = disk::partitions_physical();
    while let Some(part) = partitions.next().await {
        let part = part.unwrap();
        let usage = disk::usage(part.mount_point().to_path_buf()).await.unwrap();

        let mut dict = SpannedDictBuilder::new(span);

        dict.insert(
            "device",
            Value::string(
                part.device()
                    .unwrap_or_else(|| OsStr::new("N/A"))
                    .to_string_lossy(),
            ),
        );

        dict.insert("type", Value::string(part.file_system().as_str()));
        dict.insert("mount", Value::string(part.mount_point().to_string_lossy()));
        dict.insert("total", Value::bytes(usage.total().get()));
        dict.insert("used", Value::bytes(usage.used().get()));
        dict.insert("free", Value::bytes(usage.free().get()));

        output.push(dict.into_spanned_value());
    }

    Value::List(output)
}

async fn sysinfo(span: Span) -> Vec<Spanned<Value>> {
    let mut sysinfo = SpannedDictBuilder::new(span);

    sysinfo.insert_spanned("host", host(span).await);
    if let Some(cpu) = cpu(span).await {
        sysinfo.insert_spanned("cpu", cpu);
    }
    sysinfo.insert("disks", disks(span).await);
    sysinfo.insert_spanned("mem", mem(span).await);

    vec![sysinfo.into_spanned_value()]
}

impl Plugin for Sys {
    fn config(&mut self) -> Result<CommandConfig, ShellError> {
        Ok(CommandConfig {
            name: "sys".to_string(),
            positional: vec![],
            is_filter: true,
            is_sink: false,
            named: IndexMap::new(),
            rest_positional: true,
        })
    }
    fn begin_filter(&mut self, callinfo: CallInfo) -> Result<Vec<ReturnValue>, ShellError> {
        Ok(block_on(sysinfo(callinfo.name_span.unwrap()))
            .into_iter()
            .map(|x| ReturnSuccess::value(x))
            .collect())
    }

    fn filter(&mut self, _: Spanned<Value>) -> Result<Vec<ReturnValue>, ShellError> {
        Ok(vec![])
    }
}

fn main() {
    serve_plugin(&mut Sys::new());
}