// ! Python context module to store python modules

use pyo3::prelude::*;
use pyo3::types::PyString;
use std::collections::HashMap;
use std::sync::{OnceLock, RwLock};

// Global context instance
static PYTHON_CONTEXT: OnceLock<PythonContext> = OnceLock::new();

// Context for python app
pub struct PythonContext {
    // store python objects thread-safe ref
    pub store: RwLock<HashMap<String, Py<PyAny>>>,
}

// Context Implementation
impl PythonContext {
    pub fn new() -> Self {
        Self {
            store: RwLock::new(HashMap::new()),
        }
    }

    // Global Context instance
    pub fn global() -> &'static PythonContext {
        PYTHON_CONTEXT.get_or_init(|| PythonContext {
            store: RwLock::new(HashMap::new()),
        })
    }

    /// Store a Python object (Module, Classe, etc.)
    pub fn register(&self, key: &str, value: Py<PyAny>) {
        let mut store = self.store.write().unwrap();
        store.insert(key.to_string(), value);
    }

    /// Get Stored object
    pub fn get(&self, key: &str) -> Option<Py<PyAny>> {
        let store = self.store.read().unwrap();
        store.get(key).cloned()
    }

    /// Get or insert a cached Python string object
    pub fn get_or_insert_str(&self, key: &str, py: Python<'_>) -> Py<PyAny> {
        if let Some(v) = self.get(key) {
            return v;
        }
        let s = PyString::new(py, key).into_any().unbind();
        self.register(key, s.clone_ref(py));
        s
    }

    /// Get or insert cached "HTTP_*" header key for WSGI environ.
    pub fn get_or_insert_http_header_key(&self, name_lc: &[u8], py: Python<'_>) -> Py<PyAny> {
        // Cache key in store to avoid repeated Python string allocations.
        let mut key_buf = Vec::with_capacity(5 + name_lc.len());
        key_buf.extend_from_slice(b"HTTP_");
        for b in name_lc {
            if *b == b'-' {
                key_buf.push(b'_');
            } else {
                key_buf.push(b.to_ascii_uppercase());
            }
        }

        // Safe because we only have ASCII bytes.
        let key_str = unsafe { String::from_utf8_unchecked(key_buf) };
        let cache_key = format!("wsgi.http_key.{}", key_str);

        if let Some(v) = self.get(&cache_key) {
            return v;
        }

        let s = PyString::new(py, &key_str).into_any().unbind();
        self.register(&cache_key, s.clone_ref(py));
        s
    }
}
