use prometheus::{histogram_opts, opts, TextEncoder};
use rand::Rng;
use std::time::Duration;
use tokio::sync::OnceCell;
use trait_net::metrics::{
    prometheus::{
        ActiveCount, Elapsed, EndedCount, IntConstant, Latency, LatencyByStatus, StartedCount,
        StatusCount,
    },
    AsStatusLabel, ObservedFutureExt,
};

struct Metrics {
    constant: IntConstant,
    elapsed: Elapsed,
    started_count: StartedCount,
    active_count: ActiveCount,
    ended_count: EndedCount,
    status: StatusCount,
    latency: Latency,
    latency_by_status: LatencyByStatus,
}

impl Metrics {
    async fn global() -> &'static Self {
        static METRICS: OnceCell<Metrics> = OnceCell::const_new();
        METRICS
            .get_or_init(|| async {
                let constant = IntConstant::new(opts!("some_constant", "just 42"), 42).unwrap();
                let elapsed = Elapsed::new(opts!("elapsed", "example runtime")).unwrap();
                let started_count =
                    StartedCount::new(opts!("started_count", "started computations count"), &[])
                        .unwrap();
                let active_count =
                    ActiveCount::new(opts!("active_count", "active computations count"), &[])
                        .unwrap();
                let ended_count =
                    EndedCount::new(opts!("ended_count", "ended computations count"), &[]).unwrap();
                let status =
                    StatusCount::new(opts!("status", "computation status count"), &[]).unwrap();
                let latency =
                    Latency::new(histogram_opts!("latency", "computation latency"), &[]).unwrap();
                let latency_by_status = LatencyByStatus::new(
                    histogram_opts!("latency_by_status", "computation latency by status"),
                    &[],
                )
                .unwrap();
                Metrics {
                    constant,
                    elapsed,
                    started_count,
                    active_count,
                    ended_count,
                    latency,
                    status,
                    latency_by_status,
                }
            })
            .await
    }

    fn register(&self) {
        let registry = prometheus::default_registry();
        registry.register(Box::new(self.constant.clone())).unwrap();
        registry.register(Box::new(self.elapsed.clone())).unwrap();
        registry
            .register(Box::new(self.started_count.clone()))
            .unwrap();
        registry
            .register(Box::new(self.active_count.clone()))
            .unwrap();
        registry
            .register(Box::new(self.ended_count.clone()))
            .unwrap();
        registry.register(Box::new(self.latency.clone())).unwrap();
        registry.register(Box::new(self.status.clone())).unwrap();
        registry
            .register(Box::new(self.latency_by_status.clone()))
            .unwrap();
    }
}

enum Error {
    Ooops,
    BigOoops,
}

impl AsStatusLabel for Error {
    fn as_status_label(&self) -> String {
        match self {
            Error::Ooops => "ooops",
            Error::BigOoops => "big_ooops",
        }
        .to_owned()
    }
}

async fn computation() -> Result<(), Error> {
    let millis = rand::thread_rng().gen_range(0..2000);
    tokio::time::sleep(Duration::from_millis(millis)).await;
    match rand::thread_rng().gen_range(0..3) {
        0 => Ok(()),
        1 => Err(Error::Ooops),
        2 => Err(Error::BigOoops),
        _ => unreachable!(),
    }
}

async fn computation_loop() {
    loop {
        let metrics = Metrics::global().await;
        let _ = computation()
            .observe((
                metrics.started_count.make_observer(&[]),
                metrics.active_count.make_observer(&[]),
                metrics.ended_count.make_observer(&[]),
                metrics.latency.make_observer(&[]),
                metrics.status.make_observer(&[]),
                metrics.latency_by_status.make_observer(&[]),
            ))
            .await;
    }
}

async fn collect_loop() {
    loop {
        let text = TextEncoder::new()
            .encode_to_string(&prometheus::default_registry().gather())
            .unwrap();
        println!("{text}");
        println!("--------------------------------------------------");
        tokio::time::sleep(Duration::from_secs(1)).await;
    }
}

#[tokio::main]
async fn main() {
    Metrics::global().await.register();
    tokio::select! {
        _ = tokio::spawn(computation_loop()) => {},
        _ = tokio::spawn(computation_loop()) => {},
        _ = tokio::spawn(collect_loop()) => {},
    }
}