use std::sync::{Arc, LazyLock, RwLock};

use prometheus_client::encoding::text::encode;
use prometheus_client::registry::Registry;
use pyo3::prelude::*;

pub static DEFAULT_REGISTRY: LazyLock<PyRegistry> = LazyLock::new(PyRegistry::default);

#[derive(Clone)]
#[pyclass(name = "CollectorRegistry", from_py_object)]
pub struct PyRegistry {
    pub inner: Arc<RwLock<Registry>>,
}

impl Default for PyRegistry {
    fn default() -> Self {
        Self {
            inner: Arc::new(RwLock::new(Registry::default())),
        }
    }
}

#[pymethods]
impl PyRegistry {
    #[new]
    fn new() -> Self {
        Self::default()
    }
}

pub fn get_registry(registry: Option<&PyRegistry>) -> &PyRegistry {
    registry.unwrap_or(&DEFAULT_REGISTRY)
}

pub struct MetricOpts<'a> {
    pub name: &'a str,
    pub help: &'a str,
    pub namespace: &'a str,
    pub subsystem: &'a str,
    pub unit: &'a str,
    pub registry: Option<&'a PyRegistry>,
}

pub fn register_metric(
    opts: &MetricOpts<'_>,
    metric: impl prometheus_client::registry::Metric,
) -> PyResult<()> {
    let registry = get_registry(opts.registry);
    let full_name = crate::labels::build_full_name(opts.namespace, opts.subsystem, opts.name);
    let mut reg = registry
        .inner
        .write()
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
    if opts.unit.is_empty() {
        reg.register(full_name, opts.help, metric);
    } else {
        reg.register_with_unit(
            full_name,
            opts.help,
            prometheus_client::registry::Unit::Other(opts.unit.to_owned()),
            metric,
        );
    }
    drop(reg);
    Ok(())
}

#[pyfunction]
#[pyo3(signature = (registry=None))]
pub fn generate_latest(py: Python<'_>, registry: Option<&PyRegistry>) -> PyResult<Vec<u8>> {
    let registry = get_registry(registry);
    let inner = Arc::clone(&registry.inner);
    let result: Result<Vec<u8>, String> = py.detach(|| {
        let reg = inner.read().map_err(|e| e.to_string())?;
        let mut buf = String::new();
        encode(&mut buf, &reg).map_err(|e| e.to_string())?;
        Ok(buf.into_bytes())
    });
    result.map_err(pyo3::exceptions::PyRuntimeError::new_err)
}
