use anyhow::{Context, Result};
use async_trait::async_trait;
use prometheus_client::encoding::text::Encode;
use prometheus_client::metrics::family::Family;
use prometheus_client::registry::Registry;
use serde::Deserialize;
use tokio::process::Command;
use tracing::trace;

use crate::{F64Gauge, Metrics, HOSTNAME};

#[derive(Clone, Hash, PartialEq, Eq, Encode)]
struct Label {
    hostname: &'static str,
    cpu: String,
}

#[derive(Default)]
pub struct Exporter {
    usr: Family<Label, F64Gauge>,
    nice: Family<Label, F64Gauge>,
    sys: Family<Label, F64Gauge>,
    iowait: Family<Label, F64Gauge>,
    irq: Family<Label, F64Gauge>,
    soft: Family<Label, F64Gauge>,
    steal: Family<Label, F64Gauge>,
    guest: Family<Label, F64Gauge>,
    gnice: Family<Label, F64Gauge>,
    idle: Family<Label, F64Gauge>,
}

#[async_trait]
impl Metrics for Exporter {
    fn prefix(&self) -> &'static str {
        "cpu"
    }

    async fn should_collect(&self) -> Result<()> {
        Command::new("mpstat")
            .output()
            .await
            .context("While checking mpstat")
            .map(drop)
    }

    fn register(&self, sub_registry: &mut Registry) -> Result<()> {
        sub_registry.register("usr", "usr time (mpstat)", Box::new(self.usr.clone()));
        sub_registry.register("nice", "nice time (mpstat)", Box::new(self.nice.clone()));
        sub_registry.register("sys", "sys time (mpstat)", Box::new(self.sys.clone()));
        sub_registry.register(
            "iowait",
            "iowait time (mpstat)",
            Box::new(self.iowait.clone()),
        );
        sub_registry.register("irq", "irq time (mpstat)", Box::new(self.irq.clone()));
        sub_registry.register("soft", "soft time (mpstat)", Box::new(self.soft.clone()));
        sub_registry.register("steal", "steal time (mpstat)", Box::new(self.steal.clone()));
        sub_registry.register("guest", "guest time (mpstat)", Box::new(self.guest.clone()));
        sub_registry.register(
            "gnice",
            "guest nice time (mpstat)",
            Box::new(self.gnice.clone()),
        );
        sub_registry.register("idle ", "idle time (mpstat)", Box::new(self.idle.clone()));
        Ok(())
    }

    async fn collect(&self) -> Result<()> {
        trace!("Started");

        let json = Command::new("mpstat")
            .args(["-o", "JSON", "-P", "ALL", "5", "1"])
            .output()
            .await
            .context("While collecting mpstat")?
            .stdout;
        let data: MpStatOutput = serde_json::from_slice(&json)?;
        let statistics = data
            .sysstat
            .hosts
            .into_iter()
            .next()
            .context("Getting the first host")?
            .statistics;

        let cpus = statistics
            .into_iter()
            .next()
            .context("getting first stat measurement")?
            .cpu_load;

        for cpu in cpus {
            let label = Label {
                hostname: &HOSTNAME,
                cpu: cpu.identifier,
            };
            self.usr.get_or_create(&label).set(cpu.usr);
            self.nice.get_or_create(&label).set(cpu.nice);
            self.sys.get_or_create(&label).set(cpu.sys);
            self.iowait.get_or_create(&label).set(cpu.iowait);
            self.irq.get_or_create(&label).set(cpu.irq);
            self.soft.get_or_create(&label).set(cpu.soft);
            self.steal.get_or_create(&label).set(cpu.steal);
            self.guest.get_or_create(&label).set(cpu.guest);
            self.gnice.get_or_create(&label).set(cpu.gnice);
            self.idle.get_or_create(&label).set(cpu.idle);
        }

        trace!("Done");

        Ok(())
    }
}

#[derive(Deserialize)]
struct MpStatOutput {
    sysstat: SysStat,
}

#[derive(Deserialize)]
struct SysStat {
    hosts: Vec<Host>,
}

#[derive(Deserialize)]
struct Host {
    statistics: Vec<Statistic>,
}

#[derive(Deserialize)]
#[serde(rename_all = "kebab-case")]
struct Statistic {
    cpu_load: Vec<CpuLoad>,
}

#[derive(Deserialize)]
struct CpuLoad {
    #[serde(rename = "cpu")]
    identifier: String,
    usr: f64,
    nice: f64,
    sys: f64,
    iowait: f64,
    irq: f64,
    soft: f64,
    steal: f64,
    guest: f64,
    gnice: f64,
    idle: f64,
}

#[cfg(test)]
mod mp_stat {
    use anyhow::Result;
    use serde_json::json;

    use super::MpStatOutput;

    #[test]
    fn deserializes() -> Result<(), serde_json::Error> {
        let raw = json!({
            "sysstat": {
                "hosts": [{
                    "nodename": "kurante",
                    "sysname": "Linux",
                    "release": "5.17.7-zen1-1-zen",
                    "machine": "x86_64",
                    "number-of-cpus": 16,
                    "date": "05/27/2022",
                    "statistics": [{
                        "timestamp": "12:46:14 PM",
                        "cpu-load": [
                            {"cpu": "all", "usr": 3.19, "nice": 0.03, "sys": 0.73, "iowait": 0.08, "irq": 0.14, "soft": 0.08, "steal": 0.00, "guest": 0.00, "gnice": 0.00, "idle": 95.77},
                            {"cpu": "0", "usr": 2.41, "nice": 0.00, "sys": 0.80, "iowait": 0.00, "irq": 0.80, "soft": 0.20, "steal": 0.00, "guest": 0.00, "gnice": 0.00, "idle": 95.78},
                            {"cpu": "1", "usr": 1.20, "nice": 0.20, "sys": 0.80, "iowait": 1.00, "irq": 0.00, "soft": 0.20, "steal": 0.00, "guest": 0.00, "gnice": 0.00, "idle": 96.61},
                            {"cpu": "2", "usr": 3.58, "nice": 0.00, "sys": 1.19, "iowait": 0.20, "irq": 0.60, "soft": 0.40, "steal": 0.00, "guest": 0.00, "gnice": 0.00, "idle": 94.04},
                            {"cpu": "3", "usr": 3.21, "nice": 0.00, "sys": 1.00, "iowait": 0.00, "irq": 0.20, "soft": 0.00, "steal": 0.00, "guest": 0.00, "gnice": 0.00, "idle": 95.59},
                            {"cpu": "4", "usr": 3.62, "nice": 0.00, "sys": 0.40, "iowait": 0.00, "irq": 0.00, "soft": 0.00, "steal": 0.00, "guest": 0.00, "gnice": 0.00, "idle": 95.98},
                            {"cpu": "5", "usr": 3.21, "nice": 0.00, "sys": 0.40, "iowait": 0.00, "irq": 0.00, "soft": 0.00, "steal": 0.00, "guest": 0.00, "gnice": 0.00, "idle": 96.39},
                            {"cpu": "6", "usr": 4.22, "nice": 0.00, "sys": 0.60, "iowait": 0.00, "irq": 0.00, "soft": 0.00, "steal": 0.00, "guest": 0.00, "gnice": 0.00, "idle": 95.18},
                            {"cpu": "7", "usr": 2.81, "nice": 0.00, "sys": 0.40, "iowait": 0.00, "irq": 0.00, "soft": 0.00, "steal": 0.00, "guest": 0.00, "gnice": 0.00, "idle": 96.79},
                            {"cpu": "8", "usr": 2.01, "nice": 0.00, "sys": 0.40, "iowait": 0.00, "irq": 0.20, "soft": 0.00, "steal": 0.00, "guest": 0.00, "gnice": 0.00, "idle": 97.38},
                            {"cpu": "9", "usr": 2.18, "nice": 0.00, "sys": 1.39, "iowait": 0.00, "irq": 0.20, "soft": 0.00, "steal": 0.00, "guest": 0.00, "gnice": 0.00, "idle": 96.23},
                            {"cpu": "10", "usr": 2.80, "nice": 0.00, "sys": 0.80, "iowait": 0.00, "irq": 0.00, "soft": 0.20, "steal": 0.00, "guest": 0.00, "gnice": 0.00, "idle": 96.20},
                            {"cpu": "11", "usr": 5.18, "nice": 0.00, "sys": 1.00, "iowait": 0.00, "irq": 0.20, "soft": 0.20, "steal": 0.00, "guest": 0.00, "gnice": 0.00, "idle": 93.43},
                            {"cpu": "12", "usr": 2.61, "nice": 0.00, "sys": 0.40, "iowait": 0.00, "irq": 0.00, "soft": 0.00, "steal": 0.00, "guest": 0.00, "gnice": 0.00, "idle": 96.99},
                            {"cpu": "13", "usr": 4.83, "nice": 0.00, "sys": 0.40, "iowait": 0.00, "irq": 0.00, "soft": 0.00, "steal": 0.00, "guest": 0.00, "gnice": 0.00, "idle": 94.77},
                            {"cpu": "14", "usr": 1.79, "nice": 0.20, "sys": 0.80, "iowait": 0.00, "irq": 0.00, "soft": 0.00, "steal": 0.00, "guest": 0.00, "gnice": 0.00, "idle": 97.21},
                            {"cpu": "15", "usr": 5.41, "nice": 0.00, "sys": 0.80, "iowait": 0.00, "irq": 0.00, "soft": 0.00, "steal": 0.00, "guest": 0.00, "gnice": 0.00, "idle": 93.79}
                        ]
                    }]
                }]
            }
        });

        serde_json::from_value::<MpStatOutput>(raw).map(drop)
    }
}