use std::sync::Arc;
use std::time::Duration;

use async_trait::async_trait;
use tokio::sync::{Mutex, OwnedMutexGuard, watch};
use tokio::time::sleep;

use crate::gateway::session::Session;
use crate::storage::SessionStore;
use crate::tools::task::SubagentResult;
use crate::tools::{WaitCoordinator, WaitEvent};

/// 基于 Session 的 wait 协调器实现。
///
/// 持有 serial_lock 的 guard slot（可通过 take/drop 释放锁，通过 put 回填新 guard），
/// 以及 Session 引用（用于管理 sub_done_receiver、wait_wakeup、waiting_flag）。
///
/// wait() 流程：
/// 1. 设置 waiting=true（让 process_one 走注入路径）
/// 2. 取出 sub_done_receiver（从 Session 中 take，select! 期间不持有 Session 锁）
/// 3. 获取 wait_wakeup（Arc<Notify>，clone 后不持有 Session 锁）
/// 4. 释放 serial_lock（从 guard_slot take 并 drop guard）
/// 5. select! { sub_done_q.recv(), wakeup.notified(), timeout }
/// 6. 重新获取 serial_lock（serial_lock.lock_owned().await）
/// 7. 回填 guard 到 guard_slot
/// 8. 设置 waiting=false（先获取锁后清除，避免 TOCTOU）
/// 9. 归还 sub_done_receiver
pub struct SessionWaitCoordinator {
    /// Session 引用（Arc<Mutex<Session>>），用于访问 SessionHistory 的队列和状态
    session: Arc<Mutex<Session>>,
    /// serial_lock guard 的存储槽。
    /// 执行路径（execution.rs）获取锁后将 guard 存入此槽；
    /// wait() 取出并 drop 以释放锁，重获取后回填新 guard。
    guard_slot: Arc<Mutex<Option<OwnedMutexGuard<()>>>>,
    /// serial_lock 本体（Arc<tokio::sync::Mutex<()>>），用于重获取锁
    serial_lock: Arc<tokio::sync::Mutex<()>>,
    /// SessionStore 引用，用于查询 pending_subagents
    store: Arc<SessionStore>,
    /// 当前 topic_id
    topic_id: String,
}

impl SessionWaitCoordinator {
    pub fn new(
        session: Arc<Mutex<Session>>,
        guard_slot: Arc<Mutex<Option<OwnedMutexGuard<()>>>>,
        serial_lock: Arc<tokio::sync::Mutex<()>>,
        store: Arc<SessionStore>,
        topic_id: String,
    ) -> Self {
        Self {
            session,
            guard_slot,
            serial_lock,
            store,
            topic_id,
        }
    }
}

#[async_trait]
impl WaitCoordinator for SessionWaitCoordinator {
    fn query_pending_task_ids(&self) -> Vec<String> {
        self.store
            .list_pending_subagents(&self.topic_id, Some("running"))
            .map(|records| records.into_iter().map(|r| r.task_id).collect())
            .unwrap_or_default()
    }

    async fn try_drain_queued_results(&self) -> Vec<SubagentResult> {
        // 取出 receiver → try_recv 排空 → 归还 receiver
        // 安全性：此方法在执行路径中被调用（serial_lock 已持有），
        // 无其他代码并发访问 receiver。
        let rx = {
            let mut session = self.session.lock().await;
            session.take_sub_done_receiver(&self.topic_id)
        };

        let mut results = Vec::new();
        if let Some(mut rx) = rx {
            while let Ok(result) = rx.try_recv() {
                results.push(result);
            }
            // 归还 receiver（即使已排空，仍需放回供后续 wait() 使用）
            let mut session = self.session.lock().await;
            session.restore_sub_done_receiver(&self.topic_id, rx);
        }

        if !results.is_empty() {
            tracing::debug!(
                topic_id = %self.topic_id,
                drained_count = results.len(),
                "Drained buffered subagent results from sub_done_q"
            );
        }

        results
    }

    async fn wait(&self, timeout: Duration, cancel_rx: Option<watch::Receiver<()>>) -> WaitEvent {
        // 1. 设置 waiting=true
        {
            let mut session = self.session.lock().await;
            session.set_waiting(&self.topic_id, true);
        }

        // 2. 取出 sub_done_receiver（select! 期间不持有 Session 锁）
        let receiver = {
            let mut session = self.session.lock().await;
            session.take_sub_done_receiver(&self.topic_id)
        };

        // 3. 获取 wait_wakeup（Arc<Notify>，clone 后不持有 Session 锁）
        let wakeup = {
            let mut session = self.session.lock().await;
            session.wait_wakeup(&self.topic_id)
        };

        // 3.5. 记录等待前的用户消息数量（用于 wakeup 后提取新注入的消息）
        //     直接从 SQLite 读取，不持有任何锁
        let user_msg_count_before = self
            .store
            .load_messages_for_topic(&self.topic_id, None)
            .map(|msgs| msgs.iter().filter(|m| m.role == "user").count())
            .unwrap_or(0);

        // 4. 释放 serial_lock（取出 guard 并 drop）
        {
            let mut slot = self.guard_slot.lock().await;
            let _ = slot.take(); // drop guard → 释放 serial_lock
        }

        tracing::debug!(
            topic_id = %self.topic_id,
            timeout_secs = timeout.as_secs(),
            user_msg_count_before,
            has_cancel_rx = cancel_rx.is_some(),
            "SessionWaitCoordinator: lock released, entering select!"
        );

        // 5. select! 等待（不持有任何锁）
        //
        //    cancel 分支放在最后（biased 排序中最后被 poll），
        //    确保子代理结果和用户消息优先于取消信号被处理。
        //    场景：/stop 后用户立即发消息 → process_one 注入消息 + wakeup，
        //    select! 优先消费 wakeup（UserMessage），而非 cancel（Cancelled），
        //    使已注入的用户消息能被 Agent 处理而非丢失。
        //
        //    但 cancel_rx.changed() 不会无限阻塞——若无子代理结果、无用户消息，
        //    cancel 仍是唯一就绪分支，等待被优雅终止。
        let event = if let Some(mut rx) = receiver {
            let mut cancel_rx = cancel_rx;
            tokio::select! {
                biased;
                result = rx.recv() => {
                    match result {
                        Some(subagent_result) => {
                            let event = WaitEvent::SubagentResult(subagent_result);
                            let mut session = self.session.lock().await;
                            session.restore_sub_done_receiver(&self.topic_id, rx);
                            event
                        }
                        None => {
                            // sender 全部 drop（所有 sub_done_sender 被释放）
                            WaitEvent::Timeout
                        }
                    }
                }
                _ = wakeup.notified() => {
                    // 用户消息到达（process_one 已注入 history 并 wakeup）
                    let mut session = self.session.lock().await;
                    session.restore_sub_done_receiver(&self.topic_id, rx);
                    // 提取等待期间新注入的用户消息内容
                    let new_messages = self.fetch_new_user_messages(user_msg_count_before);
                    tracing::info!(
                        topic_id = %self.topic_id,
                        new_msg_count = new_messages.len(),
                        "SessionWaitCoordinator: woke up by user message"
                    );
                    WaitEvent::UserMessage(new_messages)
                }
                _ = async {
                    if let Some(ref mut crx) = cancel_rx {
                        let _ = crx.changed().await;
                    } else {
                        std::future::pending::<()>().await;
                    }
                } => {
                    // 取消信号到达（/stop）→ 归还 receiver，返回 Cancelled
                    tracing::info!(
                        topic_id = %self.topic_id,
                        "SessionWaitCoordinator: cancelled by /stop during wait"
                    );
                    let mut session = self.session.lock().await;
                    session.restore_sub_done_receiver(&self.topic_id, rx);
                    WaitEvent::Cancelled
                }
                _ = sleep(timeout) => {
                    let mut session = self.session.lock().await;
                    session.restore_sub_done_receiver(&self.topic_id, rx);
                    WaitEvent::Timeout
                }
            }
        } else {
            // 无 receiver（topic 无 sub_done 队列），直接等待 timeout 或 wakeup
            let mut cancel_rx = cancel_rx;
            tokio::select! {
                biased;
                _ = wakeup.notified() => {
                    let new_messages = self.fetch_new_user_messages(user_msg_count_before);
                    tracing::info!(
                        topic_id = %self.topic_id,
                        new_msg_count = new_messages.len(),
                        "SessionWaitCoordinator: woke up by user message (no receiver)"
                    );
                    WaitEvent::UserMessage(new_messages)
                }
                _ = async {
                    if let Some(ref mut crx) = cancel_rx {
                        let _ = crx.changed().await;
                    } else {
                        std::future::pending::<()>().await;
                    }
                } => {
                    tracing::info!(
                        topic_id = %self.topic_id,
                        "SessionWaitCoordinator: cancelled by /stop during wait (no receiver)"
                    );
                    WaitEvent::Cancelled
                }
                _ = sleep(timeout) => WaitEvent::Timeout,
            }
        };

        // 6. 重新获取 serial_lock
        //    lock_owned 消费 Arc<Self>，需 clone 保留 self.serial_lock 供后续可能的重入
        //
        //    取消场景下此处可能阻塞——如果 process_one 正持有锁注入用户消息，
        //    需等其释放后才能重获取。这是正确行为：确保 is_waiting 清除与
        //    process_one 的注入互斥，避免 TOCTOU。
        let new_guard = self.serial_lock.clone().lock_owned().await;

        // 7. 回填 guard 到 guard_slot
        {
            let mut slot = self.guard_slot.lock().await;
            *slot = Some(new_guard);
        }

        // 8. 设置 waiting=false（先获取锁后清除，避免 TOCTOU）
        //    注意：serial_lock 已在步骤 6 获取，此时 process_one 无法获取锁，
        //    所以清除 waiting 是安全的。
        {
            let mut session = self.session.lock().await;
            session.set_waiting(&self.topic_id, false);
        }

        tracing::debug!(
            topic_id = %self.topic_id,
            "SessionWaitCoordinator: lock reacquired, waiting cleared"
        );

        event
    }
}

impl SessionWaitCoordinator {
    /// 提取等待期间新注入的用户消息内容。
    /// 通过对比等待前的用户消息数量，从 SQLite 中取出新增的用户消息。
    fn fetch_new_user_messages(&self, count_before: usize) -> Vec<String> {
        match self.store.load_messages_for_topic(&self.topic_id, None) {
            Ok(msgs) => msgs
                .iter()
                .filter(|m| m.role == "user")
                .skip(count_before)
                .map(|m| m.content.clone())
                .collect(),
            Err(e) => {
                tracing::warn!(
                    error = %e,
                    topic_id = %self.topic_id,
                    "Failed to load messages for fetching new user messages"
                );
                Vec::new()
            }
        }
    }
}