//! Orchestrator Integration for RAG System (Phase 8b)
//!
//! Integrates RAG system with provisioning orchestrator service for:
//! - Async batch job submission
//! - Conversation management as orchestrator tasks
//! - Tool execution coordination
//! - Result polling and progress tracking
//! - Health monitoring

use std::sync::Arc;

use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use uuid::Uuid;

use crate::{
    batch_processing::BatchJob,
    error::{RagError, Result},
    tools::{ToolInput, ToolOutput},
};

/// Orchestrator task types supported by RAG system
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum OrchestratorTaskType {
    /// Batch query processing task
    BatchQuery,
    /// Conversation management task
    Conversation,
    /// Tool execution task
    ToolExecution,
    /// Health check task
    HealthCheck,
}

/// Task status in orchestrator
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum TaskStatus {
    /// Task submitted, waiting to start
    Pending,
    /// Task is currently executing
    Running,
    /// Task completed successfully
    Completed,
    /// Task failed
    Failed,
    /// Task was cancelled
    Cancelled,
}

/// Orchestrator task for batch query processing
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BatchQueryTask {
    /// Unique task ID
    pub task_id: String,
    /// Associated batch job ID
    pub job_id: String,
    /// Queries to process
    pub queries: Vec<String>,
    /// Maximum concurrent queries
    pub max_concurrent: usize,
    /// Timeout in seconds
    pub timeout_secs: u64,
    /// Current task status
    pub status: TaskStatus,
    /// Number of completed queries
    pub completed: usize,
    /// Number of failed queries
    pub failed: usize,
}

impl BatchQueryTask {
    /// Create a new batch query task
    pub fn new(job_id: String, batch_job: &BatchJob) -> Self {
        Self {
            task_id: Uuid::new_v4().to_string(),
            job_id,
            queries: batch_job.queries.iter().map(|q| q.query.clone()).collect(),
            max_concurrent: batch_job.max_concurrent,
            timeout_secs: batch_job.timeout_secs,
            status: TaskStatus::Pending,
            completed: 0,
            failed: 0,
        }
    }

    /// Update task progress
    pub fn update_progress(&mut self, completed: usize, failed: usize) {
        self.completed = completed;
        self.failed = failed;
    }

    /// Mark task as running
    pub fn start(&mut self) {
        self.status = TaskStatus::Running;
    }

    /// Mark task as completed
    pub fn complete(&mut self) {
        self.status = TaskStatus::Completed;
    }

    /// Mark task as failed
    pub fn fail(&mut self) {
        self.status = TaskStatus::Failed;
    }

    /// Get progress percentage
    pub fn progress_percent(&self) -> u8 {
        if self.queries.is_empty() {
            0
        } else {
            ((self.completed as f32 / self.queries.len() as f32) * 100.0) as u8
        }
    }
}

/// Orchestrator task for conversation management
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConversationTask {
    /// Unique task ID
    pub task_id: String,
    /// Conversation ID
    pub conversation_id: String,
    /// Current message
    pub message: String,
    /// Task status
    pub status: TaskStatus,
    /// Generated response (once completed)
    pub response: Option<String>,
    /// Turn count in conversation
    pub turn_count: usize,
}

impl ConversationTask {
    /// Create a new conversation task
    pub fn new(conversation_id: String, message: String) -> Self {
        Self {
            task_id: Uuid::new_v4().to_string(),
            conversation_id,
            message,
            status: TaskStatus::Pending,
            response: None,
            turn_count: 0,
        }
    }

    /// Mark task as running
    pub fn start(&mut self) {
        self.status = TaskStatus::Running;
    }

    /// Complete task with response
    pub fn complete_with_response(&mut self, response: String) {
        self.response = Some(response);
        self.status = TaskStatus::Completed;
    }

    /// Mark task as failed
    pub fn fail(&mut self) {
        self.status = TaskStatus::Failed;
    }
}

/// Orchestrator task for tool execution
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolExecutionTask {
    /// Unique task ID
    pub task_id: String,
    /// Tool ID to execute
    pub tool_id: String,
    /// Tool input parameters
    pub input: ToolInput,
    /// Task status
    pub status: TaskStatus,
    /// Tool execution result
    pub result: Option<ToolOutput>,
}

impl ToolExecutionTask {
    /// Create a new tool execution task
    pub fn new(tool_id: String, input: ToolInput) -> Self {
        Self {
            task_id: Uuid::new_v4().to_string(),
            tool_id,
            input,
            status: TaskStatus::Pending,
            result: None,
        }
    }

    /// Mark task as running
    pub fn start(&mut self) {
        self.status = TaskStatus::Running;
    }

    /// Complete task with result
    pub fn complete_with_result(&mut self, result: ToolOutput) {
        self.result = Some(result);
        self.status = TaskStatus::Completed;
    }

    /// Mark task as failed
    pub fn fail(&mut self) {
        self.status = TaskStatus::Failed;
    }
}

/// Unified orchestrator task
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum OrchestratorTask {
    BatchQuery(BatchQueryTask),
    Conversation(ConversationTask),
    ToolExecution(ToolExecutionTask),
}

impl OrchestratorTask {
    /// Get task ID
    pub fn task_id(&self) -> &str {
        match self {
            OrchestratorTask::BatchQuery(t) => &t.task_id,
            OrchestratorTask::Conversation(t) => &t.task_id,
            OrchestratorTask::ToolExecution(t) => &t.task_id,
        }
    }

    /// Get task status
    pub fn status(&self) -> &TaskStatus {
        match self {
            OrchestratorTask::BatchQuery(t) => &t.status,
            OrchestratorTask::Conversation(t) => &t.status,
            OrchestratorTask::ToolExecution(t) => &t.status,
        }
    }

    /// Get task type
    pub fn task_type(&self) -> OrchestratorTaskType {
        match self {
            OrchestratorTask::BatchQuery(_) => OrchestratorTaskType::BatchQuery,
            OrchestratorTask::Conversation(_) => OrchestratorTaskType::Conversation,
            OrchestratorTask::ToolExecution(_) => OrchestratorTaskType::ToolExecution,
        }
    }
}

/// Orchestrator task manager
pub struct OrchestratorManager {
    /// In-memory task store (in production, would persist to database)
    tasks: Arc<RwLock<std::collections::HashMap<String, OrchestratorTask>>>,
}

impl OrchestratorManager {
    /// Create a new orchestrator manager
    pub fn new() -> Self {
        Self {
            tasks: Arc::new(RwLock::new(std::collections::HashMap::new())),
        }
    }

    /// Submit a batch query task
    pub async fn submit_batch_task(&self, batch_job: &BatchJob) -> Result<String> {
        let job_id = batch_job.job_id.clone();
        let task = BatchQueryTask::new(job_id, batch_job);
        let task_id = task.task_id.clone();

        let mut tasks = self.tasks.write().await;
        tasks.insert(task_id.clone(), OrchestratorTask::BatchQuery(task));

        Ok(task_id)
    }

    /// Submit a conversation task
    pub async fn submit_conversation_task(
        &self,
        conversation_id: String,
        message: String,
    ) -> Result<String> {
        let task = ConversationTask::new(conversation_id, message);
        let task_id = task.task_id.clone();

        let mut tasks = self.tasks.write().await;
        tasks.insert(task_id.clone(), OrchestratorTask::Conversation(task));

        Ok(task_id)
    }

    /// Submit a tool execution task
    pub async fn submit_tool_task(&self, tool_id: String, input: ToolInput) -> Result<String> {
        let task = ToolExecutionTask::new(tool_id, input);
        let task_id = task.task_id.clone();

        let mut tasks = self.tasks.write().await;
        tasks.insert(task_id.clone(), OrchestratorTask::ToolExecution(task));

        Ok(task_id)
    }

    /// Get task status
    pub async fn get_task_status(&self, task_id: &str) -> Result<TaskStatus> {
        let tasks = self.tasks.read().await;
        tasks
            .get(task_id)
            .map(|t| t.status().clone())
            .ok_or_else(|| RagError::NotFound(format!("Task {} not found", task_id)))
    }

    /// Get task details
    pub async fn get_task(&self, task_id: &str) -> Result<OrchestratorTask> {
        let tasks = self.tasks.read().await;
        tasks
            .get(task_id)
            .cloned()
            .ok_or_else(|| RagError::NotFound(format!("Task {} not found", task_id)))
    }

    /// Update batch task progress
    pub async fn update_batch_progress(
        &self,
        task_id: &str,
        completed: usize,
        failed: usize,
    ) -> Result<()> {
        let mut tasks = self.tasks.write().await;
        if let Some(OrchestratorTask::BatchQuery(task)) = tasks.get_mut(task_id) {
            task.update_progress(completed, failed);
            Ok(())
        } else {
            Err(RagError::NotFound(format!(
                "Batch task {} not found",
                task_id
            )))
        }
    }

    /// Mark batch task as completed
    pub async fn complete_batch_task(&self, task_id: &str) -> Result<()> {
        let mut tasks = self.tasks.write().await;
        if let Some(OrchestratorTask::BatchQuery(task)) = tasks.get_mut(task_id) {
            task.complete();
            Ok(())
        } else {
            Err(RagError::NotFound(format!(
                "Batch task {} not found",
                task_id
            )))
        }
    }

    /// Complete conversation task with response
    pub async fn complete_conversation_task(&self, task_id: &str, response: String) -> Result<()> {
        let mut tasks = self.tasks.write().await;
        if let Some(OrchestratorTask::Conversation(task)) = tasks.get_mut(task_id) {
            task.complete_with_response(response);
            Ok(())
        } else {
            Err(RagError::NotFound(format!(
                "Conversation task {} not found",
                task_id
            )))
        }
    }

    /// Complete tool execution task with result
    pub async fn complete_tool_task(&self, task_id: &str, result: ToolOutput) -> Result<()> {
        let mut tasks = self.tasks.write().await;
        if let Some(OrchestratorTask::ToolExecution(task)) = tasks.get_mut(task_id) {
            task.complete_with_result(result);
            Ok(())
        } else {
            Err(RagError::NotFound(format!(
                "Tool task {} not found",
                task_id
            )))
        }
    }

    /// List all tasks
    pub async fn list_tasks(&self) -> Result<Vec<OrchestratorTask>> {
        let tasks = self.tasks.read().await;
        Ok(tasks.values().cloned().collect())
    }

    /// List tasks by status
    pub async fn list_tasks_by_status(&self, status: TaskStatus) -> Result<Vec<OrchestratorTask>> {
        let tasks = self.tasks.read().await;
        Ok(tasks
            .values()
            .filter(|t| t.status() == &status)
            .cloned()
            .collect())
    }

    /// Get statistics
    pub async fn get_stats(&self) -> Result<OrchestratorStats> {
        let tasks = self.tasks.read().await;

        let total = tasks.len();
        let pending = tasks
            .values()
            .filter(|t| t.status() == &TaskStatus::Pending)
            .count();
        let running = tasks
            .values()
            .filter(|t| t.status() == &TaskStatus::Running)
            .count();
        let completed = tasks
            .values()
            .filter(|t| t.status() == &TaskStatus::Completed)
            .count();
        let failed = tasks
            .values()
            .filter(|t| t.status() == &TaskStatus::Failed)
            .count();

        Ok(OrchestratorStats {
            total_tasks: total,
            pending,
            running,
            completed,
            failed,
            success_rate: if completed + failed > 0 {
                (completed as f32 / (completed + failed) as f32) * 100.0
            } else {
                0.0
            },
        })
    }
}

impl Default for OrchestratorManager {
    fn default() -> Self {
        Self::new()
    }
}

/// Orchestrator statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OrchestratorStats {
    pub total_tasks: usize,
    pub pending: usize,
    pub running: usize,
    pub completed: usize,
    pub failed: usize,
    pub success_rate: f32,
}

/// Health status for orchestrator
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OrchestratorHealth {
    pub status: String,
    pub tasks_pending: usize,
    pub tasks_running: usize,
    pub uptime_seconds: u64,
}

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

    #[test]
    fn test_batch_query_task_creation() {
        let batch_job = BatchJob::new(vec![BatchQuery::new("test".into())]);
        let task = BatchQueryTask::new(batch_job.job_id.clone(), &batch_job);

        assert_eq!(task.status, TaskStatus::Pending);
        assert_eq!(task.queries.len(), 1);
        assert_eq!(task.completed, 0);
    }

    #[test]
    fn test_batch_query_task_progress() {
        let batch_job = BatchJob::new(vec![
            BatchQuery::new("q1".into()),
            BatchQuery::new("q2".into()),
            BatchQuery::new("q3".into()),
        ]);
        let mut task = BatchQueryTask::new(batch_job.job_id.clone(), &batch_job);

        task.update_progress(2, 0);
        assert_eq!(task.completed, 2);
        assert_eq!(task.progress_percent(), 66);
    }

    #[test]
    fn test_conversation_task_creation() {
        let task = ConversationTask::new("conv-1".into(), "Hello".into());

        assert_eq!(task.status, TaskStatus::Pending);
        assert_eq!(task.response, None);
    }

    #[test]
    fn test_conversation_task_completion() {
        let mut task = ConversationTask::new("conv-1".into(), "Hello".into());
        task.complete_with_response("Hi there!".into());

        assert_eq!(task.status, TaskStatus::Completed);
        assert_eq!(task.response, Some("Hi there!".into()));
    }

    #[test]
    fn test_tool_execution_task_creation() {
        let input = ToolInput {
            params: serde_json::json!({}),
        };
        let task = ToolExecutionTask::new("create_server".into(), input);

        assert_eq!(task.status, TaskStatus::Pending);
        assert_eq!(task.tool_id, "create_server");
    }

    #[test]
    fn test_orchestrator_task_types() {
        let batch_job = BatchJob::new(vec![BatchQuery::new("test".into())]);
        let batch_task = BatchQueryTask::new(batch_job.job_id.clone(), &batch_job);
        let orch_task = OrchestratorTask::BatchQuery(batch_task);

        assert_eq!(orch_task.task_type(), OrchestratorTaskType::BatchQuery);
    }

    #[tokio::test]
    async fn test_orchestrator_manager_submit_batch() {
        let manager = OrchestratorManager::new();
        let batch_job = BatchJob::new(vec![BatchQuery::new("test".into())]);

        let task_id = manager
            .submit_batch_task(&batch_job)
            .await
            .expect("submit failed");

        assert!(!task_id.is_empty());

        let status = manager
            .get_task_status(&task_id)
            .await
            .expect("status failed");
        assert_eq!(status, TaskStatus::Pending);
    }

    #[tokio::test]
    async fn test_orchestrator_manager_list_tasks() {
        let manager = OrchestratorManager::new();
        let batch_job = BatchJob::new(vec![BatchQuery::new("test".into())]);

        let _ = manager.submit_batch_task(&batch_job).await;

        let tasks = manager.list_tasks().await.expect("list failed");
        assert_eq!(tasks.len(), 1);
    }

    #[tokio::test]
    async fn test_orchestrator_manager_stats() {
        let manager = OrchestratorManager::new();
        let batch_job = BatchJob::new(vec![BatchQuery::new("test".into())]);

        let task_id = manager
            .submit_batch_task(&batch_job)
            .await
            .expect("submit failed");
        let _ = manager.complete_batch_task(&task_id).await;

        let stats = manager.get_stats().await.expect("stats failed");
        assert_eq!(stats.total_tasks, 1);
        assert_eq!(stats.completed, 1);
    }
}
