use std::sync::Arc;

use chrono::Utc;
use tokio::sync::mpsc;
use tracing::{debug, info, warn};
use vapora_knowledge_graph::{ExecutionRecord, KGPersistence, PersistedExecution};
use vapora_llm_router::{EmbeddingProvider, LLMRouter};

use super::state_machine::{Agent, ExecutionResult, Idle};
use crate::messages::TaskAssignment;
use crate::registry::AgentMetadata;

/// Per-agent executor handling task processing with KG persistence.
pub struct AgentExecutor {
    agent: Agent<Idle>,
    task_rx: mpsc::Receiver<TaskAssignment>,
    kg_persistence: Option<Arc<KGPersistence>>,
    embedding_provider: Option<Arc<dyn EmbeddingProvider>>,
    router: Option<Arc<LLMRouter>>,
}

impl AgentExecutor {
    /// Create a new executor without persistence or LLM routing.
    pub fn new(metadata: AgentMetadata, task_rx: mpsc::Receiver<TaskAssignment>) -> Self {
        Self {
            agent: Agent::new(metadata),
            task_rx,
            kg_persistence: None,
            embedding_provider: None,
            router: None,
        }
    }

    /// Create executor with KG persistence.
    pub fn with_persistence(
        metadata: AgentMetadata,
        task_rx: mpsc::Receiver<TaskAssignment>,
        kg_persistence: Arc<KGPersistence>,
        embedding_provider: Arc<dyn EmbeddingProvider>,
    ) -> Self {
        Self {
            agent: Agent::new(metadata),
            task_rx,
            kg_persistence: Some(kg_persistence),
            embedding_provider: Some(embedding_provider),
            router: None,
        }
    }

    /// Attach an LLM router to this executor.
    ///
    /// When set, [`AgentExecutor::run`] forwards each task to the router using
    /// `AgentMetadata::system_prompt` as the system message and the task
    /// description (+ context when non-empty) as the user message.
    /// Budget-aware routing uses `AgentMetadata::role` for spend tracking.
    ///
    /// When absent, the executor produces a stub result — useful in tests that
    /// must not make network calls.
    pub fn with_router(mut self, router: Arc<LLMRouter>) -> Self {
        self.router = Some(router);
        self
    }

    /// Run the executor loop, processing tasks until the channel closes.
    pub async fn run(mut self) {
        info!(
            "AgentExecutor started for agent: {} (stable_id: {})",
            self.agent.metadata.id,
            self.agent.metadata.stable_id_or_role()
        );
        // Use stable_id so KG records survive hot-reloads and are correlated
        // across agent restarts for the same role.
        let agent_id = self.agent.metadata.stable_id_or_role().to_string();

        while let Some(task) = self.task_rx.recv().await {
            debug!("Received task: {}", task.id);

            let execution_start = Utc::now();

            // Execute before state transitions — state machine consumes self.agent via
            // partial move, which would prevent borrowing self.router afterwards.
            let (result, provider_used) = self.execute_task(&task).await;

            let agent_assigned = self.agent.assign_task(task.clone());
            let agent_executing = agent_assigned.start_execution();
            let completed_agent = agent_executing.complete(result.clone());
            self.agent = completed_agent.reset();

            self.persist_execution_internal(
                &task,
                &result,
                execution_start,
                &agent_id,
                &provider_used,
            )
            .await;

            info!("Task {} completed", task.id);
        }

        info!("AgentExecutor stopped for agent: {}", agent_id);
    }

    /// Execute a task, returning the result and the provider name that handled
    /// it.
    ///
    /// When a router is attached, calls `complete_with_budget` with:
    /// - `system_prompt` from `AgentMetadata` → sent as the LLM system message
    /// - `task.description` (plus `task.context` when meaningful) → user
    ///   message
    /// - `task.required_role` → task-type routing key
    /// - `agent.role` → budget-tracking role
    ///
    /// Falls back to a stub result when no router is configured.
    async fn execute_task(&self, task: &TaskAssignment) -> (ExecutionResult, String) {
        let Some(ref router) = self.router else {
            return (
                ExecutionResult {
                    output: "Task executed successfully".to_string(),
                    input_tokens: 100,
                    output_tokens: 50,
                    duration_ms: 500,
                },
                "unknown".to_string(),
            );
        };

        let prompt = build_prompt(&task.description, &task.context);
        let system_prompt = self.agent.metadata.system_prompt.clone();
        let provider_name = self.agent.metadata.llm_provider.clone();
        let wall_start = std::time::Instant::now();

        match router
            .complete_with_budget(
                &task.required_role,
                prompt,
                system_prompt,
                None,
                Some(self.agent.metadata.role.as_str()),
            )
            .await
        {
            Ok(response) => {
                let duration_ms = wall_start.elapsed().as_millis() as u64;
                (
                    ExecutionResult {
                        output: response.text,
                        input_tokens: response.input_tokens,
                        output_tokens: response.output_tokens,
                        duration_ms,
                    },
                    provider_name,
                )
            }
            Err(e) => {
                warn!("LLM call failed for task {}: {}", task.id, e);
                let duration_ms = wall_start.elapsed().as_millis() as u64;
                (
                    ExecutionResult {
                        output: format!("LLM error: {e}"),
                        input_tokens: 0,
                        output_tokens: 0,
                        duration_ms,
                    },
                    provider_name,
                )
            }
        }
    }

    async fn persist_execution_internal(
        &self,
        task: &TaskAssignment,
        result: &ExecutionResult,
        execution_start: chrono::DateTime<Utc>,
        agent_id: &str,
        provider: &str,
    ) {
        if let Some(ref kg_persistence) = self.kg_persistence {
            if let Some(ref embedding_provider) = self.embedding_provider {
                let embedding = match embedding_provider.embed(&task.description).await {
                    Ok(emb) => emb,
                    Err(e) => {
                        warn!("Failed to generate embedding for task {}: {}", task.id, e);
                        vec![0.0; 1536]
                    }
                };

                let execution_record = ExecutionRecord {
                    id: task.id.clone(),
                    task_id: task.id.clone(),
                    agent_id: agent_id.to_string(),
                    agent_role: None,
                    task_type: task.required_role.clone(),
                    description: task.description.clone(),
                    duration_ms: result.duration_ms,
                    input_tokens: result.input_tokens,
                    output_tokens: result.output_tokens,
                    cost_cents: 0,
                    provider: provider.to_string(),
                    success: true,
                    error: None,
                    solution: Some(result.output.clone()),
                    root_cause: None,
                    timestamp: execution_start,
                };

                let persisted =
                    PersistedExecution::from_execution_record(&execution_record, embedding);

                if let Err(e) = kg_persistence.persist_execution(persisted).await {
                    warn!("Failed to persist execution: {}", e);
                } else {
                    debug!("Persisted execution {} to KG", task.id);
                }

                if let Err(e) = kg_persistence
                    .record_event(
                        "task_completed",
                        agent_id,
                        "duration_ms",
                        result.duration_ms as f64,
                    )
                    .await
                {
                    warn!("Failed to record event: {}", e);
                }

                if let Err(e) = kg_persistence
                    .record_event(
                        "token_usage",
                        agent_id,
                        "tokens",
                        (result.input_tokens + result.output_tokens) as f64,
                    )
                    .await
                {
                    warn!("Failed to record token event: {}", e);
                }
            } else {
                warn!(
                    "KG persistence available but no embedding provider for task {}",
                    task.id
                );
            }
        }
    }
}

/// Internal state for executor result handling
pub enum ExecutorState {
    Completed(Agent<super::state_machine::Completed>),
    Failed(Agent<super::state_machine::Failed>, String),
}

/// Build the user-facing LLM prompt from a task description and its JSON
/// context.
///
/// The `context` field in [`TaskAssignment`] carries supplementary JSON that
/// may be empty or contain only `"{}"` when no additional context is available.
/// Those cases are skipped so the model receives a clean description-only
/// prompt.
fn build_prompt(description: &str, context: &str) -> String {
    let ctx = context.trim();
    if ctx.is_empty() || ctx == "{}" {
        description.to_string()
    } else {
        format!("{description}\n\nContext:\n{ctx}")
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::registry::AgentStatus;

    #[tokio::test]
    async fn test_executor_creation() {
        let metadata = AgentMetadata {
            id: "test-executor".to_string(),
            stable_id: "developer".to_string(),
            role: "developer".to_string(),
            name: "Test Executor".to_string(),
            version: "0.1.0".to_string(),
            status: AgentStatus::Active,
            capabilities: vec!["code_generation".to_string()],
            llm_provider: "claude".to_string(),
            llm_model: "claude-sonnet-4".to_string(),
            max_concurrent_tasks: 3,
            current_tasks: 0,
            created_at: Utc::now(),
            last_heartbeat: Utc::now(),
            uptime_percentage: 100.0,
            total_tasks_completed: 0,
            system_prompt: None,
        };

        let (_tx, rx) = mpsc::channel(10);
        let executor = AgentExecutor::new(metadata.clone(), rx);

        assert_eq!(executor.agent.metadata.id, "test-executor");
        assert!(executor.kg_persistence.is_none());
        assert!(executor.router.is_none());
    }

    #[test]
    fn test_executor_persistence_disabled_by_default() {
        let metadata = AgentMetadata {
            id: "test-no-persist".to_string(),
            stable_id: "reviewer".to_string(),
            role: "reviewer".to_string(),
            name: "Test No Persist".to_string(),
            version: "0.1.0".to_string(),
            status: AgentStatus::Active,
            capabilities: vec!["review".to_string()],
            llm_provider: "openai".to_string(),
            llm_model: "gpt-4".to_string(),
            max_concurrent_tasks: 5,
            current_tasks: 0,
            created_at: Utc::now(),
            last_heartbeat: Utc::now(),
            uptime_percentage: 99.5,
            total_tasks_completed: 100,
            system_prompt: None,
        };

        let (_tx, rx) = mpsc::channel(10);
        let executor = AgentExecutor::new(metadata, rx);

        assert!(!executor.agent.metadata.role.is_empty());
        assert!(executor.embedding_provider.is_none());
        assert!(executor.router.is_none());
    }

    #[test]
    fn build_prompt_no_context() {
        let p = build_prompt("Fix the bug", "");
        assert_eq!(p, "Fix the bug");
    }

    #[test]
    fn build_prompt_empty_json_context() {
        let p = build_prompt("Fix the bug", "{}");
        assert_eq!(p, "Fix the bug");
    }

    #[test]
    fn build_prompt_with_context() {
        let p = build_prompt("Fix the bug", r#"{"file":"src/lib.rs"}"#);
        assert!(p.contains("Fix the bug"));
        assert!(p.contains("src/lib.rs"));
        assert!(p.contains("Context:"));
    }
}
