use async_trait::async_trait;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyModule};
use pyo3_async_runtimes::tokio::future_into_py;
use runtime::core::agent::{Agent, AgentContext};
use runtime::core::memory::{Memory, MemoryEntry, MemoryQuery};
use runtime::core::mesh::Mesh;
use runtime::core::message::Message;
use runtime::llm::LLMConfig;
use runtime::memory::InMemoryBackend;
use runtime::memory::RedisBackend;
use runtime::LocalMesh;
use serde_json::Value;
use distributed::registry::{AgentMetadata, HealthStatus, InMemoryRegistry, Registry};
use distributed::DistributedMesh;
use runtime::logging::{init_logging, LoggingConfig, LoggingGuards};
use std::path::PathBuf;
use std::sync::Arc;
use tokio::runtime::Runtime;

// Create a static Tokio runtime for async execution
lazy_static::lazy_static! {
    static ref RUNTIME: Runtime = Runtime::new().unwrap();
}

/// Python wrapper for LoggingConfig
#[pyclass(subclass)]
#[derive(Clone)]
struct PyLoggingConfig {
    inner: LoggingConfig,
}

#[pymethods]
impl PyLoggingConfig {
    #[new]
    fn new(log_level: String, log_file_path: Option<String>, json_output: bool) -> Self {
        PyLoggingConfig {
            inner: LoggingConfig {
                log_level,
                log_file_path: log_file_path.map(PathBuf::from),
                json_output,
            },
        }
    }
}

/// Python wrapper for LoggingGuards to keep file appenders alive
#[pyclass(subclass)]
struct PyLoggingHandle {
    _inner: LoggingGuards,
}

/// Initialize logging from Python
#[pyfunction]
fn init_logging_py(config: PyLoggingConfig) -> PyLoggingHandle {
    let guards = init_logging(&config.inner);
    PyLoggingHandle { _inner: guards }
}

/// Python wrapper for LocalMesh
#[pyclass(subclass)]
struct PyLocalMesh {
    inner: Arc<LocalMesh>,
}

#[pymethods]
impl PyLocalMesh {
    #[new]
    fn new(name: String) -> Self {
        PyLocalMesh {
            inner: Arc::new(LocalMesh::new(name)),
        }
    }

    fn start(&self) -> PyResult<()> {
        let mesh = self.inner.clone();
        RUNTIME.block_on(async move {
            mesh.start()
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
        })
    }

    fn add_agent(&self, agent: Py<PyAgent>) -> PyResult<()> {
        let mesh = self.inner.clone();
        let agent_wrapper = Box::new(PythonAgentWrapper { agent });

        RUNTIME.block_on(async move {
            mesh.add_agent(agent_wrapper)
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
        })
    }

    fn send_to(&self, target: String, payload: String) -> PyResult<()> {
        let mesh = self.inner.clone();
        let msg = Message::new("system", payload.into_bytes(), target.clone());

        RUNTIME.block_on(async move {
            mesh.send(msg, &target)
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
        })
    }
}

/// Python wrapper for DistributedMesh
#[pyclass(subclass)]
struct PyDistributedMesh {
    inner: Arc<distributed::DistributedMesh>,
}

#[pymethods]
impl PyDistributedMesh {
    #[new]
    fn new(name: String, port: u16) -> Self {
        PyDistributedMesh {
            inner: Arc::new(distributed::DistributedMesh::new(name, port)),
        }
    }

    fn start(&self) -> PyResult<()> {
        let mesh = self.inner.clone();
        RUNTIME.block_on(async move {
            mesh.start()
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
        })
    }

    fn stop(&self) -> PyResult<()> {
        let mesh = self.inner.clone();
        RUNTIME.block_on(async move {
            mesh.stop()
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
        })
    }

    fn add_agent(&self, agent: Py<PyAgent>) -> PyResult<()> {
        let mesh = self.inner.clone();
        let agent_wrapper = Box::new(PythonAgentWrapper { agent });

        RUNTIME.block_on(async move {
            mesh.add_agent(agent_wrapper)
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
        })
    }

    fn connect_peer(&self, agent_name: String, url: String) -> PyResult<()> {
        let mesh = self.inner.clone();
        RUNTIME.block_on(async move {
            mesh.connect_peer(agent_name, url).await;
            Ok(())
        })
    }

    fn send_to(&self, target: String, payload: String) -> PyResult<()> {
        let mesh = self.inner.clone();
        let msg = Message::new("system", payload.into_bytes(), target.clone());

        RUNTIME.block_on(async move {
            mesh.send(msg, &target)
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
        })
    }

    #[staticmethod]
    fn with_registry(name: String, port: u16, registry: &PyInMemoryRegistry) -> Self {
        let reg = registry.inner.clone();
        PyDistributedMesh {
            inner: Arc::new(distributed::DistributedMesh::with_registry(name, port, reg)),
        }
    }
}

/// Python wrapper for AgentMetadata
#[pyclass]
#[derive(Clone)]
struct PyAgentMetadata {
    inner: AgentMetadata,
}

#[pymethods]
impl PyAgentMetadata {
    #[new]
    fn new(id: String, name: String, address: String) -> Self {
        PyAgentMetadata {
            inner: AgentMetadata {
                id,
                name,
                address,
                capabilities: vec![],
                health_status: HealthStatus::Healthy,
                last_heartbeat: 0,
            },
        }
    }

    #[getter]
    fn id(&self) -> String {
        self.inner.id.clone()
    }

    #[getter]
    fn name(&self) -> String {
        self.inner.name.clone()
    }

    #[getter]
    fn address(&self) -> String {
        self.inner.address.clone()
    }
}

/// Python wrapper for InMemoryRegistry
#[pyclass]
#[derive(Clone)]
struct PyInMemoryRegistry {
    inner: Arc<InMemoryRegistry>,
}

#[pymethods]
impl PyInMemoryRegistry {
    #[new]
    fn new() -> Self {
        PyInMemoryRegistry {
            inner: Arc::new(InMemoryRegistry::new()),
        }
    }

    fn register(&self, metadata: PyAgentMetadata) -> PyResult<()> {
        let registry = self.inner.clone();
        let meta = metadata.inner;
        RUNTIME.block_on(async move {
            registry.register(meta).await.map_err(|e| {
                PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string())
            })
        })
    }

    fn get_agent(&self, agent_id: String) -> PyResult<Option<PyAgentMetadata>> {
        let registry = self.inner.clone();
        RUNTIME.block_on(async move {
            match registry.get_agent(&agent_id).await {
                Ok(Some(meta)) => Ok(Some(PyAgentMetadata { inner: meta })),
                Ok(None) => Ok(None),
                Err(e) => Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                    e.to_string(),
                )),
            }
        })
    }
}

/// Base class for Python Agents
#[derive(Clone)]
#[pyclass(subclass)]
struct PyAgent {}

#[pymethods]
impl PyAgent {
    #[new]
    #[pyo3(signature = (_name = String::new()))]
    fn new(_name: String) -> Self {
        PyAgent {}
    }

    fn name(&self) -> String {
        "agent".to_string() // Default name, should be overridden by subclasses
    }

    fn act(self_: Py<Self>, action_name: String, inputs: Py<PyDict>) -> PyResult<Py<PyAny>> {
        Python::with_gil(|py| {
            let bound = self_.bind(py);
            if let Ok(tool_invoker) = bound.getattr("tool_invoker") {
                let json = py.import("json")?;
                let inputs_str: String = json.call_method1("dumps", (inputs,))?.extract()?;

                let result = tool_invoker.call_method1("invoke", (action_name, inputs_str))?;
                Ok(result.into())
            } else {
                Ok(py.None())
            }
        })
    }
}

/// Wrapper to adapt Python Agents to Rust Agent trait
struct PythonAgentWrapper {
    agent: Py<PyAgent>,
}

#[async_trait::async_trait]
impl Agent for PythonAgentWrapper {
    fn name(&self) -> String {
        Python::with_gil(|py| {
            let agent = self.agent.bind(py);
            // Try to call the name() method on the Python object
            if let Ok(name) = agent.call_method0("name") {
                name.extract().unwrap_or_else(|_| "unknown".to_string())
            } else {
                "unknown".to_string()
            }
        })
    }

    async fn on_start(&mut self, _ctx: &mut AgentContext) -> runtime::core::error::Result<()> {
        // Call on_start on the Python object if it exists
        Python::with_gil(|py| {
            let agent = self.agent.bind(py);
            if agent.hasattr("on_start")? {
                // TODO: Pass context to Python
                // We need to wrap AgentContext into a Python object to pass it here
                agent.call_method1("on_start", (py.None(),))?;
            }
            Ok(())
        })
        .map_err(|e: PyErr| runtime::core::error::Error::MeshError(e.to_string()))
    }

    async fn on_message(
        &mut self,
        msg: Message,
        _ctx: &mut AgentContext,
    ) -> runtime::core::error::Result<()> {
        let result = Python::with_gil(|py| {
            let agent = self.agent.bind(py);
            if agent.hasattr("on_message")? {
                // Pass proper message object and context
                // Convert bytes to string
                let payload_str = String::from_utf8_lossy(&msg.payload).to_string();
                let result = agent.call_method1("on_message", (payload_str, py.None()))?;

                // Check if the result is awaitable (coroutine)
                let asyncio = py.import("asyncio")?;
                let is_coroutine = asyncio
                    .call_method1("iscoroutine", (result.clone(),))?
                    .extract::<bool>()?;

                if is_coroutine {
                    // Create a new event loop for this thread
                    let new_loop = asyncio.call_method0("new_event_loop")?;

                    // Get the current event loop (if any) to restore later
                    let old_loop = asyncio.call_method0("get_event_loop").ok();

                    // Temporarily set the new loop as the current loop for this thread
                    asyncio.call_method1("set_event_loop", (new_loop.clone(),))?;

                    // Run the coroutine on this thread-local event loop
                    let coro_result = new_loop.call_method1("run_until_complete", (result,));

                    // Restore the old event loop (or set to None)
                    if let Some(old) = old_loop {
                        asyncio.call_method1("set_event_loop", (old,))?;
                    } else {
                        asyncio.call_method1("set_event_loop", (py.None(),))?;
                    }

                    // Close the loop to clean up
                    new_loop.call_method0("close")?;

                    // Check if run_until_complete succeeded
                    coro_result?;

                    // Return None since we already executed the coroutine
                    Ok::<Option<Py<PyAny>>, PyErr>(None)
                } else {
                    Ok::<Option<Py<PyAny>>, PyErr>(None)
                }
            } else {
                Ok::<Option<Py<PyAny>>, PyErr>(None)
            }
        })
        .map_err(|e: PyErr| runtime::core::error::Error::MeshError(e.to_string()))?;

        Ok(())
    }

    async fn on_stop(&mut self, _ctx: &mut AgentContext) -> runtime::core::error::Result<()> {
        Python::with_gil(|py| {
            let agent = self.agent.bind(py);
            if agent.hasattr("on_stop")? {
                agent.call_method1("on_stop", (py.None(),))?;
            }
            Ok(())
        })
        .map_err(|e: PyErr| runtime::core::error::Error::MeshError(e.to_string()))
    }
}

use runtime::core::action::{ActionInvoker, ActionMetadata, ToolInvoker};

/// Base class for Python Actions
#[pyclass(subclass, name = "_PyAction")]
struct PyAction {
    metadata: ActionMetadata,
}

#[pymethods]
impl PyAction {
    #[new]
    #[pyo3(signature = (name, description, input_schema, output_schema=None))]
    fn new(
        name: String,
        description: String,
        input_schema: String,
        output_schema: Option<String>,
    ) -> PyResult<Self> {
        let input_json: Value = serde_json::from_str(&input_schema)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

        let output_json: Option<Value> = match output_schema {
            Some(s) => Some(
                serde_json::from_str(&s)
                    .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?,
            ),
            None => None,
        };

        Ok(PyAction {
            metadata: ActionMetadata {
                name,
                description,
                input_schema: input_json,
                output_schema: output_json,
            },
        })
    }

    fn metadata(&self) -> PyResult<String> {
        serde_json::to_string(&self.metadata)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))
    }

    fn execute(&self, _context: Py<PyAgentContext>, _inputs: Py<PyDict>) -> PyResult<Py<PyAny>> {
        Python::with_gil(|py| Ok(py.None()))
    }
}

/// Wrapper to adapt Python Actions to Rust ActionInvoker trait
struct PythonActionWrapper {
    action: Py<PyAction>,
    metadata: ActionMetadata,
}

/// Python wrapper for AgentContext
#[pyclass]
#[derive(Clone)]
struct PyAgentContext {
    #[pyo3(get)]
    mesh_name: String,
}

#[pymethods]
impl PyAgentContext {
    #[new]
    fn new(mesh_name: String) -> Self {
        PyAgentContext { mesh_name }
    }
}

#[async_trait::async_trait]
impl ActionInvoker for PythonActionWrapper {
    async fn execute(
        &self,
        ctx: &mut AgentContext,
        inputs: Value,
    ) -> runtime::core::error::Result<Value> {
        let value_obj = Python::with_gil(|py| {
            let action = self.action.bind(py);
            let inputs_str = serde_json::to_string(&inputs).unwrap();
            let json = py.import("json")?;
            let py_inputs = json.call_method1("loads", (inputs_str,))?;
            let py_ctx = PyAgentContext {
                mesh_name: ctx.mesh_name.clone(),
            };
            let py_ctx_obj = py_ctx.into_pyobject(py)?;

            let result = action.call_method1("execute", (py_ctx_obj, py_inputs))?;

            // Check if the result is awaitable (coroutine)
            let asyncio = py.import("asyncio")?;
            let is_coroutine = asyncio
                .call_method1("iscoroutine", (result.clone(),))?
                .extract::<bool>()?;

            if is_coroutine {
                // Create a new event loop for this thread
                let new_loop = asyncio.call_method0("new_event_loop")?;

                // Get the current event loop (if any) to restore later
                let old_loop = asyncio.call_method0("get_event_loop").ok();

                // Temporarily set the new loop as the current loop for this thread
                asyncio.call_method1("set_event_loop", (new_loop.clone(),))?;

                // Run the coroutine on this thread-local event loop
                let coro_result = new_loop.call_method1("run_until_complete", (result,));

                // Restore the old event loop (or set to None)
                if let Some(old) = old_loop {
                    asyncio.call_method1("set_event_loop", (old,))?;
                } else {
                    asyncio.call_method1("set_event_loop", (py.None(),))?;
                }

                // Close the loop to clean up
                new_loop.call_method0("close")?;

                // Check if run_until_complete succeeded
                let final_result = coro_result?;

                Ok(final_result.unbind())
            } else {
                Ok(result.unbind())
            }
        })
        .map_err(|e: PyErr| runtime::core::error::Error::ActionExecutionError(e.to_string()))?;

        // Convert value_obj to Value
        Python::with_gil(|py| {
            let val = value_obj.bind(py);
            if let Ok(result_str) = val.extract::<String>() {
                if let Ok(v) = serde_json::from_str::<Value>(&result_str) {
                    Ok(v)
                } else {
                    Ok(Value::String(result_str))
                }
            } else {
                let result_str = val.to_string();
                Ok(Value::String(result_str))
            }
        })
        .map_err(|e: PyErr| runtime::core::error::Error::ActionExecutionError(e.to_string()))
    }

    fn metadata(&self) -> &ActionMetadata {
        &self.metadata
    }
}

/// Python wrapper for ToolInvoker
#[pyclass]
struct PyToolInvoker {
    inner: Arc<std::sync::Mutex<ToolInvoker>>,
}

#[pymethods]
impl PyToolInvoker {
    #[new]
    fn new() -> Self {
        PyToolInvoker {
            inner: Arc::new(std::sync::Mutex::new(ToolInvoker::new())),
        }
    }

    fn register(&self, action: Py<PyAction>) -> PyResult<()> {
        let (metadata, action_wrapper) = Python::with_gil(|py| {
            let bound_action = action.bind(py);
            let metadata_str: String = bound_action.call_method0("metadata")?.extract()?;
            let metadata: ActionMetadata = serde_json::from_str(&metadata_str)
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

            let wrapper = PythonActionWrapper {
                action: action.clone_ref(py),
                metadata: metadata.clone(),
            };
            Ok::<(ActionMetadata, PythonActionWrapper), PyErr>((metadata, wrapper))
        })?;

        let mut invoker = self.inner.lock().unwrap();
        invoker.register(Box::new(action_wrapper));
        Ok(())
    }

    fn invoke(&self, name: String, inputs: String) -> PyResult<String> {
        let invoker = self.inner.clone();
        let inputs_json: Value = serde_json::from_str(&inputs)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

        // We need a dummy context for now as we don't have one exposed to Python yet in this call
        // In a real app, this would come from the agent
        let mut ctx = AgentContext::new("system".to_string());

        let result = RUNTIME.block_on(async move {
            let invoker = invoker.lock().unwrap();
            invoker.invoke(&name, &mut ctx, inputs_json).await
        });

        match result {
            Ok(val) => serde_json::to_string(&val)
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string())),
            Err(e) => Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                e.to_string(),
            )),
        }
    }
}

/// Python wrapper for LLMConfig
#[pyclass(subclass)]
#[derive(Clone)]
struct PyLlmConfig {
    inner: LLMConfig,
}

#[pymethods]
impl PyLlmConfig {
    #[staticmethod]
    fn builder() -> Self {
        PyLlmConfig {
            inner: LLMConfig::default(),
        }
    }

    fn provider(mut slf: PyRefMut<'_, Self>, provider: String) -> PyResult<PyRefMut<'_, Self>> {
        // In the Rust config, model is "provider::model".
        // We'll store provider temporarily or just prepend it if model is set?
        // Actually, let's just assume model() will be called with the model name,
        // and we combine them. Or we can store them separately in a builder state if we were building from scratch.
        // But here we are wrapping LLMConfig.
        // Let's assume the user calls provider() then model().
        // If model is already set, we might need to handle it.
        // For simplicity, let's just store it in the model string if it's empty, or prepend it.
        // But a cleaner way is to just let the user pass "provider::model" to model().
        // However, the Python example uses .provider("ollama").model("llama3.2:latest").

        // Let's try to handle this by checking if model already has "::".
        let current_model = slf.inner.model.clone();
        if current_model.contains("::") {
            let parts: Vec<&str> = current_model.split("::").collect();
            slf.inner.model = format!("{}::{}", provider, parts[1]);
        } else {
            slf.inner.model = format!("{}::{}", provider, current_model);
        }
        Ok(slf)
    }

    fn model(mut slf: PyRefMut<'_, Self>, model: String) -> PyResult<PyRefMut<'_, Self>> {
        let current_model = slf.inner.model.clone();
        if current_model.contains("::") {
            let parts: Vec<&str> = current_model.split("::").collect();
            slf.inner.model = format!("{}::{}", parts[0], model);
        } else {
            // If provider was set (e.g. "ollama::"), append model.
            // If not, just set model.
            if current_model.ends_with("::") {
                slf.inner.model = format!("{}{}", current_model, model);
            } else {
                slf.inner.model = model;
            }
        }
        Ok(slf)
    }

    fn base_url(mut slf: PyRefMut<'_, Self>, base_url: String) -> PyResult<PyRefMut<'_, Self>> {
        slf.inner.base_url = Some(base_url);
        Ok(slf)
    }

    fn temperature(mut slf: PyRefMut<'_, Self>, temperature: f32) -> PyResult<PyRefMut<'_, Self>> {
        slf.inner.temperature = Some(temperature);
        Ok(slf)
    }

    fn max_tokens(mut slf: PyRefMut<'_, Self>, max_tokens: u32) -> PyResult<PyRefMut<'_, Self>> {
        slf.inner.max_tokens = Some(max_tokens);
        Ok(slf)
    }

    fn build(slf: PyRefMut<'_, Self>) -> PyResult<PyRefMut<'_, Self>> {
        Ok(slf)
    }
}

/// Python wrapper for LlmAgent with builder pattern
#[pyclass]
struct PyLlmAgent {
    inner: Option<Arc<tokio::sync::Mutex<runtime::llm::LlmAgent>>>,
    // Builder state (used before build() is called)
    name: String,
    model: String,
    api_key: Option<String>,
    system_prompt: Option<String>,
    temperature: Option<f32>,
    max_tokens: Option<u32>,
    config: Option<LLMConfig>,
    memory: Option<Py<PyAny>>,
    actions: Vec<Py<PyAction>>,
}

#[pymethods]
impl PyLlmAgent {
    #[new]
    fn new(name: String, model: String) -> Self {
        PyLlmAgent {
            inner: None,
            name,
            model,
            api_key: None,
            system_prompt: None,
            temperature: None,
            max_tokens: None,
            config: None,
            memory: None,
            actions: Vec::new(),
        }
    }

    #[staticmethod]
    fn with_config(name: String, config: PyLlmConfig) -> Self {
        PyLlmAgent {
            inner: None,
            name,
            model: config.inner.model.clone(),
            api_key: None,
            system_prompt: None,
            temperature: None,
            max_tokens: None,

            config: Some(config.inner),
            memory: None,
            actions: Vec::new(),
        }
    }

    fn with_api_key(mut slf: PyRefMut<'_, Self>, api_key: String) -> PyResult<PyRefMut<'_, Self>> {
        slf.api_key = Some(api_key);
        Ok(slf)
    }

    fn with_system_prompt(
        mut slf: PyRefMut<'_, Self>,
        prompt: String,
    ) -> PyResult<PyRefMut<'_, Self>> {
        slf.system_prompt = Some(prompt);
        Ok(slf)
    }

    fn with_temperature(
        mut slf: PyRefMut<'_, Self>,
        temperature: f32,
    ) -> PyResult<PyRefMut<'_, Self>> {
        slf.temperature = Some(temperature);
        Ok(slf)
    }

    fn with_max_tokens(
        mut slf: PyRefMut<'_, Self>,
        max_tokens: u32,
    ) -> PyResult<PyRefMut<'_, Self>> {
        slf.max_tokens = Some(max_tokens);
        Ok(slf)
    }

    fn with_memory(mut slf: PyRefMut<'_, Self>, memory: Py<PyAny>) -> PyResult<PyRefMut<'_, Self>> {
        slf.memory = Some(memory);
        Ok(slf)
    }

    fn add_action(mut slf: PyRefMut<'_, Self>, action: Py<PyAction>) -> PyResult<PyRefMut<'_, Self>> {
        // If agent is already built, register directly
        if let Some(ref inner) = slf.inner {
            let (metadata, action_wrapper) = Python::with_gil(|py| {
                let bound_action = action.bind(py);
                let metadata_str: String = bound_action.call_method0("metadata")?.extract()?;
                let metadata: ActionMetadata = serde_json::from_str(&metadata_str)
                    .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

                let wrapper = PythonActionWrapper {
                    action: action.clone_ref(py),
                    metadata: metadata.clone(),
                };
                Ok::<(ActionMetadata, PythonActionWrapper), PyErr>((metadata, wrapper))
            })?;

            let agent_arc = inner.clone();
            RUNTIME.block_on(async move {
                let mut agent = agent_arc.lock().await;
                if let Some(invoker) = agent.tool_invoker_mut() {
                    invoker.register(Box::new(action_wrapper));
                }
            });
        } else {
            // Otherwise store for later
            slf.actions.push(action);
        }
        Ok(slf)
    }

    fn add_tool(slf: PyRefMut<'_, Self>, action: Py<PyAction>) -> PyResult<PyRefMut<'_, Self>> {
        Self::add_action(slf, action)
    }

    fn build(mut slf: PyRefMut<'_, Self>) -> PyResult<PyRefMut<'_, Self>> {
        // Helper function to convert Python memory to Rust Memory
        fn convert_python_memory(py_memory: &Py<PyAny>) -> PyResult<Arc<dyn runtime::core::memory::Memory>> {
            Python::with_gil(|py| {
                let bound = py_memory.bind(py);

                // Try to extract as PyInMemoryBackend first
                if let Ok(backend) = bound.extract::<PyInMemoryBackend>() {
                    let rust_memory: Arc<dyn runtime::core::memory::Memory> = backend.inner;
                    return Ok(rust_memory);
                }

                // Otherwise, wrap it as a custom Python Memory implementation
                let wrapper = PythonMemoryWrapper {
                    memory: py_memory.clone_ref(py),
                };
                Ok(Arc::new(wrapper) as Arc<dyn runtime::core::memory::Memory>)
            })
        }
        
        // Build the Rust LlmAgent using the builder pattern
        let mut builder = runtime::llm::LlmAgent::builder(&slf.name, &slf.model);

        // If config is present, use it as base (though LlmAgent::builder doesn't take config directly,
        // we might need to use new_with_config if we want to use the full config)
        if let Some(ref config) = slf.config {
            // If we have a full config, we should use LlmAgent::new_with_config
            // But we need to handle the system prompt which might be separate
            let system_prompt = slf
                .system_prompt
                .clone()
                .unwrap_or_else(|| "You are a helpful AI assistant.".to_string());


            // If memory is set, we might need to manually inject it or warn that it's not supported with full config yet
            // unless we update new_with_config to take memory. 
            // I updated new_with_config in Rust to take memory.
            // So I should update the call above.
            
            // Re-do the call with memory
            let memory = if let Some(ref py_mem) = slf.memory {
                Some(convert_python_memory(py_mem)?)
            } else {
                None
            };
            
            let agent =
                runtime::llm::LlmAgent::new_with_config(&slf.name, config.clone(), system_prompt, memory)
                    .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;

            // Register actions
            let mut agent_guard = agent;
            if let Some(invoker) = agent_guard.tool_invoker_mut() {
                for action in &slf.actions {
                    let (metadata, action_wrapper) = Python::with_gil(|py| {
                        let bound_action = action.bind(py);
                        let metadata_str: String = bound_action.call_method0("metadata")?.extract()?;
                        let metadata: ActionMetadata = serde_json::from_str(&metadata_str)
                            .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

                        let wrapper = PythonActionWrapper {
                            action: action.clone_ref(py),
                            metadata: metadata.clone(),
                        };
                        Ok::<(ActionMetadata, PythonActionWrapper), PyErr>((metadata, wrapper))
                    })?;
                    invoker.register(Box::new(action_wrapper));
                }
            }
            
            slf.inner = Some(Arc::new(tokio::sync::Mutex::new(agent_guard)));
            return Ok(slf);
        }

        let mut builder = runtime::llm::LlmAgent::builder(&slf.name, &slf.model);

        if let Some(ref api_key) = slf.api_key {
            builder = builder.with_api_key(api_key);
        }

        if let Some(ref prompt) = slf.system_prompt {
            builder = builder.with_system_prompt(prompt);
        }

        if let Some(temp) = slf.temperature {
            builder = builder.with_temperature(temp);
        }

        if let Some(tokens) = slf.max_tokens {
            builder = builder.with_max_tokens(tokens);
        }

        if let Some(ref py_mem) = slf.memory {
            let memory = convert_python_memory(py_mem)?;
            builder = builder.with_memory(memory);
        }

        let mut agent = builder
            .build()
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;

        // Register actions
        if let Some(invoker) = agent.tool_invoker_mut() {
            for action in &slf.actions {
                let (metadata, action_wrapper) = Python::with_gil(|py| {
                    let bound_action = action.bind(py);
                    let metadata_str: String = bound_action.call_method0("metadata")?.extract()?;
                    let metadata: ActionMetadata = serde_json::from_str(&metadata_str)
                        .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

                    let wrapper = PythonActionWrapper {
                        action: action.clone_ref(py),
                        metadata: metadata.clone(),
                    };
                    Ok::<(ActionMetadata, PythonActionWrapper), PyErr>((metadata, wrapper))
                })?;
                invoker.register(Box::new(action_wrapper));
            }
        }

        slf.inner = Some(Arc::new(tokio::sync::Mutex::new(agent)));
        Ok(slf)
    }

    fn send_message(&self, message: String) -> PyResult<String> {
        let agent = self.inner.as_ref().ok_or_else(|| {
            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                "Agent not built. Call build() first.",
            )
        })?;

        let agent_clone = agent.clone();

        RUNTIME.block_on(async move {
            let mut agent_guard = agent_clone.lock().await;
            let mut ctx = runtime::core::agent::AgentContext::new("python".to_string());

            // Use the new send_message_and_get_response method
            let response = agent_guard
                .send_message_and_get_response(message, &mut ctx)
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;

            Ok(response)
        })
    }

    fn send_message_async<'py>(
        &self,
        py: Python<'py>,
        message: String,
    ) -> PyResult<Bound<'py, PyAny>> {
        let agent = self.inner.as_ref().ok_or_else(|| {
            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                "Agent not built. Call build() first.",
            )
        })?;

        let agent_clone = agent.clone();

        future_into_py(py, async move {
            let mut agent_guard = agent_clone.lock().await;
            let mut ctx = runtime::core::agent::AgentContext::new("python".to_string());

            let response = agent_guard
                .send_message_and_get_response(message, &mut ctx)
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;

            Ok(response)
        })
    }

    fn register_action(&self, action: Py<PyAction>) -> PyResult<()> {
        let agent = self.inner.as_ref().ok_or_else(|| {
            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                "Agent not built. Call build() first.",
            )
        })?;

        let (metadata, action_wrapper) = Python::with_gil(|py| {
            let bound_action = action.bind(py);
            let metadata_str: String = bound_action.call_method0("metadata")?.extract()?;
            let metadata: runtime::core::action::ActionMetadata =
                serde_json::from_str(&metadata_str)
                    .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

            let wrapper = PythonActionWrapper {
                action: action.clone_ref(py),
                metadata: metadata.clone(),
            };
            Ok::<(runtime::core::action::ActionMetadata, PythonActionWrapper), PyErr>((
                metadata, wrapper,
            ))
        })?;

        let mut agent_guard = agent.blocking_lock();
        if let Some(tool_invoker) = agent_guard.tool_invoker_mut() {
            tool_invoker.register(Box::new(action_wrapper));
        }

        Ok(())
    }
   
    fn with_react(&self, config: PyReActConfig) -> PyResult<()> {
        let agent = self.inner.as_ref().ok_or_else(|| {
            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                "Agent not built. Call build() first.",
            )
        })?;
       
        let mut agent_guard = agent.blocking_lock();
        agent_guard.with_react(config.inner);
       
        Ok(())
    }
   
    fn send_message_react(&self, message: String) -> PyResult<PyReActResult> {
        let agent = self.inner.as_ref().ok_or_else(|| {
            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                "Agent not built. Call build() first.",
            )
        })?;
       
        let agent_clone = agent.clone();
       
        RUNTIME.block_on(async move {
            let mut agent_guard = agent_clone.lock().await;
            let mut ctx = runtime::core::agent::AgentContext::new("python".to_string());
           
            let result = agent_guard
                .send_message_react(message, &mut ctx)
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
           
            Ok(PyReActResult { inner: result })
        })
    }
}

/// Python wrapper for MemoryEntry
#[pyclass(subclass)]
#[derive(Clone)]
struct PyMemoryEntry {
    inner: MemoryEntry,
}

#[pymethods]
impl PyMemoryEntry {
    #[new]
    fn new(content: String) -> Self {
        PyMemoryEntry {
            inner: MemoryEntry::new(content),
        }
    }

    #[getter]
    fn id(&self) -> String {
        self.inner.id.clone()
    }

    #[getter]
    fn content(&self) -> String {
        self.inner.content.clone()
    }

    #[getter]
    fn metadata(&self, py: Python) -> PyResult<Py<PyDict>> {
        let dict = PyDict::new(py);
        for (key, value) in &self.inner.metadata {
            let json_str = serde_json::to_string(&value)
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
            let json = py.import("json")?;
            let py_value = json.call_method1("loads", (json_str,))?;
            dict.set_item(key, py_value)?;
        }
        Ok(dict.unbind())
    }

    #[getter]
    fn created_at(&self) -> String {
        self.inner.created_at.to_rfc3339()
    }

    #[getter]
    fn expires_at(&self) -> Option<String> {
        self.inner.expires_at.map(|dt| dt.to_rfc3339())
    }

    fn with_metadata(
        mut slf: PyRefMut<'_, Self>,
        key: String,
        value: Py<PyAny>,
    ) -> PyResult<PyRefMut<'_, Self>> {
        Python::with_gil(|py| {
            let val = value.bind(py);
            let json = py.import("json")?;
            let json_str: String = json.call_method1("dumps", (val,))?.extract()?;
            let json_value: Value = serde_json::from_str(&json_str)
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

            slf.inner = slf.inner.clone().with_metadata(key, json_value);
            Ok(slf)
        })
    }

    fn with_ttl_seconds(mut slf: PyRefMut<'_, Self>, seconds: i64) -> PyResult<PyRefMut<'_, Self>> {
        slf.inner = slf.inner.clone().with_ttl_seconds(seconds);
        Ok(slf)
    }

    fn is_expired(&self) -> bool {
        self.inner.is_expired()
    }

    fn __repr__(&self) -> String {
        format!(
            "MemoryEntry(id='{}', content='{}', metadata={:?})",
            self.inner.id, self.inner.content, self.inner.metadata
        )
    }
}

/// Python wrapper for MemoryQuery
#[pyclass(subclass)]
#[derive(Clone)]
struct PyMemoryQuery {
    inner: MemoryQuery,
}

#[pymethods]
impl PyMemoryQuery {
    #[new]
    fn new() -> Self {
        PyMemoryQuery {
            inner: MemoryQuery::new(),
        }
    }

    fn with_filter(
        mut slf: PyRefMut<'_, Self>,
        key: String,
        value: Py<PyAny>,
    ) -> PyResult<PyRefMut<'_, Self>> {
        Python::with_gil(|py| {
            let val = value.bind(py);
            let json = py.import("json")?;
            let json_str: String = json.call_method1("dumps", (val,))?.extract()?;
            let json_value: Value = serde_json::from_str(&json_str)
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

            slf.inner = slf.inner.clone().with_filter(key, json_value);
            Ok(slf)
        })
    }

    fn with_limit(mut slf: PyRefMut<'_, Self>, limit: usize) -> PyResult<PyRefMut<'_, Self>> {
        slf.inner = slf.inner.clone().with_limit(limit);
        Ok(slf)
    }

    fn __repr__(&self) -> String {
        format!(
            "MemoryQuery(filters={:?}, limit={:?})",
            self.inner.filters, self.inner.limit
        )
    }
}

/// Wrapper to adapt Python Memory to Rust Memory trait
struct PythonMemoryWrapper {
    memory: Py<PyAny>,
}

unsafe impl Send for PythonMemoryWrapper {}
unsafe impl Sync for PythonMemoryWrapper {}

#[async_trait]
impl runtime::core::memory::Memory for PythonMemoryWrapper {
    async fn store(&self, entry: runtime::core::memory::MemoryEntry) -> runtime::core::error::Result<String> {
        // Clone the reference within the GIL context before moving to spawn_blocking
        let memory = Python::with_gil(|py| self.memory.clone_ref(py));

        tokio::task::spawn_blocking(move || {
            Python::with_gil(|py| {
                let py_entry = PyMemoryEntry { inner: entry };
                let id = memory.bind(py).call_method1("store", (py_entry,))
                    .map_err(|e| runtime::core::error::Error::MeshError(e.to_string()))?;
                id.extract::<String>()
                    .map_err(|e| runtime::core::error::Error::MeshError(e.to_string()))
            })
        })
        .await
        .map_err(|e| runtime::core::error::Error::MeshError(e.to_string()))?
    }

    async fn get(&self, id: &str) -> runtime::core::error::Result<Option<runtime::core::memory::MemoryEntry>> {
        let id_owned = id.to_string();
        let memory = Python::with_gil(|py| self.memory.clone_ref(py));

        tokio::task::spawn_blocking(move || {
            Python::with_gil(|py| -> runtime::core::error::Result<Option<runtime::core::memory::MemoryEntry>> {
                let obj = memory.bind(py).call_method1("get", (id_owned,))
                    .map_err(|e: PyErr| runtime::core::error::Error::MeshError(e.to_string()))?;

                if obj.is_none() {
                    return Ok(None);
                }

                // Try to extract as PyMemoryEntry
                match obj.extract::<PyMemoryEntry>() {
                    Ok(py_entry) => Ok(Some(py_entry.inner)),
                    Err(e) => Err(runtime::core::error::Error::MeshError(e.to_string()))
                }
            })
        })
        .await
        .map_err(|e| runtime::core::error::Error::MeshError(e.to_string()))?
    }

    async fn search(&self, query: runtime::core::memory::MemoryQuery) -> runtime::core::error::Result<Vec<runtime::core::memory::MemoryEntry>> {
        let memory = Python::with_gil(|py| self.memory.clone_ref(py));

        tokio::task::spawn_blocking(move || {
            Python::with_gil(|py| {
                let py_query = PyMemoryQuery { inner: query };
                let list = memory.bind(py).call_method1("search", (py_query,))
                    .map_err(|e: PyErr| runtime::core::error::Error::MeshError(e.to_string()))?;

                let entries: Vec<PyMemoryEntry> = list.extract()
                    .map_err(|e: PyErr| runtime::core::error::Error::MeshError(e.to_string()))?;
                Ok(entries.into_iter().map(|e| e.inner).collect())
            })
        })
        .await
        .map_err(|e| runtime::core::error::Error::MeshError(e.to_string()))?
    }

    async fn delete(&self, id: &str) -> runtime::core::error::Result<bool> {
        let id_owned = id.to_string();
        let memory = Python::with_gil(|py| self.memory.clone_ref(py));

        tokio::task::spawn_blocking(move || {
            Python::with_gil(|py| {
                let deleted = memory.bind(py).call_method1("delete", (id_owned,))
                    .map_err(|e| runtime::core::error::Error::MeshError(e.to_string()))?;
                deleted.extract::<bool>()
                    .map_err(|e| runtime::core::error::Error::MeshError(e.to_string()))
            })
        })
        .await
        .map_err(|e| runtime::core::error::Error::MeshError(e.to_string()))?
    }

    async fn clear(&self) -> runtime::core::error::Result<()> {
        let memory = Python::with_gil(|py| self.memory.clone_ref(py));

        tokio::task::spawn_blocking(move || {
            Python::with_gil(|py| {
                memory.bind(py).call_method0("clear")
                    .map_err(|e| runtime::core::error::Error::MeshError(e.to_string()))?;
                Ok(())
            })
        })
        .await
        .map_err(|e| runtime::core::error::Error::MeshError(e.to_string()))?
    }

    async fn count(&self) -> runtime::core::error::Result<usize> {
        let memory = Python::with_gil(|py| self.memory.clone_ref(py));

        tokio::task::spawn_blocking(move || {
            Python::with_gil(|py| {
                let count = memory.bind(py).call_method0("count")
                    .map_err(|e| runtime::core::error::Error::MeshError(e.to_string()))?;
                count.extract::<usize>()
                    .map_err(|e| runtime::core::error::Error::MeshError(e.to_string()))
            })
        })
        .await
        .map_err(|e| runtime::core::error::Error::MeshError(e.to_string()))?
    }
}


/// Python wrapper for InMemoryBackend
#[pyclass(subclass)]
#[derive(Clone)]
struct PyInMemoryBackend {
    inner: Arc<InMemoryBackend>,
}

#[pymethods]
impl PyInMemoryBackend {
    #[new]
    fn new() -> Self {
        PyInMemoryBackend {
            inner: Arc::new(InMemoryBackend::new()),
        }
    }

    #[staticmethod]
    fn with_max_entries(max: usize) -> Self {
        PyInMemoryBackend {
            inner: Arc::new(InMemoryBackend::new().with_max_entries(max)),
        }
    }

    #[staticmethod]
    fn with_ttl_seconds(seconds: i64) -> Self {
        PyInMemoryBackend {
            inner: Arc::new(InMemoryBackend::new().with_ttl_seconds(seconds)),
        }
    }

    fn store(&self, entry: PyMemoryEntry) -> PyResult<String> {
        let backend = self.inner.clone();
        let entry_inner = entry.inner.clone();

        RUNTIME.block_on(async move {
            backend
                .store(entry_inner)
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
        })
    }

    fn get(&self, id: String) -> PyResult<Option<PyMemoryEntry>> {
        let backend = self.inner.clone();

        RUNTIME.block_on(async move {
            let result = backend
                .get(&id)
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;

            Ok(result.map(|entry| PyMemoryEntry { inner: entry }))
        })
    }

    fn search(&self, query: PyMemoryQuery) -> PyResult<Vec<PyMemoryEntry>> {
        let backend = self.inner.clone();
        let query_inner = query.inner.clone();

        RUNTIME.block_on(async move {
            let results = backend
                .search(query_inner)
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;

            Ok(results
                .into_iter()
                .map(|entry| PyMemoryEntry { inner: entry })
                .collect())
        })
    }

    fn delete(&self, id: String) -> PyResult<bool> {
        let backend = self.inner.clone();

        RUNTIME.block_on(async move {
            backend
                .delete(&id)
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
        })
    }

    fn clear(&self) -> PyResult<()> {
        let backend = self.inner.clone();

        RUNTIME.block_on(async move {
            backend
                .clear()
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
        })
    }

    fn count(&self) -> PyResult<usize> {
        let backend = self.inner.clone();

        RUNTIME.block_on(async move {
            backend
                .count()
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
        })
    }

    fn __repr__(&self) -> String {
        "InMemoryBackend()".to_string()
    }
}

// ============================================================================
// Redis Backend Python Bindings
// ============================================================================

#[pyclass(subclass)]
#[derive(Clone)]
struct PyRedisBackend {
    inner: Arc<RedisBackend>,
}

#[pymethods]
impl PyRedisBackend {
    #[new]
    fn new(redis_url: String) -> PyResult<Self> {
        RUNTIME.block_on(async move {
            let backend = RedisBackend::new(redis_url)
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
            Ok(PyRedisBackend {
                inner: Arc::new(backend),
            })
        })
    }

    fn with_prefix(&self, prefix: String) -> Self {
        let backend = (*self.inner).clone().with_prefix(prefix);
        PyRedisBackend {
            inner: Arc::new(backend),
        }
    }

    fn with_ttl_seconds(&self, seconds: i64) -> Self {
        let backend = (*self.inner).clone().with_ttl_seconds(seconds);
        PyRedisBackend {
            inner: Arc::new(backend),
        }
    }

    fn store(&self, entry: PyMemoryEntry) -> PyResult<String> {
        let backend = self.inner.clone();
        let entry_inner = entry.inner.clone();

        RUNTIME.block_on(async move {
            backend
                .store(entry_inner)
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
        })
    }

    fn get(&self, id: String) -> PyResult<Option<PyMemoryEntry>> {
        let backend = self.inner.clone();

        RUNTIME.block_on(async move {
            let result = backend
                .get(&id)
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;

            Ok(result.map(|entry| PyMemoryEntry { inner: entry }))
        })
    }

    fn search(&self, query: PyMemoryQuery) -> PyResult<Vec<PyMemoryEntry>> {
        let backend = self.inner.clone();
        let query_inner = query.inner.clone();

        RUNTIME.block_on(async move {
            let results = backend
                .search(query_inner)
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;

            Ok(results
                .into_iter()
                .map(|entry| PyMemoryEntry { inner: entry })
                .collect())
        })
    }

    fn delete(&self, id: String) -> PyResult<bool> {
        let backend = self.inner.clone();

        RUNTIME.block_on(async move {
            backend
                .delete(&id)
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
        })
    }

    fn clear(&self) -> PyResult<()> {
        let backend = self.inner.clone();

        RUNTIME.block_on(async move {
            backend
                .clear()
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
        })
    }

    fn count(&self) -> PyResult<usize> {
        let backend = self.inner.clone();

        RUNTIME.block_on(async move {
            backend
                .count()
                .await
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
        })
    }

    fn __repr__(&self) -> String {
        "RedisBackend()".to_string()
    }
}

// ============================================================================
// ReAct Framework Python Bindings
// ============================================================================

/// Python wrapper for ReActConfig
#[pyclass]
#[derive(Clone)]
struct PyReActConfig {
    inner: runtime::llm::ReActConfig,
}

#[pymethods]
impl PyReActConfig {
    #[new]
    fn new() -> Self {
        PyReActConfig {
            inner: runtime::llm::ReActConfig::default(),
        }
    }
   
    fn with_max_iterations(mut slf: PyRefMut<'_, Self>, max_iterations: usize) -> PyResult<PyRefMut<'_, Self>> {
        slf.inner.max_iterations = max_iterations;
        Ok(slf)
    }
   
    fn with_thought_prefix(mut slf: PyRefMut<'_, Self>, prefix: String) -> PyResult<PyRefMut<'_, Self>> {
        slf.inner.thought_prefix = prefix;
        Ok(slf)
    }
   
    fn with_action_prefix(mut slf: PyRefMut<'_, Self>, prefix: String) -> PyResult<PyRefMut<'_, Self>> {
        slf.inner.action_prefix = prefix;
        Ok(slf)
    }
}

/// Python wrapper for ReActStep
#[pyclass]
#[derive(Clone)]
struct PyReActStep {
    inner: runtime::llm::ReActStep,
}

#[pymethods]
impl PyReActStep {
    #[getter]
    fn iteration(&self) -> usize {
        self.inner.iteration
    }
   
    #[getter]
    fn thought(&self) -> String {
        self.inner.thought.clone()
    }
   
    #[getter]
    fn action(&self) -> Option<String> {
        self.inner.action.clone()
    }
   
    #[getter]
    fn action_input(&self) -> Option<String> {
        self.inner.action_input.clone()
    }
   
    #[getter]
    fn observation(&self) -> Option<String> {
        self.inner.observation.clone()
    }
   
    fn __repr__(&self) -> String {
        format!("ReActStep(iteration={}, thought='{}...')", self.inner.iteration, &self.inner.thought.chars().take(50).collect::<String>())
    }
}

/// Python wrapper for ReActResult
#[pyclass]
struct PyReActResult {
    inner: runtime::llm::ReActResult,
}

#[pymethods]
impl PyReActResult {
    #[getter]
    fn answer(&self) -> String {
        self.inner.answer.clone()
    }
   
    #[getter]
    fn iterations(&self) -> usize {
        self.inner.iterations
    }
   
    #[getter]
    fn finish_reason(&self) -> String {
        format!("{:?}", self.inner.finish_reason)
    }
   
    fn get_steps(&self) -> Vec<PyReActStep> {
        self.inner.steps.iter().map(|s| PyReActStep { inner: s.clone() }).collect()
    }
   
    fn __repr__(&self) -> String {
        format!("ReActResult(answer='...', iterations={})", self.inner.iterations)
    }
}

#[pyfunction]
fn get_metrics() -> PyResult<Py<PyDict>> {
    let snapshot = runtime::metrics::metrics().snapshot();
    Python::with_gil(|py| {
        let dict = PyDict::new(py);
        dict.set_item("message_throughput", snapshot.message_throughput)?;
        dict.set_item("avg_message_latency_us", snapshot.avg_message_latency_us)?;
        dict.set_item("avg_agent_execution_time_us", snapshot.avg_agent_execution_time_us)?;
        dict.set_item("total_llm_tokens", snapshot.total_llm_tokens)?;
        dict.set_item("avg_llm_latency_us", snapshot.avg_llm_latency_us)?;
        dict.set_item("total_llm_cost_us", snapshot.total_llm_cost_us)?;
        dict.set_item("memory_hits", snapshot.memory_hits)?;
        dict.set_item("memory_misses", snapshot.memory_misses)?;
        dict.set_item("memory_writes", snapshot.memory_writes)?;
        
        let errors = PyDict::new(py);
        for (k, v) in snapshot.errors {
            errors.set_item(k, v)?;
        }
        dict.set_item("errors", errors)?;
        
        Ok(dict.into())
    })
}

#[pymodule]
fn ceylonai_next(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyLocalMesh>()?;
    m.add_class::<PyDistributedMesh>()?;
    m.add_class::<PyAgent>()?;
    m.add_class::<PyAction>()?;
    m.add_class::<PyToolInvoker>()?;
    m.add_class::<PyLlmAgent>()?;
    m.add_class::<PyAgentContext>()?;
    m.add_class::<PyMemoryEntry>()?;
    m.add_class::<PyMemoryQuery>()?;
    m.add_class::<PyInMemoryBackend>()?;
    m.add_class::<PyRedisBackend>()?;
    m.add_class::<PyLlmConfig>()?;
    m.add_class::<PyReActConfig>()?;
    m.add_class::<PyReActStep>()?;
    m.add_class::<PyReActResult>()?;
    m.add_class::<PyAgentMetadata>()?;
    m.add_class::<PyInMemoryRegistry>()?;
    m.add_class::<PyLoggingConfig>()?;
    m.add_class::<PyLoggingHandle>()?;
    m.add_function(wrap_pyfunction!(init_logging_py, m)?)?;
    m.add_function(wrap_pyfunction!(get_metrics, m)?)?;
    Ok(())
}
