//! Python bindings for State types (SharedState, PlainState)

use pyo3::prelude::*;
use pyo3::exceptions::{PyKeyError, PyRuntimeError, PyValueError};
use pyo3::types::PyDict;
use serde_json::Value;

use flowgentra_ai::core::state::{PlainState, SharedState};
use flowgentra_ai::core::state::State;

use crate::{json_to_py, py_to_json};

// ─── PySharedState ──────────────────────────────────────────────────────────

/// Thread-safe shared state (default state type).
///
/// A dict-like object backed by Arc<RwLock<PlainState>> in Rust.
/// Supports get/set/delete operations with automatic JSON conversion.
///
/// Example:
///     state = SharedState()
///     state["input"] = "hello"
///     print(state["input"])  # "hello"
///     print(state.to_dict()) # {"input": "hello"}
#[pyclass(name = "SharedState")]
#[derive(Clone)]
pub struct PySharedState {
    pub(crate) inner: SharedState,
}

#[pymethods]
impl PySharedState {
    /// Create a new empty SharedState
    #[new]
    #[pyo3(signature = (initial=None))]
    fn new(initial: Option<&Bound<'_, PyDict>>) -> PyResult<Self> {
        let shared = SharedState::new(PlainState::new());
        if let Some(dict) = initial {
            for (k, v) in dict.iter() {
                let key: String = k.extract()?;
                let val = py_to_json(&v)?;
                shared.set(key, val);
            }
        }
        Ok(PySharedState { inner: shared })
    }

    /// Get a value by key
    fn get(&self, py: Python<'_>, key: &str) -> PyResult<PyObject> {
        match self.inner.get(key) {
            Some(val) => json_to_py(py, &val),
            None => Ok(py.None()),
        }
    }

    /// Set a value by key
    fn set(&self, key: &str, value: &Bound<'_, PyAny>) -> PyResult<()> {
        let val = py_to_json(value)?;
        self.inner.set(key, val);
        Ok(())
    }

    /// Check if a key exists
    fn contains_key(&self, key: &str) -> bool {
        self.inner.contains_key(key)
    }

    /// Remove a key and return its value
    fn remove(&self, py: Python<'_>, key: &str) -> PyResult<PyObject> {
        match self.inner.remove(key) {
            Some(val) => json_to_py(py, &val),
            None => Ok(py.None()),
        }
    }

    /// Get all keys
    fn keys(&self) -> Vec<String> {
        self.inner.keys().collect()
    }

    /// Convert state to a Python dict
    fn to_dict(&self, py: Python<'_>) -> PyResult<PyObject> {
        let val = self.inner.to_value();
        json_to_py(py, &val)
    }

    /// Convert state to a JSON string
    fn to_json(&self) -> PyResult<String> {
        self.inner
            .to_json_string()
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))
    }

    /// Create from a Python dict
    #[staticmethod]
    fn from_dict(dict: &Bound<'_, PyDict>) -> PyResult<Self> {
        let shared = SharedState::new(PlainState::new());
        for (k, v) in dict.iter() {
            let key: String = k.extract()?;
            let val = py_to_json(&v)?;
            shared.set(key, val);
        }
        Ok(PySharedState { inner: shared })
    }

    /// Create from a JSON string
    #[staticmethod]
    fn from_json(json_str: &str) -> PyResult<Self> {
        let value: Value = serde_json::from_str(json_str)
            .map_err(|e| PyValueError::new_err(format!("Invalid JSON: {}", e)))?;
        let plain = PlainState::from_json(value)
            .map_err(|e| PyValueError::new_err(format!("{}", e)))?;
        Ok(PySharedState {
            inner: SharedState::new(plain),
        })
    }

    /// Get a string value or None
    fn get_string(&self, key: &str) -> Option<String> {
        self.inner.get_string(key)
    }

    /// Deep clone (creates independent copy, not shared reference)
    fn deep_clone(&self) -> Self {
        PySharedState {
            inner: self.inner.deep_clone(),
        }
    }

    // ── Python dunder methods ──

    fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<PyObject> {
        match self.inner.get(key) {
            Some(val) => json_to_py(py, &val),
            None => Err(PyKeyError::new_err(key.to_string())),
        }
    }

    fn __setitem__(&self, key: &str, value: &Bound<'_, PyAny>) -> PyResult<()> {
        self.set(key, value)
    }

    fn __delitem__(&self, key: &str) -> PyResult<()> {
        if self.inner.contains_key(key) {
            self.inner.remove(key);
            Ok(())
        } else {
            Err(PyKeyError::new_err(key.to_string()))
        }
    }

    fn __contains__(&self, key: &str) -> bool {
        self.inner.contains_key(key)
    }

    fn __len__(&self) -> usize {
        self.inner.keys().count()
    }

    fn __repr__(&self) -> String {
        format!(
            "SharedState({})",
            self.inner
                .to_json_string()
                .unwrap_or_else(|_| "{}".to_string())
        )
    }

    fn __str__(&self) -> String {
        self.inner
            .to_json_string()
            .unwrap_or_else(|_| "{}".to_string())
    }
}

// ─── PyPlainState ───────────────────────────────────────────────────────────

/// Owned state (not thread-safe, for direct manipulation).
///
/// Use SharedState for most cases. PlainState is for advanced use
/// when you need direct, non-shared ownership of state data.
#[pyclass(name = "PlainState")]
#[derive(Clone)]
pub struct PyPlainState {
    pub(crate) inner: PlainState,
}

#[pymethods]
impl PyPlainState {
    #[new]
    #[pyo3(signature = (initial=None))]
    fn new(initial: Option<&Bound<'_, PyDict>>) -> PyResult<Self> {
        let plain = PlainState::new();
        if let Some(dict) = initial {
            for (k, v) in dict.iter() {
                let key: String = k.extract()?;
                let val = py_to_json(&v)?;
                plain.set(key, val);
            }
        }
        Ok(PyPlainState { inner: plain })
    }

    fn get(&self, py: Python<'_>, key: &str) -> PyResult<PyObject> {
        match self.inner.get(key) {
            Some(val) => json_to_py(py, val),
            None => Ok(py.None()),
        }
    }

    fn set(&mut self, key: &str, value: &Bound<'_, PyAny>) -> PyResult<()> {
        let val = py_to_json(value)?;
        self.inner.set(key, val);
        Ok(())
    }

    fn contains_key(&self, key: &str) -> bool {
        self.inner.contains_key(key)
    }

    fn remove(&mut self, py: Python<'_>, key: &str) -> PyResult<PyObject> {
        match self.inner.remove(key) {
            Some(val) => json_to_py(py, &val),
            None => Ok(py.None()),
        }
    }

    fn keys(&self) -> Vec<String> {
        self.inner.keys().cloned().collect()
    }

    fn to_dict(&self, py: Python<'_>) -> PyResult<PyObject> {
        let val = self.inner.to_value();
        json_to_py(py, &val)
    }

    fn to_json(&self) -> PyResult<String> {
        self.inner
            .to_json_string()
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))
    }

    #[staticmethod]
    fn from_dict(dict: &Bound<'_, PyDict>) -> PyResult<Self> {
        let plain = PlainState::new();
        for (k, v) in dict.iter() {
            let key: String = k.extract()?;
            let val = py_to_json(&v)?;
            plain.set(key, val);
        }
        Ok(PyPlainState { inner: plain })
    }

    #[staticmethod]
    fn from_json(json_str: &str) -> PyResult<Self> {
        let value: Value = serde_json::from_str(json_str)
            .map_err(|e| PyValueError::new_err(format!("Invalid JSON: {}", e)))?;
        let plain = PlainState::from_json(value)
            .map_err(|e| PyValueError::new_err(format!("{}", e)))?;
        Ok(PyPlainState { inner: plain })
    }

    fn get_string(&self, key: &str) -> Option<String> {
        self.inner.get_string(key)
    }

    fn get_int(&self, key: &str) -> Option<i64> {
        self.inner.get_int(key)
    }

    fn get_float(&self, key: &str) -> Option<f64> {
        self.inner.get_float(key)
    }

    fn get_bool(&self, key: &str) -> Option<bool> {
        self.inner.get_bool(key)
    }

    fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<PyObject> {
        match self.inner.get(key) {
            Some(val) => json_to_py(py, val),
            None => Err(PyKeyError::new_err(key.to_string())),
        }
    }

    fn __setitem__(&mut self, key: &str, value: &Bound<'_, PyAny>) -> PyResult<()> {
        self.set(key, value)
    }

    fn __contains__(&self, key: &str) -> bool {
        self.inner.contains_key(key)
    }

    fn __len__(&self) -> usize {
        self.inner.keys().count()
    }

    fn __repr__(&self) -> String {
        format!(
            "PlainState({})",
            self.inner
                .to_json_string()
                .unwrap_or_else(|_| "{}".to_string())
        )
    }
}
