#![allow(non_local_definitions)] // PyO3 #[pymethods] generates non-local impl blocks

pub mod agent;
pub mod python_tool;
pub mod python_event_handler;

use pyo3::prelude::*;
use lumina_engine::{Runtime, Memory, EventType};
use lumina_tools::{ToolRegistry, ToolCache, ToolSystem, Tool as RustTool};
use lumina_adapters::{LLMAdapter, GeminiAdapter, LLMEventHandler};
use std::sync::Arc;
use uuid::Uuid;
use serde_json::Value;

// Re-export for Python
use agent::{Agent, WorkflowStep};
use python_tool::PythonTool;
use python_event_handler::PythonEventHandler;

/// Python wrapper for Lumina Runtime
#[pyclass]
pub struct LuminaRuntime {
    runtime: Arc<Runtime>,
    memory: Arc<Memory>,
    tool_registry: Arc<ToolRegistry>,
    tool_cache: Arc<ToolCache>,
    tool_system: Arc<ToolSystem>,
}

#[pymethods]
impl LuminaRuntime {
    #[new]
    fn new() -> PyResult<Self> {
        // Initialize runtime
        let runtime = Arc::new(Runtime::default());
        let memory = Arc::new(Memory::new());
        let tool_registry = Arc::new(ToolRegistry::new());
        let tool_cache = Arc::new(ToolCache::default());
        let event_bus = runtime.event_bus();
        
        // Create ToolSystem (event-driven tool dispatcher)
        let tool_system = Arc::new(ToolSystem::new(
            tool_registry.clone(),
            tool_cache.clone(),
            event_bus.clone(),
            5000, // 5s timeout
            3,    // 3 retries
            1000, // 1s retry delay
        ));
        
        // Register ToolSystem as an event handler for ToolCallRequested events
        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async {
            event_bus.subscribe(tool_system.clone() as Arc<dyn lumina_engine::EventHandler>).await;
        });
        
        Ok(Self {
            runtime,
            memory,
            tool_registry,
            tool_cache,
            tool_system,
        })
    }
    
    /// Start the runtime
    fn start(&self) -> PyResult<()> {
        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async {
            self.runtime.start().await;
        });
        Ok(())
    }
    
    /// Stop the runtime
    fn stop(&self) -> PyResult<()> {
        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async {
            self.runtime.stop().await;
        });
        Ok(())
    }
    
    /// Get runtime statistics
    fn stats(&self) -> PyResult<String> {
        let rt = tokio::runtime::Runtime::new().unwrap();
        let stats = rt.block_on(async {
            self.runtime.stats().await
        });
        Ok(format!("{:?}", stats))
    }
    
    /// Register a Python tool
    fn register_python_tool(
        &self,
        name: String,
        description: String,
        parameters: String, // JSON string
        py_func: PyObject,
    ) -> PyResult<()> {
        let params: Value = serde_json::from_str(&parameters)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(
                format!("Invalid parameters JSON: {}", e)
            ))?;
        
        let tool: Arc<dyn RustTool> = Arc::new(PythonTool::new(
            name,
            description,
            params,
            py_func,
        ));
        
        let registry = self.tool_registry.clone();
        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async {
            registry.register(tool).await;
        });
        
        Ok(())
    }
    
    /// Get tool registry (for advanced usage)
    fn get_tool_registry(&self) -> PyResult<PyObject> {
        // Return a reference that can be used for advanced operations
        Python::with_gil(|py| Ok(py.None()))
    }
    
    /// Set memory at a specific layer
    fn set_memory(
        &self,
        layer: String,
        agent_id: String,
        key: String,
        value: String, // JSON string
    ) -> PyResult<()> {
        let value_json: Value = serde_json::from_str(&value)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(
                format!("Invalid value JSON: {}", e)
            ))?;
        
        let memory = self.memory.clone();
        let rt = tokio::runtime::Runtime::new().unwrap();
        
        rt.block_on(async {
            match layer.as_str() {
                "global" => {
                    memory.write_global(key, value_json).await;
                    Ok(())
                }
                "agent" => {
                    memory.write_agent(&agent_id, key, value_json).await;
                    Ok(())
                }
                _ => {
                    Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                        format!("Invalid memory layer: {}", layer)
                    ))
                }
            }
        })
    }
    
    /// Get memory from a specific layer
    fn get_memory(
        &self,
        layer: String,
        agent_id: String,
        key: String,
    ) -> PyResult<Option<String>> {
        let memory = self.memory.clone();
        let rt = tokio::runtime::Runtime::new().unwrap();
        
        let result = rt.block_on(async {
            match layer.as_str() {
                "global" => {
                    memory.read_global(&key).await
                }
                "agent" => {
                    memory.read_agent(&agent_id, &key).await
                }
                _ => None,
            }
        });
        
        Ok(result.map(|v| serde_json::to_string(&v).unwrap_or_default()))
    }
    
    /// Register a Python event handler
    fn register_event_handler(
        &self,
        event_type_str: String,
        py_handler: PyObject,
    ) -> PyResult<()> {
        // Parse event type
        let event_type = match event_type_str.as_str() {
            "AgentCreated" => EventType::AgentCreated,
            "AgentStarted" => EventType::AgentStarted,
            "AgentStepCompleted" => EventType::AgentStepCompleted,
            "AgentStepFailed" => EventType::AgentStepFailed,
            "AgentCompleted" => EventType::AgentCompleted,
            "AgentFailed" => EventType::AgentFailed,
            "ToolCallRequested" => EventType::ToolCallRequested,
            "ToolCallStarted" => EventType::ToolCallStarted,
            "ToolCallCompleted" => EventType::ToolCallCompleted,
            "ToolCallFailed" => EventType::ToolCallFailed,
            "ToolCallRetrying" => EventType::ToolCallRetrying,
            "MemoryRead" => EventType::MemoryRead,
            "MemoryWrite" => EventType::MemoryWrite,
            "MemoryUpdated" => EventType::MemoryUpdated,
            "LLMCallRequested" => EventType::LLMCallRequested,
            "LLMCallStarted" => EventType::LLMCallStarted,
            "LLMCallCompleted" => EventType::LLMCallCompleted,
            "LLMCallFailed" => EventType::LLMCallFailed,
            "WorkflowStarted" => EventType::WorkflowStarted,
            "WorkflowStepCompleted" => EventType::WorkflowStepCompleted,
            "WorkflowCompleted" => EventType::WorkflowCompleted,
            "WorkflowFailed" => EventType::WorkflowFailed,
            _ => EventType::Custom(event_type_str.clone()),
        };
        
        let handler: Arc<dyn lumina_engine::EventHandler> = Arc::new(
            PythonEventHandler::new(py_handler, vec![event_type.clone()])
        );
        
        let event_bus = self.runtime.event_bus();
        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async {
            event_bus.subscribe(handler).await;
        });
        
        Ok(())
    }
}

/// Python wrapper for Agent
#[pyclass]
pub struct LuminaAgent {
    agent: Arc<Agent>,
    runtime: Arc<LuminaRuntime>,
}

#[pymethods]
impl LuminaAgent {
    #[new]
    #[pyo3(signature = (name, seed, runtime, gemini_api_key=None))]
    fn new(
        name: String,
        seed: u64,
        runtime: PyRef<LuminaRuntime>,
        gemini_api_key: Option<String>,
    ) -> PyResult<Self> {
        let agent_id = Uuid::new_v4().to_string();
        
        // Create LLM adapter
        let api_key = gemini_api_key.unwrap_or_else(|| {
            std::env::var("GEMINI_API_KEY").unwrap_or_default()
        });
        
        let llm_adapter: Arc<dyn LLMAdapter> = Arc::new(GeminiAdapter::new(api_key));
        
        // Register LLM event handler to listen for LLMCallRequested events
        let event_bus_for_llm = runtime.runtime.event_bus();
        let llm_handler = Arc::new(LLMEventHandler::new(
            llm_adapter.clone(),
            event_bus_for_llm.clone(),
        ));
        
        // Register LLM handler with event bus
        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async {
            event_bus_for_llm.subscribe(llm_handler.clone() as Arc<dyn lumina_engine::EventHandler>).await;
        });
        
        // For now, we need to create a ToolExecutor for backward compatibility
        // TODO: Refactor Agent to use event-driven tool calls
        use lumina_tools::ToolExecutor;
        let event_bus_for_executor = runtime.runtime.event_bus();
        let tool_executor = Arc::new(ToolExecutor::new(
            runtime.tool_registry.clone(),
            runtime.tool_cache.clone(),
            event_bus_for_executor,
            5000,
            3,
            1000,
        ));
        
        // Create agent with empty workflow (can be set later)
        let agent = Arc::new(Agent::new(
            agent_id,
            name,
            seed,
            runtime.runtime.clone(),
            runtime.memory.clone(),
            tool_executor,
            llm_adapter,
            Vec::new(), // Empty workflow initially
        ));
        
        Ok(Self {
            agent,
            runtime: Arc::new(LuminaRuntime {
                runtime: runtime.runtime.clone(),
                memory: runtime.memory.clone(),
                tool_registry: runtime.tool_registry.clone(),
                tool_cache: runtime.tool_cache.clone(),
                tool_system: runtime.tool_system.clone(),
            }),
        })
    }
    
    /// Get agent ID
    fn id(&self) -> PyResult<String> {
        let rt = tokio::runtime::Runtime::new().unwrap();
        let id = rt.block_on(async {
            self.agent.id().await
        });
        Ok(id)
    }
    
    /// Get agent state
    fn get_state(&self) -> PyResult<String> {
        let rt = tokio::runtime::Runtime::new().unwrap();
        let state = rt.block_on(async {
            self.agent.get_state().await
        });
        Ok(format!("{:?}", state))
    }
    
    /// Run agent with optional input
    fn run(&self, input: Option<String>) -> PyResult<String> {
        let rt = tokio::runtime::Runtime::new().unwrap();
        
        // Add overall timeout to prevent hanging forever
        let result = rt.block_on(async {
            tokio::time::timeout(
                std::time::Duration::from_secs(60), // 60 second overall timeout
                self.agent.run(input)
            ).await
        });
        
        // Handle timeout or result
        let agent_result = match result {
            Ok(Ok(results)) => Ok(results),
            Ok(Err(e)) => Err(e),
            Err(_) => {
                // Overall timeout - agent took too long
                return Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                    "Agent execution timed out after 60 seconds"
                ));
            }
        };
        
        // Serialize results to JSON
        let results = agent_result.map_err(|e: agent::AgentError| {
            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                format!("Agent execution failed: {}", e)
            )
        })?;
        
        let json_result = serde_json::to_string(&results)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(
                format!("Failed to serialize results: {}", e)
            ))?;
        
        Ok(json_result)
    }
    
    /// Set workflow steps for the agent
    fn set_workflow(&self, workflow_json: String) -> PyResult<()> {
        // Parse workflow from JSON
        let workflow: Vec<WorkflowStep> = serde_json::from_str(&workflow_json)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(
                format!("Invalid workflow JSON: {}", e)
            ))?;
        
        // Set workflow on the agent
        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async {
            self.agent.set_workflow(workflow).await;
        });
        
        Ok(())
    }
    
    /// Register a tool for this agent
    fn register_tool(
        &self,
        name: String,
        description: String,
        parameters: String,
        py_func: PyObject,
    ) -> PyResult<()> {
        // Register tool in runtime's tool registry
        self.runtime.register_python_tool(name, description, parameters, py_func)
    }
    
    /// Set memory for this agent
    fn set_agent_memory(&self, key: String, value: String) -> PyResult<()> {
        let rt = tokio::runtime::Runtime::new().unwrap();
        let agent_id = rt.block_on(async {
            self.agent.id().await
        });
        
        self.runtime.set_memory("agent".to_string(), agent_id, key, value)
    }
    
    /// Get memory for this agent
    fn get_agent_memory(&self, key: String) -> PyResult<Option<String>> {
        let agent_id = {
            let rt = tokio::runtime::Runtime::new().unwrap();
            rt.block_on(async {
                self.agent.id().await
            })
        };
        
        self.runtime.get_memory("agent".to_string(), agent_id, key)
    }
    
    /// Register event handler for this agent
    fn register_event_handler(&self, event_type: String, py_handler: PyObject) -> PyResult<()> {
        self.runtime.register_event_handler(event_type, py_handler)
    }
}

// Module definition for PyO3
// Note: The function name must match the library name "lumina_core" from Cargo.toml
#[pymodule]
fn lumina_core(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_class::<LuminaRuntime>()?;
    m.add_class::<LuminaAgent>()?;
    Ok(())
}
