mofa-foundation 0.1.1

use mofa_kernel::message::{AgentEvent, AgentMessage, SchedulingStatus, TaskPriority, TaskRequest};
use mofa_kernel::{AgentBus, CommunicationMode};
use std::collections::{BinaryHeap, HashMap};
use std::sync::Arc;
use tokio::sync::RwLock;

// 带优先级的任务包装器（BinaryHeap 默认是最大堆，直接支持优先级排序）
#[derive(Debug, Clone, Eq, PartialEq)]
struct PriorityTask {
    priority: TaskPriority,
    task: TaskRequest,
    submit_time: std::time::Instant,
}

impl Ord for PriorityTask {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.priority
            .cmp(&other.priority)
            .then_with(|| other.submit_time.cmp(&self.submit_time)) // 同优先级先提交先执行
    }
}

impl PartialOrd for PriorityTask {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

// 优先级调度器
pub struct PriorityScheduler {
    task_queue: Arc<RwLock<BinaryHeap<PriorityTask>>>, // 优先级任务队列
    agent_load: Arc<RwLock<HashMap<String, usize>>>,   // 智能体当前负载（执行中的任务数）
    bus: Arc<AgentBus>,
    task_status: Arc<RwLock<HashMap<String, SchedulingStatus>>>, // 任务状态跟踪
    role_mapping: Arc<RwLock<HashMap<String, Vec<String>>>>,     // 角色-智能体映射
}

impl PriorityScheduler {
    pub async fn new(bus: Arc<AgentBus>) -> Self {
        Self {
            task_queue: Arc::new(RwLock::new(BinaryHeap::new())),
            agent_load: Arc::new(RwLock::new(HashMap::new())),
            bus,
            task_status: Arc::new(RwLock::new(HashMap::new())),
            role_mapping: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// 1. 提交任务到优先级队列
    pub async fn submit_task(&self, task: TaskRequest) -> anyhow::Result<()> {
        let priority_task = PriorityTask {
            priority: task.priority.clone(),
            task: task.clone(),
            submit_time: std::time::Instant::now(),
        };
        self.task_queue.write().await.push(priority_task);
        self.task_status
            .write()
            .await
            .insert(task.task_id, SchedulingStatus::Pending);
        // 提交后立即触发调度
        self.schedule().await?;
        Ok(())
    }

    /// 2. 核心调度逻辑：选高优先级任务 + 选负载最低的智能体
    pub async fn schedule(&self) -> anyhow::Result<()> {
        let mut task_queue = self.task_queue.write().await;
        let mut agent_load = self.agent_load.write().await;
        let mut task_status = self.task_status.write().await;

        while let Some(priority_task) = task_queue.pop() {
            let task = priority_task.task.clone(); // Clone instead of moving
            let task_id = task.task_id.clone();

            // 跳过已被处理的任务
            if task_status.get(&task_id) != Some(&SchedulingStatus::Pending) {
                continue;
            }

            // 选择负载最低的可用智能体（同角色内）
            let target_agent = self.select_low_load_agent("worker").await?;
            if target_agent.is_empty() {
                // 无可用智能体，重新入队
                task_queue.push(priority_task);
                break;
            }
            let target_agent = target_agent[0].clone();

            // 检查是否需要抢占：如果目标智能体有低优先级任务在运行
            self.preempt_low_priority_task(&target_agent, &task).await?;

            // 发送任务给目标智能体
            // Convert TaskRequest to AgentMessage::TaskRequest variant
            let task_msg = AgentMessage::TaskRequest {
                task_id: task.task_id.clone(),
                content: task.content.clone(),
            };
            self.bus
                .send_message(
                    "scheduler",
                    CommunicationMode::PointToPoint(target_agent.clone()),
                    &task_msg,
                )
                .await?;

            // 更新状态和负载
            task_status.insert(task_id, SchedulingStatus::Running);
            *agent_load.entry(target_agent).or_insert(0) += 1;
        }
        Ok(())
    }

    /// 3. 负载均衡：选择同角色内负载最低的智能体
    async fn select_low_load_agent(&self, role: &str) -> anyhow::Result<Vec<String>> {
        let role_map = self.role_mapping.read().await;
        let agents = role_map
            .get(role)
            .ok_or_else(|| anyhow::anyhow!("No agent for role: {}", role))?;
        let agent_load = self.agent_load.read().await;

        // 按负载升序排序，取负载最低的
        let mut sorted_agents = agents.clone();
        sorted_agents.sort_by_key(|agent_id| agent_load.get(agent_id).cloned().unwrap_or(0));
        Ok(sorted_agents)
    }

    /// 4. 任务抢占：高优先级任务抢占低优先级任务的执行资源
    async fn preempt_low_priority_task(
        &self,
        agent_id: &str,
        _high_priority_task: &TaskRequest,
    ) -> anyhow::Result<()> {
        // 简化实现：直接发送抢占事件
        let agent_load = self.agent_load.read().await;
        let task_status = self.task_status.read().await;

        // 检查目标智能体当前运行的任务
        if let Some(&load) = agent_load.get(agent_id)
            && load > 0
        {
            // 找到该智能体上优先级最低的运行中任务
            let low_priority_task_id = task_status
                .iter()
                .find(|(_, status_ref)| **status_ref == SchedulingStatus::Running)
                .map(|(task_id, _)| task_id.clone())
                .ok_or_else(|| anyhow::anyhow!("No running task on agent: {}", agent_id))?;

            // 发送抢占指令，标记低优先级任务为 Preempted
            let preempt_msg =
                AgentMessage::Event(AgentEvent::TaskPreempted(low_priority_task_id.clone()));
            self.bus
                .send_message(
                    "scheduler",
                    CommunicationMode::PointToPoint(agent_id.to_string()),
                    &preempt_msg,
                )
                .await?;
        }
        Ok(())
    }

    /// 5. 任务完成后更新状态和负载
    pub async fn on_task_completed(&self, agent_id: &str, task_id: &str) -> anyhow::Result<()> {
        let mut agent_load = self.agent_load.write().await;
        let mut task_status = self.task_status.write().await;

        agent_load
            .entry(agent_id.to_string())
            .and_modify(|count| *count -= 1);
        task_status.insert(task_id.to_string(), SchedulingStatus::Completed);

        // 任务完成后再次触发调度，处理队列中的下一个任务
        self.schedule().await?;
        Ok(())
    }
}