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

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

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

#[derive(Default)]
pub struct Exporter {
    size: Family<Label, U64Gauge>,
    allocated: Family<Label, U64Gauge>,
    free: Family<Label, U64Gauge>,
    fragmentation: Family<Label, U64Gauge>,
    capacity: Family<Label, U64Gauge>,
    dedupe: Family<Label, F64Gauge>,
    health: Family<Label, U32Gauge>,
}

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

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

    fn register(&self, sub_registry: &mut Registry) -> Result<()> {
        sub_registry.register("size", "zpool size", Box::new(self.size.clone()));
        sub_registry.register(
            "allocated",
            "zpool allocated",
            Box::new(self.allocated.clone()),
        );
        sub_registry.register("free", "zpool free", Box::new(self.free.clone()));
        sub_registry.register(
            "fragmentation",
            "zpool fragmentation",
            Box::new(self.fragmentation.clone()),
        );
        sub_registry.register(
            "capacity",
            "zpool capacity",
            Box::new(self.capacity.clone()),
        );
        sub_registry.register("dedupe", "zpool dedupe", Box::new(self.dedupe.clone()));
        sub_registry.register("health", "zpool health", Box::new(self.health.clone()));
        Ok(())
    }

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

        let zpool = Command::new("zpool")
            .args([
                "list",
                "-H",
                "-p",
                "-o",
                "name,size,alloc,free,frag,cap,dedup,health",
            ])
            .output()
            .await?;
        for string_data in String::from_utf8(zpool.stdout)?.lines() {
            let mut info = string_data.split_whitespace();
            let zpool_name = info.next().context("getting the zpool name")?;
            let [size, alloc, free, frag, cap, dedup, health] = info
                .collect::<ArrayVec<_, 7>>()
                .into_inner()
                .map_err(|_| anyhow::anyhow!("parsing zpool info"))?;

            let label = Label {
                hostname: &*HOSTNAME,
                zpool: zpool_name.to_string(),
            };

            self.size.get_or_create(&label).set(size.parse()?);
            self.allocated.get_or_create(&label).set(alloc.parse()?);
            self.free.get_or_create(&label).set(free.parse()?);
            self.fragmentation.get_or_create(&label).set(frag.parse()?);
            self.capacity.get_or_create(&label).set(cap.parse()?);
            self.dedupe.get_or_create(&label).set(dedup.parse()?);
            self.health
                .get_or_create(&label)
                .set(if health == "ONLINE" { 0 } else { 1 });
        }

        trace!("Done");

        Ok(())
    }
}