use pyo3::prelude::*;

pub mod types;
pub mod nvidia;
pub mod detect;

pub const LLMHOSTS_GPU_VERSION: &str = env!("CARGO_PKG_VERSION");

#[pyfunction]
fn detect_gpus() -> PyResult<String> {
    let topology = detect::detect_all();
    serde_json::to_string(&topology)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
}

/// Register llmhosts_gpu functions into a parent module.
pub fn register_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add("LLMHOSTS_GPU_VERSION", LLMHOSTS_GPU_VERSION)?;
    m.add_function(wrap_pyfunction!(detect_gpus, m)?)?;
    Ok(())
}

#[pymodule]
fn _llmhosts_gpu(m: &Bound<'_, PyModule>) -> PyResult<()> {
    register_module(m)
}
