use crate::event::Event;
use std::collections::{HashMap, BinaryHeap};
use std::sync::Arc;
use tokio::sync::RwLock;
use tokio::task::JoinHandle;
use tracing::{debug, info, warn};
use uuid::Uuid;
use chrono::{DateTime, Utc};

/// Task priority (higher = more priority)
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct TaskPriority(pub u8);

impl Default for TaskPriority {
    fn default() -> Self {
        TaskPriority(5)
    }
}

/// Task for concurrent execution
#[derive(Debug, Clone)]
pub struct ScheduledTask {
    pub id: Uuid,
    pub name: String,
    pub priority: TaskPriority,
    pub scheduled_at: DateTime<Utc>,
    pub event: Option<Event>,
}

impl PartialEq for ScheduledTask {
    fn eq(&self, other: &Self) -> bool {
        self.id == other.id
    }
}

impl Eq for ScheduledTask {}

impl PartialOrd for ScheduledTask {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        // Higher priority first, then earlier scheduled time
        match other.priority.0.cmp(&self.priority.0) {
            std::cmp::Ordering::Equal => {
                self.scheduled_at.partial_cmp(&other.scheduled_at)
            }
            other => Some(other),
        }
    }
}

impl Ord for ScheduledTask {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.partial_cmp(other).unwrap_or(std::cmp::Ordering::Equal)
    }
}

/// Premium task scheduler for concurrent agent execution
/// 
/// Features:
/// - Priority-based task scheduling
/// - Concurrent task execution with limits
/// - Task cancellation support
/// - Scheduled task execution
pub struct TaskScheduler {
    /// Priority queue of scheduled tasks
    task_queue: Arc<RwLock<BinaryHeap<ScheduledTask>>>,
    
    /// Active task handles
    active_tasks: Arc<RwLock<HashMap<Uuid, JoinHandle<()>>>>,
    
    /// Maximum concurrent tasks
    max_concurrent: usize,
    
    /// Current number of active tasks
    active_count: Arc<RwLock<usize>>,
}

impl TaskScheduler {
    /// Create a new task scheduler
    pub fn new(max_concurrent: usize) -> Self {
        Self {
            task_queue: Arc::new(RwLock::new(BinaryHeap::new())),
            active_tasks: Arc::new(RwLock::new(HashMap::new())),
            max_concurrent,
            active_count: Arc::new(RwLock::new(0)),
        }
    }
    
    /// Schedule a task for execution
    pub async fn schedule_task(
        &self,
        task_id: Uuid,
        name: String,
        priority: TaskPriority,
        scheduled_at: DateTime<Utc>,
        event: Option<Event>,
    ) {
        let task = ScheduledTask {
            id: task_id,
            name,
            priority,
            scheduled_at,
            event,
        };
        
        let mut queue = self.task_queue.write().await;
        queue.push(task);
        debug!("Scheduled task: {} (priority: {})", task_id, priority.0);
    }
    
    /// Schedule a task to run immediately
    pub async fn schedule_immediate(
        &self,
        task_id: Uuid,
        name: String,
        priority: TaskPriority,
        event: Option<Event>,
    ) {
        self.schedule_task(task_id, name, priority, Utc::now(), event).await;
    }
    
    /// Get next ready task (if any and if we have capacity)
    pub async fn get_next_task(&self) -> Option<ScheduledTask> {
        let active = *self.active_count.read().await;
        if active >= self.max_concurrent {
            return None; // At capacity
        }
        
        let mut queue = self.task_queue.write().await;
        let now = Utc::now();
        
        // Find the first task that's ready to execute
        let mut ready_tasks = Vec::new();
        let mut remaining_tasks = Vec::new();
        
        while let Some(task) = queue.pop() {
            if task.scheduled_at <= now {
                ready_tasks.push(task);
            } else {
                remaining_tasks.push(task);
            }
        }
        
        // Put remaining tasks back
        for task in remaining_tasks {
            queue.push(task);
        }
        
        // Return the highest priority ready task
        ready_tasks.into_iter().max()
    }
    
    /// Register an active task
    pub async fn register_task(&self, task_id: Uuid, handle: JoinHandle<()>) {
        let mut tasks = self.active_tasks.write().await;
        let mut count = self.active_count.write().await;
        
        tasks.insert(task_id, handle);
        *count += 1;
        
        info!("Registered active task: {} (active: {}/{})", task_id, *count, self.max_concurrent);
    }
    
    /// Unregister a completed task
    pub async fn unregister_task(&self, task_id: Uuid) {
        let mut tasks = self.active_tasks.write().await;
        let mut count = self.active_count.write().await;
        
        if tasks.remove(&task_id).is_some() {
            *count = count.saturating_sub(1);
            debug!("Unregistered task: {} (active: {}/{})", task_id, *count, self.max_concurrent);
        }
    }
    
    /// Cancel a task
    pub async fn cancel_task(&self, task_id: Uuid) -> bool {
        // Remove from queue if not started
        let mut queue = self.task_queue.write().await;
        let mut new_queue = BinaryHeap::new();
        let mut found = false;
        
        while let Some(task) = queue.pop() {
            if task.id == task_id {
                found = true;
            } else {
                new_queue.push(task);
            }
        }
        
        *queue = new_queue;
        
        // Abort if already running
        let mut tasks = self.active_tasks.write().await;
        if let Some(handle) = tasks.remove(&task_id) {
            handle.abort();
            let mut count = self.active_count.write().await;
            *count = count.saturating_sub(1);
            found = true;
        }
        
        if found {
            info!("Cancelled task: {}", task_id);
        }
        
        found
    }
    
    /// Get number of active tasks
    pub async fn active_count(&self) -> usize {
        *self.active_count.read().await
    }
    
    /// Get number of queued tasks
    pub async fn queued_count(&self) -> usize {
        self.task_queue.read().await.len()
    }
    
    /// Get maximum concurrent tasks
    pub fn max_concurrent(&self) -> usize {
        self.max_concurrent
    }
}

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

