//! WSGI-specific context (caches + environ pool)

use crate::errors::FalcornError;
use crate::python::py_context::PythonContext;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyMapping, PyString};
use std::collections::HashMap;
use std::sync::{OnceLock, RwLock};

static WSGI_CONTEXT: OnceLock<WsgiContext> = OnceLock::new();

pub struct WsgiContext {
    // cached Python strings (keys)
    key_cache: RwLock<HashMap<String, Py<PyAny>>>,
    // cached HTTP_* header keys
    http_key_cache: RwLock<HashMap<Vec<u8>, Py<PyAny>>>,
    // base environ
    base_environ: RwLock<Option<Py<PyDict>>>,
    // pool of reusable environ dicts
    environ_pool: RwLock<Vec<Py<PyDict>>>,
}

impl WsgiContext {
    pub fn global() -> &'static WsgiContext {
        WSGI_CONTEXT.get_or_init(|| WsgiContext {
            key_cache: RwLock::new(HashMap::new()),
            http_key_cache: RwLock::new(HashMap::new()),
            base_environ: RwLock::new(None),
            environ_pool: RwLock::new(Vec::new()),
        })
    }

    pub fn get_or_insert_str(&self, key: &str, py: Python<'_>) -> Py<PyAny> {
        if let Some(v) = self.key_cache.read().unwrap().get(key).cloned() {
            return v;
        }
        let s = PyString::new(py, key).into_any().unbind();
        self.key_cache
            .write()
            .unwrap()
            .insert(key.to_string(), s.clone_ref(py));
        s
    }

    pub fn get_or_insert_http_header_key(&self, name_lc: &[u8], py: Python<'_>) -> Py<PyAny> {
        if let Some(v) = self.http_key_cache.read().unwrap().get(name_lc).cloned() {
            return v;
        }

        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());
            }
        }

        let key_str = unsafe { String::from_utf8_unchecked(key_buf) };
        let s = PyString::new(py, &key_str).into_any().unbind();
        self.http_key_cache
            .write()
            .unwrap()
            .insert(name_lc.to_vec(), s.clone_ref(py));
        s
    }

    pub fn get_or_insert_base_environ(&self, py: Python<'_>) -> Result<Py<PyDict>, FalcornError> {
        if let Some(env) = self.base_environ.read().unwrap().clone() {
            return Ok(env);
        }

        let dict = PyDict::new(py);

        // WSGI VERSION
        dict.set_item(self.get_or_insert_str("wsgi.version", py), (1, 0))
            .map_err(|_| FalcornError::runtime("Failed to set wsgi.version"))?;

        // WSGI URL SCHEME
        dict.set_item(self.get_or_insert_str("wsgi.url_scheme", py), "http")
            .map_err(|_| FalcornError::runtime("Failed to set wsgi.url_scheme"))?;

        // WSGI MULTITHREAD
        dict.set_item(self.get_or_insert_str("wsgi.multithread", py), true)
            .map_err(|_| FalcornError::runtime("Failed to set wsgi.multithread"))?;

        // WSGI MULTIPROCESS
        dict.set_item(self.get_or_insert_str("wsgi.multiprocess", py), true)
            .map_err(|_| FalcornError::runtime("Failed to set wsgi.multiprocess"))?;

        // WSGI RUNONCE
        dict.set_item(self.get_or_insert_str("wsgi.run_once", py), false)
            .map_err(|_| FalcornError::runtime("Failed to set wsgi.run_once"))?;

        // SCRIPT NAME
        dict.set_item(self.get_or_insert_str("SCRIPT_NAME", py), "")
            .map_err(|_| FalcornError::runtime("Failed to set SCRIPT_NAME"))?;

        // QUERY STRING
        dict.set_item(self.get_or_insert_str("QUERY_STRING", py), "")
            .map_err(|_| FalcornError::runtime("Failed to set QUERY_STRING"))?;

        // STDERR
        if let Some(stderr) = PythonContext::global().get("sys.stderr") {
            dict.set_item(self.get_or_insert_str("wsgi.errors", py), stderr.bind(py))
                .map_err(|_| FalcornError::runtime("Failed to set wsgi.errors"))?;
        }

        let env = dict.unbind();
        *self.base_environ.write().unwrap() = Some(env.clone_ref(py));
        Ok(env)
    }

    pub fn take_environ(&self, py: Python<'_>) -> Result<Py<PyDict>, FalcornError> {
        if let Some(env) = self.environ_pool.write().unwrap().pop() {
            return Ok(env);
        }
        let base = self.get_or_insert_base_environ(py)?;
        let env_any = base.bind(py).call_method0("copy")?;
        let dict = env_any
            .cast_into::<PyDict>()
            .map_err(|_| FalcornError::runtime("Failed to clone base environ"))?;
        Ok(dict.unbind())
    }

    pub fn recycle_environ(&self, py: Python<'_>, env: Py<PyDict>) -> Result<(), FalcornError> {
        let base = self.get_or_insert_base_environ(py)?;
        let env_b = env.bind(py);
        env_b.clear();
        env_b
            .update(
                base.bind(py)
                    .cast::<PyMapping>()
                    .map_err(|_| FalcornError::runtime("Failed to reset environ"))?,
            )
            .map_err(|_| FalcornError::runtime("Failed to reset environ"))?;

        const MAX_POOL: usize = 1024;
        let mut pool = self.environ_pool.write().unwrap();
        if pool.len() < MAX_POOL {
            pool.push(env);
        }
        Ok(())
    }
}
