use lumina_engine::{
    Event, EventBus, EventType, Memory, Runtime,
};
use lumina_tools::{ToolContext, ToolExecutor, ToolResult};
use lumina_adapters::{LLMAdapter, LLMRequest, LLMResponse, TokenUsage};
use rand::{Rng, SeedableRng};
use rand_xoshiro::Xoshiro256PlusPlus;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{error, info};

/// Agent state
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentState {
    pub id: String,
    pub name: String,
    pub seed: u64,
    pub current_step: u32,
    pub status: AgentStatus,
    pub metadata: HashMap<String, String>,
}

/// Agent status
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum AgentStatus {
    Created,
    Running,
    Completed,
    Failed,
    Paused,
}

/// Agent step result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StepResult {
    pub step_id: String,
    pub step_number: u32,
    pub success: bool,
    pub output: serde_json::Value,
    pub error: Option<String>,
    pub tool_results: Vec<ToolResult>,
    pub llm_response: Option<LLMResponse>,
}

/// Agent workflow definition
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkflowStep {
    pub step_id: String,
    pub step_type: StepType,
    pub prompt: Option<String>,
    pub tool_calls: Vec<ToolCall>,
    pub condition: Option<String>, // For conditional steps
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum StepType {
    LLMCall,
    ToolExecution,
    Parallel, // Execute multiple steps in parallel
    Conditional, // Conditional branching
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
    pub tool_id: String,
    pub parameters: serde_json::Value,
}

/// Core Agent structure
pub struct Agent {
    state: Arc<RwLock<AgentState>>,
    #[allow(dead_code)]
    runtime: Arc<Runtime>,
    memory: Arc<Memory>,
    tool_executor: Arc<ToolExecutor>,
    llm_adapter: Arc<dyn LLMAdapter>,
    workflow: Arc<RwLock<Vec<WorkflowStep>>>,
    event_bus: Arc<EventBus>,
    rng: Arc<RwLock<Xoshiro256PlusPlus>>,
}

impl Agent {
    pub fn new(
        id: String,
        name: String,
        seed: u64,
        runtime: Arc<Runtime>,
        memory: Arc<Memory>,
        tool_executor: Arc<ToolExecutor>,
        llm_adapter: Arc<dyn LLMAdapter>,
        workflow: Vec<WorkflowStep>,
    ) -> Self {
        let state = AgentState {
            id: id.clone(),
            name,
            seed,
            current_step: 0,
            status: AgentStatus::Created,
            metadata: HashMap::new(),
        };
        
        let event_bus = runtime.event_bus();
        let rng = Arc::new(RwLock::new(Xoshiro256PlusPlus::seed_from_u64(seed)));
        
        Self {
            state: Arc::new(RwLock::new(state)),
            runtime,
            memory,
            tool_executor,
            llm_adapter,
            workflow: Arc::new(RwLock::new(workflow)),
            event_bus,
            rng,
        }
    }
    
    /// Set the agent's workflow
    pub async fn set_workflow(&self, workflow: Vec<WorkflowStep>) {
        let mut wf = self.workflow.write().await;
        *wf = workflow;
    }
    
    /// Get agent ID
    pub async fn id(&self) -> String {
        self.state.read().await.id.clone()
    }
    
    /// Get agent state
    pub async fn get_state(&self) -> AgentState {
        self.state.read().await.clone()
    }
    
    /// Run the agent workflow
    pub async fn run(&self, input: Option<String>) -> Result<Vec<StepResult>, AgentError> {
        let mut state = self.state.write().await;
        state.status = AgentStatus::Running;
        drop(state);
        
        let agent_id = self.id().await;
        info!("Starting agent workflow: {}", agent_id);
        
        // Emit agent started event
        self.emit_event(EventType::AgentStarted, serde_json::json!({
            "input": input
        })).await;
        
        let mut results = Vec::new();
        
        // Get workflow
        let workflow = self.workflow.read().await.clone();
        
        // If workflow is empty and we have input, create an automatic LLM step
        let workflow_to_execute = if workflow.is_empty() && input.is_some() {
            vec![WorkflowStep {
                step_id: "auto_llm_step".to_string(),
                step_type: StepType::LLMCall,
                prompt: Some(input.clone().unwrap_or_default()),
                tool_calls: Vec::new(),
                condition: None,
            }]
        } else {
            workflow
        };
        
        for (idx, step) in workflow_to_execute.iter().enumerate() {
            let step_number = idx as u32 + 1;
            let step_id = format!("{}-step-{}", agent_id, step_number);
            
            info!("Executing step {}: {}", step_number, step.step_id);
            
            match self.execute_step(step, &step_id, step_number, input.as_deref()).await {
                Ok(result) => {
                    results.push(result.clone());
                    
                    // Emit step completed event
                    self.emit_event(EventType::AgentStepCompleted, serde_json::json!({
                        "step_id": step.step_id,
                        "step_number": step_number,
                        "success": result.success,
                    })).await;
                    
                    // Update state
                    let mut state = self.state.write().await;
                    state.current_step = step_number;
                    
                    // If step failed, stop workflow
                    if !result.success {
                        state.status = AgentStatus::Failed;
                        self.emit_event(EventType::AgentFailed, serde_json::json!({
                            "step_id": step.step_id,
                            "error": result.error,
                        })).await;
                        return Err(AgentError::StepFailed(step.step_id.clone(), result.error));
                    }
                }
                Err(e) => {
                    error!("Step {} failed: {}", step.step_id, e);
                    let mut state = self.state.write().await;
                    state.status = AgentStatus::Failed;
                    self.emit_event(EventType::AgentFailed, serde_json::json!({
                        "step_id": step.step_id,
                        "error": e.to_string(),
                    })).await;
                    return Err(e);
                }
            }
        }
        
        // Workflow completed
        let mut state = self.state.write().await;
        state.status = AgentStatus::Completed;
        drop(state);
        
        self.emit_event(EventType::AgentCompleted, serde_json::json!({
            "steps_completed": results.len(),
        })).await;
        
        info!("Agent workflow completed: {}", agent_id);
        Ok(results)
    }
    
    /// Execute a single workflow step
    async fn execute_step(
        &self,
        step: &WorkflowStep,
        step_id: &str,
        step_number: u32,
        input: Option<&str>,
    ) -> Result<StepResult, AgentError> {
        let mut tool_results = Vec::new();
        let mut llm_response = None;
        
        match step.step_type {
            StepType::LLMCall => {
                // Build prompt with context
                let prompt = step.prompt.as_ref().ok_or_else(|| {
                    AgentError::InvalidStep("LLM step requires prompt".to_string())
                })?;
                
                let agent_id = self.id().await;
                let context = self.memory.get_agent_context(&agent_id).await;
                let full_prompt = self.build_prompt_with_context(prompt, &context, input).await;
                
                // Get deterministic RNG for this step
                let seed = {
                    let mut rng = self.rng.write().await;
                    rng.gen::<u64>()
                };
                
                // Emit LLM call requested event (event-driven)
                // The LLM handler will process this and emit LLMCallCompleted or LLMCallFailed
                self.emit_event(EventType::LLMCallRequested, serde_json::json!({
                    "step_id": step_id,
                    "prompt": full_prompt,
                    "temperature": 0.7,
                    "max_tokens": 1000,
                    "seed": seed,
                })).await;
                
                // Wait briefly for event-driven response, then fallback to direct call
                // This ensures we don't hang forever if the event handler doesn't respond
                let step_id_clone = step_id.to_string();
                let agent_id_clone = self.id().await;
                let event_bus = self.event_bus.clone();
                
                // Wait a short time for event-driven response
                let wait_result = tokio::time::timeout(
                    std::time::Duration::from_secs(2),
                    async {
                        // Poll event history for completion/failure events
                        for _ in 0..20 {
                            let history = event_bus.get_history(Some(20)).await;
                            
                            // Look for LLMCallCompleted or LLMCallFailed for this step
                            for event in history.iter().rev() {
                                if event.source == agent_id_clone {
                                    match event.event_type {
                                        EventType::LLMCallCompleted => {
                                            if let Some(payload_step_id) = event.payload.get("step_id")
                                                .and_then(|v| v.as_str()) {
                                                if payload_step_id == step_id_clone {
                                                    // Extract response from event
                                                    if let Some(response_text) = event.payload.get("response")
                                                        .and_then(|v| v.as_str()) {
                                                        let response = LLMResponse {
                                                            text: response_text.to_string(),
                                                            model: event.payload.get("model")
                                                                .and_then(|v| v.as_str())
                                                                .unwrap_or("unknown")
                                                                .to_string(),
                                                            usage: TokenUsage {
                                                                prompt_tokens: event.payload.get("usage")
                                                                    .and_then(|u| u.get("prompt_tokens"))
                                                                    .and_then(|v| v.as_u64())
                                                                    .map(|u| u as u32)
                                                                    .unwrap_or(0),
                                                                completion_tokens: event.payload.get("usage")
                                                                    .and_then(|u| u.get("completion_tokens"))
                                                                    .and_then(|v| v.as_u64())
                                                                    .map(|u| u as u32)
                                                                    .unwrap_or(0),
                                                                total_tokens: event.payload.get("usage")
                                                                    .and_then(|u| u.get("total_tokens"))
                                                                    .and_then(|v| v.as_u64())
                                                                    .map(|u| u as u32)
                                                                    .unwrap_or(0),
                                                            },
                                                            finish_reason: event.payload.get("finish_reason")
                                                                .and_then(|v| v.as_str())
                                                                .map(|s| s.to_string()),
                                                            metadata: HashMap::new(),
                                                        };
                                                        return Ok(Some(response));
                                                    }
                                                }
                                            }
                                        }
                                        EventType::LLMCallFailed => {
                                            if let Some(payload_step_id) = event.payload.get("step_id")
                                                .and_then(|v| v.as_str()) {
                                                if payload_step_id == step_id_clone {
                                                    let error_msg = event.payload.get("error")
                                                        .and_then(|v| v.as_str())
                                                        .unwrap_or("LLM call failed")
                                                        .to_string();
                                                    return Err(error_msg);
                                                }
                                            }
                                        }
                                        _ => {}
                                    }
                                }
                            }
                            
                            // Small delay before checking again
                            tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
                        }
                        Ok(None) // No event found, will use fallback
                    }
                ).await;
                
                // Handle the result
                let response = match wait_result {
                    Ok(Ok(Some(resp))) => {
                        // Got completion event
                        resp
                    }
                    Ok(Ok(None)) | Err(_) => {
                        // No event or timeout - fallback to direct call
                        let request = LLMRequest {
                            prompt: full_prompt.clone(),
                            system_prompt: None,
                            temperature: 0.7,
                            max_tokens: Some(1000),
                            seed,
                            stop_sequences: Vec::new(),
                            metadata: HashMap::new(),
                        };
                        
                        // Add timeout to prevent hanging
                        let llm_result = tokio::time::timeout(
                            std::time::Duration::from_secs(30),
                            self.llm_adapter.generate(&request)
                        ).await;
                        
                        match llm_result {
                            Ok(Ok(resp)) => {
                                // Emit completion event
                                self.emit_event(EventType::LLMCallCompleted, serde_json::json!({
                                    "step_id": step_id,
                                    "response": resp.text.clone(),
                                    "usage": resp.usage,
                                })).await;
                                resp
                            }
                            Ok(Err(ref e)) => {
                                // API error
                                let error_msg = e.to_string();
                                self.emit_event(EventType::LLMCallFailed, serde_json::json!({
                                    "step_id": step_id,
                                    "error": error_msg.clone(),
                                })).await;
                                self.emit_event(EventType::AgentStepFailed, serde_json::json!({
                                    "step_id": step_id,
                                    "error": error_msg.clone(),
                                })).await;
                                return Err(AgentError::LLMError(error_msg));
                            }
                            Err(_) => {
                                // Timeout
                                let error_msg = "LLM call timed out after 30 seconds".to_string();
                                self.emit_event(EventType::LLMCallFailed, serde_json::json!({
                                    "step_id": step_id,
                                    "error": error_msg.clone(),
                                })).await;
                                self.emit_event(EventType::AgentStepFailed, serde_json::json!({
                                    "step_id": step_id,
                                    "error": error_msg.clone(),
                                })).await;
                                return Err(AgentError::LLMError(error_msg));
                            }
                        }
                    }
                    Ok(Err(e)) => {
                        // Got failure event - stop execution immediately
                        self.emit_event(EventType::AgentStepFailed, serde_json::json!({
                            "step_id": step_id,
                            "error": e.clone(),
                        })).await;
                        return Err(AgentError::LLMError(e));
                    }
                };
                
                llm_response = Some(response.clone());
                
                // Store LLM response in memory
                let agent_id = self.id().await;
                self.memory.write_step(
                    &agent_id,
                    step_id,
                    "llm_response".to_string(),
                    serde_json::json!({ "text": response.text }),
                ).await;
            }
            
            StepType::ToolExecution => {
                for tool_call in &step.tool_calls {
                    let agent_id = self.id().await;
                let context = ToolContext {
                        agent_id: agent_id.clone(),
                        step_id: Some(step_id.to_string()),
                        seed: {
                            let mut rng = self.rng.write().await;
                            rng.gen::<u64>()
                        },
                        metadata: HashMap::new(),
                    };
                    
                    let result = self.tool_executor.execute(
                        &tool_call.tool_id,
                        tool_call.parameters.clone(),
                        &context,
                    ).await
                    .map_err(|e| AgentError::ToolError(e.to_string()))?;
                    
                    tool_results.push(result.clone());
                    
                    // Store tool result in memory
                    if result.success {
                        let agent_id = self.id().await;
                        self.memory.write_tool(
                            &agent_id,
                            step_id,
                            &tool_call.tool_id,
                            "result".to_string(),
                            result.output.clone(),
                        ).await;
                    }
                }
            }
            
            StepType::Parallel => {
                // Execute multiple tool calls in parallel (premium concurrency feature)
                if step.tool_calls.is_empty() {
                    return Err(AgentError::InvalidStep("Parallel step requires at least one tool call".to_string()));
                }
                
                info!("Executing {} tools in parallel", step.tool_calls.len());
                
                let agent_id = self.id().await;
                let mut parallel_tasks = Vec::new();
                
                // Prepare all tool calls for parallel execution
                for tool_call in &step.tool_calls {
                    let tool_id = tool_call.tool_id.clone();
                    let parameters = tool_call.parameters.clone();
                    let context = ToolContext {
                        agent_id: agent_id.clone(),
                        step_id: Some(step_id.to_string()),
                        seed: {
                            let mut rng = self.rng.write().await;
                            rng.gen::<u64>()
                        },
                        metadata: HashMap::new(),
                    };
                    
                    let tool_executor = self.tool_executor.clone();
                    let tool_id_clone = tool_id.clone();
                    let parameters_clone = parameters.clone();
                    let context_clone = context.clone();
                    
                    info!("Preparing parallel task for tool: {}", tool_id_clone);
                    
                    // Spawn async task for parallel execution
                    // Note: This works perfectly for Rust tools (the main use case)
                    // Python tools will execute sequentially to avoid GIL deadlocks
                    // The whole point of Rust in Lumina is to avoid Python GIL limitations
                    let tool_executor_clone = tool_executor.clone();
                    let tool_id_for_task = tool_id_clone.clone();
                    let parameters_for_task = parameters_clone.clone();
                    let context_for_task = context_clone.clone();
                    
                    info!("[Parallel] Spawning task for tool: {}", tool_id_for_task);
                    
                    // Spawn async task - works great for Rust tools
                    let task_handle = tokio::task::spawn(async move {
                        info!("[Parallel] Executing tool: {}", tool_id_for_task);
                        let result = tool_executor_clone.execute(&tool_id_for_task, parameters_for_task, &context_for_task).await;
                        info!("[Parallel] Tool {} completed", tool_id_for_task);
                        result
                    });
                    
                    parallel_tasks.push((tool_id, task_handle));
                }
                
                info!("Waiting for {} parallel tasks to complete", parallel_tasks.len());
                
                // Wait for all parallel tool executions to complete
                for (tool_id, task) in parallel_tasks {
                    info!("Waiting for tool: {}", tool_id);
                    match task.await {
                        Ok(Ok(result)) => {
                            info!("Tool {} completed successfully", tool_id);
                            let result_clone = result.clone();
                            tool_results.push(result_clone.clone());
                            
                            // Store tool result in memory
                            if result_clone.success {
                                let agent_id = self.id().await;
                                self.memory.write_tool(
                                    &agent_id,
                                    step_id,
                                    &tool_id,
                                    "result".to_string(),
                                    result_clone.output.clone(),
                                ).await;
                            }
                        }
                        Ok(Err(e)) => {
                            error!("Parallel tool {} failed: {}", tool_id, e);
                            return Err(AgentError::ToolError(format!("Parallel tool {} failed: {}", tool_id, e)));
                        }
                        Err(e) => {
                            error!("Parallel tool {} task panicked: {}", tool_id, e);
                            return Err(AgentError::ToolError(format!("Parallel tool {} task panicked: {}", tool_id, e)));
                        }
                    }
                }
                
                info!("Parallel step completed with {} tool calls", tool_results.len());
            }
            
            StepType::Conditional => {
                return Err(AgentError::NotImplemented("Conditional steps not yet implemented".to_string()));
            }
        }
        
        Ok(StepResult {
            step_id: step_id.to_string(),
            step_number,
            success: true,
            output: serde_json::json!({
                "tool_results": tool_results.len(),
                "llm_response": llm_response.is_some(),
            }),
            error: None,
            tool_results,
            llm_response,
        })
    }
    
    /// Build prompt with context from memory
    async fn build_prompt_with_context(
        &self,
        base_prompt: &str,
        context: &HashMap<String, serde_json::Value>,
        input: Option<&str>,
    ) -> String {
        let mut prompt = String::new();
        
        if let Some(input) = input {
            prompt.push_str(&format!("Input: {}\n\n", input));
        }
        
        if !context.is_empty() {
            prompt.push_str("Context:\n");
            for (key, value) in context.iter() {
                prompt.push_str(&format!("  {}: {}\n", key, value));
            }
            prompt.push_str("\n");
        }
        
        prompt.push_str(base_prompt);
        prompt
    }
    
    /// Emit an event
    async fn emit_event(&self, event_type: EventType, payload: serde_json::Value) {
        let agent_id = self.id().await;
        let state = self.state.read().await;
        let event = Event::new(
            event_type,
            agent_id,
            payload,
            state.seed,
        );
        
        if let Err(e) = self.event_bus.publish(event).await {
            error!("Failed to emit event: {}", e);
        }
    }
    
    /// Serialize agent state for replay
    pub async fn serialize_state(&self) -> Result<Vec<u8>, bincode::Error> {
        let state = self.get_state().await;
        bincode::serialize(&state)
    }
    
    /// Deserialize agent state for replay
    pub async fn deserialize_state(&self, data: &[u8]) -> Result<(), bincode::Error> {
        let state: AgentState = bincode::deserialize(data)?;
        *self.state.write().await = state;
        Ok(())
    }
}

/// Agent errors
#[derive(Debug, thiserror::Error)]
pub enum AgentError {
    #[error("Step failed: {0} - {1:?}")]
    StepFailed(String, Option<String>),
    
    #[error("LLM error: {0}")]
    LLMError(String),
    
    #[error("Tool error: {0}")]
    ToolError(String),
    
    #[error("Invalid step: {0}")]
    InvalidStep(String),
    
    #[error("Not implemented: {0}")]
    NotImplemented(String),
    
    #[error("Serialization error: {0}")]
    SerializationError(#[from] bincode::Error),
}

