// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 LLMHosts.com Contributors

use pyo3::prelude::*;
use std::path::PathBuf;
use std::sync::{Arc, RwLock};

pub mod gguf;
pub mod quant;
pub mod tensor;
pub mod rope;
pub mod attention;
pub mod ffn;
pub mod model;
pub mod generate;
pub mod kv_cache;
pub mod compute;
pub mod memory;
#[cfg(any(feature = "cuda", feature = "wgpu"))]
pub mod kernels;
pub mod test_utils; // NOT cfg(test) — needed for generate_test_model pyfunction

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

/// Python-exposed inference engine.
///
/// # Concurrency
///
/// `model` is held behind `Arc<RwLock<Option<...>>>`:
/// - **Write lock** — acquired only during `load_model()` and `unload()` (rare).
/// - **Read lock** — acquired during `generate_tokens()`, `memory_usage()`,
///   `backend_name()`, and `is_loaded()`. Multiple concurrent read locks do
///   not block each other, enabling truly parallel inference requests.
///
/// Each `generate_tokens()` call creates its own [`InferenceSession`] with an
/// independent KV cache, so N concurrent Python async tasks can run forward()
/// simultaneously without any Rust-level serialization.
#[pyclass]
pub struct LLMHostsEngine {
    model: Arc<RwLock<Option<model::transformer::LLMHostsModel>>>,
    model_path: Option<String>,
}

#[pymethods]
impl LLMHostsEngine {
    #[new]
    fn new() -> Self {
        Self {
            model: Arc::new(RwLock::new(None)),
            model_path: None,
        }
    }

    /// Load a GGUF model. Returns JSON config info.
    /// Acquires the write lock — no concurrent inference during load.
    fn load_model(&mut self, py: Python<'_>, path: &str) -> PyResult<String> {
        let p = PathBuf::from(path);
        if !p.exists() {
            return Err(pyo3::exceptions::PyFileNotFoundError::new_err(
                path.to_string(),
            ));
        }

        // Release the GIL during GGUF parsing + mmap — heavy I/O must not stall other threads.
        let load_result = py.detach(|| -> Result<_, String> {
            let header = gguf::parse_header(&p).map_err(|e| e.to_string())?;
            let config = model::config::LLMHostsModelConfig::from_gguf(&header)
                .map_err(|e| e.to_string())?;
            let mapped = gguf::MappedWeights::open(&p, header.data_offset)
                .map_err(|e| e.to_string())?;
            let weights =
                model::transformer::LLMHostsModelWeights::from_gguf(&config, &header, &mapped)
                    .map_err(|e| e.to_string())?;
            Ok((config, weights))
        });
        let (config, weights) =
            load_result.map_err(pyo3::exceptions::PyRuntimeError::new_err)?;

        let hosts_model = model::transformer::LLMHostsModel::new(config.clone(), weights);

        *self.model.write().unwrap() = Some(hosts_model);
        self.model_path = Some(path.to_string());

        Ok(serde_json::to_string(&serde_json::json!({
            "architecture": config.architecture,
            "hidden_size": config.hidden_size,
            "num_layers": config.num_layers,
            "num_attention_heads": config.num_attention_heads,
            "num_kv_heads": config.num_kv_heads,
            "vocab_size": config.vocab_size,
            "max_seq_len": config.max_seq_len,
            "eos_token_id": config.eos_token_id,
            "bos_token_id": config.bos_token_id,
            "rope_freq_base": config.rope_freq_base,
        }))
        .unwrap())
    }

    /// Generate tokens from token IDs. Returns JSON with generated token IDs.
    ///
    /// Acquires a **read lock** — concurrent calls proceed without blocking each
    /// other. Each call creates its own [`InferenceSession`] (KV cache) so
    /// generation state is never shared between requests.
    fn generate_tokens(
        &self,
        py: Python<'_>,
        token_ids: Vec<u32>,
        max_tokens: usize,
        temperature: f32,
        top_p: f32,
    ) -> PyResult<String> {
        let guard = self.model.read().unwrap();
        let model = guard
            .as_ref()
            .ok_or_else(|| pyo3::exceptions::PyRuntimeError::new_err("No model loaded"))?;

        // Auto-include EOS token from model config as a stop token
        let mut stop_tokens = vec![];
        if let Some(eos) = model.config.eos_token_id {
            stop_tokens.push(eos);
        }
        let params = generate::GenerateParams {
            max_tokens,
            temperature,
            top_p,
            stop_tokens,
        };

        // Release the GIL during generation — does not stall the FastAPI event loop.
        // generate() creates a fresh InferenceSession internally — fully concurrent.
        let tokens = py.detach(|| generate::generate(model, &token_ids, &params));
        Ok(serde_json::to_string(&tokens).unwrap())
    }

    #[getter]
    fn is_loaded(&self) -> bool {
        self.model.read().unwrap().is_some()
    }

    #[getter]
    fn model_path(&self) -> Option<String> {
        self.model_path.clone()
    }

    /// Returns JSON with current memory usage: allocated, total, free, weights, kv_cache (all in bytes).
    fn memory_usage(&self) -> PyResult<String> {
        let guard = self.model.read().unwrap();
        let model = guard
            .as_ref()
            .ok_or_else(|| pyo3::exceptions::PyRuntimeError::new_err("No model loaded"))?;
        let usage = model.memory_usage();
        Ok(serde_json::to_string(&serde_json::json!({
            "allocated": usage.allocated,
            "total": usage.total,
            "free": usage.free,
            "weights": usage.weights,
            "kv_cache": usage.kv_cache,
        }))
        .unwrap())
    }

    /// Unload the model and free memory. Acquires write lock.
    fn unload(&mut self) {
        *self.model.write().unwrap() = None;
        self.model_path = None;
    }

    /// Return the name of the active compute backend (e.g. "cpu", "cuda:0", "wgpu:Metal").
    fn backend_name(&self) -> PyResult<String> {
        let guard = self.model.read().unwrap();
        let model = guard
            .as_ref()
            .ok_or_else(|| pyo3::exceptions::PyRuntimeError::new_err("No model loaded"))?;
        Ok(model.backend_name().to_string())
    }

    /// Return list of compute backends compiled into this build.
    #[staticmethod]
    fn available_backends() -> Vec<String> {
        compute::auto::available_backends()
            .into_iter()
            .map(|s| s.to_string())
            .collect()
    }
}

/// Generate a synthetic test GGUF model for testing.
#[pyfunction]
fn generate_test_model(output_dir: &str) -> PyResult<String> {
    let config = test_utils::TestModelConfig::default();
    let path = test_utils::generate_test_gguf(std::path::Path::new(output_dir), &config);
    Ok(path.to_string_lossy().to_string())
}

/// Register llmhosts_core classes/functions into a parent module.
pub fn register_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add("LLMHOSTS_CORE_VERSION", LLMHOSTS_CORE_VERSION)?;
    m.add_class::<LLMHostsEngine>()?;
    m.add_function(wrap_pyfunction!(generate_test_model, m)?)?;
    Ok(())
}

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