use std::collections::HashMap;
use std::sync::{Arc, Mutex, Weak};
use std::time::Duration;

use tokio::sync::{Mutex as AsyncMutex, oneshot, watch};
use tokio::time::{Instant, sleep_until, timeout};

use crate::channels::{Channel, ChannelError, LivePolicy, TurnSink, TurnTarget};
use crate::delivery::{PresentationPolicy, project_snapshot};
use crate::session::{TurnSnapshot, TurnStatus};
use crate::task_supervisor::TaskSupervisor;

const SINK_CALL_TIMEOUT: Duration = Duration::from_secs(30);
const FINAL_RETRY_DELAYS: &[Duration] = &[
    Duration::from_secs(1),
    Duration::from_secs(2),
    Duration::from_secs(4),
];

#[derive(Debug)]
pub enum DeliveryError {
    ChannelNotFound(String),
    OpenFailed(ChannelError),
    SnapshotStreamClosed,
    SupervisorStopping,
    FinalTimedOut,
    FinalFailed(ChannelError),
}

impl std::fmt::Display for DeliveryError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::ChannelNotFound(channel) => write!(formatter, "channel not found: {channel}"),
            Self::OpenFailed(error) => write!(formatter, "failed to open turn sink: {error}"),
            Self::SnapshotStreamClosed => {
                formatter.write_str("turn snapshot stream closed before a terminal state")
            }
            Self::SupervisorStopping => {
                formatter.write_str("cannot start turn delivery while Gateway is stopping")
            }
            Self::FinalTimedOut => formatter.write_str("final turn delivery timed out"),
            Self::FinalFailed(error) => write!(formatter, "final turn delivery failed: {error}"),
        }
    }
}

impl std::error::Error for DeliveryError {}

/// Shared ordering boundary for writes to one `(channel, chat_id)` target.
///
/// The registry stores weak references so inactive conversations disappear
/// without a cleanup task. Callers hold the returned lock only around one
/// external write, never for the lifetime of a Turn.
#[derive(Clone, Default)]
pub struct ConversationWriteLocks {
    locks: Arc<Mutex<HashMap<String, Weak<AsyncMutex<()>>>>>,
}

impl ConversationWriteLocks {
    pub fn for_target(&self, channel: &str, chat_id: &str) -> Arc<AsyncMutex<()>> {
        let key = format!("{channel}\0{chat_id}");
        let mut locks = self
            .locks
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        if let Some(existing) = locks.get(&key).and_then(Weak::upgrade) {
            return existing;
        }
        let lock = Arc::new(AsyncMutex::new(()));
        locks.insert(key, Arc::downgrade(&lock));
        lock
    }
}

#[derive(Clone)]
pub struct DeliveryCoordinator {
    write_locks: ConversationWriteLocks,
    sink_call_timeout: Duration,
    final_retry_delays: Arc<[Duration]>,
}

pub(crate) struct SinkRoute {
    pub channel: String,
    pub chat_id: String,
    pub live_policy: LivePolicy,
    pub presentation: PresentationPolicy,
}

impl DeliveryCoordinator {
    pub fn new(write_locks: ConversationWriteLocks) -> Self {
        Self {
            write_locks,
            sink_call_timeout: SINK_CALL_TIMEOUT,
            final_retry_delays: FINAL_RETRY_DELAYS.into(),
        }
    }

    #[cfg(test)]
    fn for_test(
        sink_call_timeout: Duration,
        final_retry_delays: impl Into<Arc<[Duration]>>,
    ) -> Self {
        Self {
            write_locks: ConversationWriteLocks::default(),
            sink_call_timeout,
            final_retry_delays: final_retry_delays.into(),
        }
    }

    pub fn write_locks(&self) -> ConversationWriteLocks {
        self.write_locks.clone()
    }

    pub async fn open_and_deliver(
        &self,
        channel: Arc<dyn Channel + Send + Sync>,
        target: TurnTarget,
        presentation: PresentationPolicy,
        snapshots: watch::Receiver<Arc<TurnSnapshot>>,
    ) -> Result<(), DeliveryError> {
        let live_policy = channel.live_policy();
        let sink = channel
            .open_turn(target.clone())
            .await
            .map_err(DeliveryError::OpenFailed)?;
        self.deliver(
            &target.channel,
            &target.chat_id,
            live_policy,
            presentation,
            snapshots,
            sink,
        )
        .await
    }

    /// Start one sink lifecycle under the Gateway's task owner and return a
    /// bounded completion report to the caller.
    pub fn spawn(
        &self,
        supervisor: &TaskSupervisor,
        channel: Arc<dyn Channel + Send + Sync>,
        target: TurnTarget,
        presentation: PresentationPolicy,
        snapshots: watch::Receiver<Arc<TurnSnapshot>>,
    ) -> Result<oneshot::Receiver<Result<(), DeliveryError>>, DeliveryError> {
        let (result_tx, result_rx) = oneshot::channel();
        let coordinator = self.clone();
        let task_name = format!("turn-delivery:{}:{}", target.channel, target.chat_id);
        let spawned = supervisor.spawn(task_name, async move {
            let result = coordinator
                .open_and_deliver(channel, target, presentation, snapshots)
                .await;
            if let Err(error) = &result {
                tracing::error!(error = %error, "Turn delivery failed");
            }
            let _ = result_tx.send(result);
        });
        if !spawned {
            return Err(DeliveryError::SupervisorStopping);
        }
        Ok(result_rx)
    }

    pub(crate) fn spawn_sink(
        &self,
        supervisor: &TaskSupervisor,
        route: SinkRoute,
        mut snapshots: watch::Receiver<Arc<TurnSnapshot>>,
        mut sink: Box<dyn TurnSink>,
    ) -> Result<oneshot::Receiver<Result<(), DeliveryError>>, DeliveryError> {
        let SinkRoute {
            channel,
            chat_id,
            live_policy,
            presentation,
        } = route;
        let (result_tx, result_rx) = oneshot::channel();
        let coordinator = self.clone();
        let task_name = format!("turn-delivery:{channel}:{chat_id}");
        let cancellation = supervisor.cancellation_token();
        let shutdown_snapshot = snapshots.clone();
        let spawned = supervisor.spawn_graceful(task_name, async move {
            let mut delivery = Box::pin(coordinator.deliver_sink(
                &channel,
                &chat_id,
                live_policy,
                presentation,
                &mut snapshots,
                &mut *sink,
            ));
            let result = tokio::select! {
                result = &mut delivery => result,
                () = cancellation.cancelled() => {
                    drop(delivery);
                    let snapshot = shutdown_snapshot.borrow().clone();
                    let projected = project_snapshot(&snapshot, presentation);
                    coordinator
                        .abort_for_shutdown(&channel, &chat_id, &mut *sink, &projected)
                        .await
                }
            };
            if let Err(error) = &result {
                tracing::error!(channel, chat_id, error = %error, "Turn delivery failed");
            }
            let _ = result_tx.send(result);
        });
        if !spawned {
            return Err(DeliveryError::SupervisorStopping);
        }
        Ok(result_rx)
    }

    pub async fn deliver(
        &self,
        channel: &str,
        chat_id: &str,
        live_policy: LivePolicy,
        presentation: PresentationPolicy,
        mut snapshots: watch::Receiver<Arc<TurnSnapshot>>,
        mut sink: Box<dyn TurnSink>,
    ) -> Result<(), DeliveryError> {
        self.deliver_sink(
            channel,
            chat_id,
            live_policy,
            presentation,
            &mut snapshots,
            &mut *sink,
        )
        .await
    }

    async fn deliver_sink(
        &self,
        channel: &str,
        chat_id: &str,
        live_policy: LivePolicy,
        presentation: PresentationPolicy,
        snapshots: &mut watch::Receiver<Arc<TurnSnapshot>>,
        sink: &mut dyn TurnSink,
    ) -> Result<(), DeliveryError> {
        let target_lock = self.write_locks.for_target(channel, chat_id);
        let min_interval = match live_policy {
            LivePolicy::FinalOnly => None,
            LivePolicy::Snapshot { min_interval } if presentation.live => Some(min_interval),
            LivePolicy::Snapshot { .. } => None,
        };
        let mut next_update_at = Instant::now();

        loop {
            let snapshot = snapshots.borrow_and_update().clone();
            if snapshot.status != TurnStatus::Running {
                let projected = project_snapshot(&snapshot, presentation);
                return self.deliver_terminal(&target_lock, sink, &projected).await;
            }

            if let Some(interval) = min_interval {
                while Instant::now() < next_update_at {
                    tokio::select! {
                        changed = snapshots.changed() => {
                            changed.map_err(|_| DeliveryError::SnapshotStreamClosed)?;
                            let latest = snapshots.borrow_and_update().clone();
                            if latest.status != TurnStatus::Running {
                                let projected = project_snapshot(&latest, presentation);
                                return self.deliver_terminal(&target_lock, sink, &projected).await;
                            }
                        }
                        () = sleep_until(next_update_at) => break,
                    }
                }

                let latest = snapshots.borrow_and_update().clone();
                if latest.status != TurnStatus::Running {
                    let projected = project_snapshot(&latest, presentation);
                    return self.deliver_terminal(&target_lock, sink, &projected).await;
                }
                let projected = project_snapshot(&latest, presentation);
                let _guard = target_lock.lock().await;
                match timeout(self.sink_call_timeout, sink.update(&projected)).await {
                    Ok(Ok(())) => {}
                    Ok(Err(error)) => {
                        tracing::warn!(error = %error, revision = projected.revision, "Live turn update failed; waiting for a newer snapshot");
                    }
                    Err(_) => {
                        tracing::warn!(
                            revision = projected.revision,
                            "Live turn update timed out; waiting for a newer snapshot"
                        );
                    }
                }
                next_update_at = Instant::now() + interval;
            }

            snapshots
                .changed()
                .await
                .map_err(|_| DeliveryError::SnapshotStreamClosed)?;
        }
    }

    async fn abort_for_shutdown(
        &self,
        channel: &str,
        chat_id: &str,
        sink: &mut dyn TurnSink,
        snapshot: &TurnSnapshot,
    ) -> Result<(), DeliveryError> {
        let target_lock = self.write_locks.for_target(channel, chat_id);
        let _guard = target_lock.lock().await;
        match timeout(self.sink_call_timeout, sink.abort(snapshot)).await {
            Ok(Ok(())) => Ok(()),
            Ok(Err(error)) => Err(DeliveryError::FinalFailed(error)),
            Err(_) => Err(DeliveryError::FinalTimedOut),
        }
    }

    async fn deliver_terminal(
        &self,
        target_lock: &Arc<AsyncMutex<()>>,
        sink: &mut dyn TurnSink,
        snapshot: &TurnSnapshot,
    ) -> Result<(), DeliveryError> {
        let attempts = self.final_retry_delays.len() + 1;
        for attempt in 0..attempts {
            let _guard = target_lock.lock().await;
            let result = if snapshot.status == TurnStatus::Completed {
                timeout(self.sink_call_timeout, sink.finish(snapshot)).await
            } else {
                timeout(self.sink_call_timeout, sink.abort(snapshot)).await
            };
            drop(_guard);

            match result {
                Ok(Ok(())) => return Ok(()),
                Ok(Err(error))
                    if error.is_transient() && attempt < self.final_retry_delays.len() =>
                {
                    sleep_until(Instant::now() + self.final_retry_delays[attempt]).await;
                }
                Ok(Err(error)) => return Err(DeliveryError::FinalFailed(error)),
                Err(_) if attempt < self.final_retry_delays.len() => {
                    sleep_until(Instant::now() + self.final_retry_delays[attempt]).await;
                }
                Err(_) => return Err(DeliveryError::FinalTimedOut),
            }
        }
        unreachable!()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use async_trait::async_trait;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use tokio::sync::{Mutex as TokioMutex, Notify};

    use crate::agent::TurnEvent;
    use crate::bus::{MessageBus, OutboundMessage};
    use crate::channels::{Channel, TurnSink};
    use crate::session::{TurnBlock, TurnController};

    #[derive(Default)]
    struct SinkState {
        updates: TokioMutex<Vec<TurnSnapshot>>,
        terminal: TokioMutex<Vec<TurnSnapshot>>,
        update_started: Notify,
        release_update: Notify,
        block_first_update: bool,
        fail_updates: AtomicUsize,
        fail_finish: AtomicUsize,
    }

    struct RecordingSink(Arc<SinkState>);

    struct SinkChannel {
        state: Arc<SinkState>,
        opened: AtomicUsize,
    }

    #[async_trait]
    impl Channel for SinkChannel {
        fn name(&self) -> &str {
            "sink-channel"
        }

        fn is_running(&self) -> bool {
            true
        }

        async fn start(&self, _bus: Arc<MessageBus>) -> Result<(), ChannelError> {
            Ok(())
        }

        async fn stop(&self) -> Result<(), ChannelError> {
            Ok(())
        }

        fn live_policy(&self) -> LivePolicy {
            LivePolicy::FinalOnly
        }

        async fn open_turn(&self, _target: TurnTarget) -> Result<Box<dyn TurnSink>, ChannelError> {
            self.opened.fetch_add(1, Ordering::SeqCst);
            Ok(sink(self.state.clone()))
        }

        async fn send(&self, _msg: OutboundMessage) -> Result<(), ChannelError> {
            Ok(())
        }
    }

    #[async_trait]
    impl TurnSink for RecordingSink {
        async fn update(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError> {
            self.0.update_started.notify_waiters();
            if self.0.block_first_update && self.0.updates.lock().await.is_empty() {
                self.0.release_update.notified().await;
            }
            if self
                .0
                .fail_updates
                .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |remaining| {
                    remaining.checked_sub(1)
                })
                .is_ok()
            {
                return Err(ChannelError::SendError("update".into()));
            }
            self.0.updates.lock().await.push(snapshot.clone());
            Ok(())
        }

        async fn finish(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError> {
            if self
                .0
                .fail_finish
                .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |remaining| {
                    remaining.checked_sub(1)
                })
                .is_ok()
            {
                return Err(ChannelError::SendError("finish".into()));
            }
            self.0.terminal.lock().await.push(snapshot.clone());
            Ok(())
        }

        async fn abort(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError> {
            self.0.terminal.lock().await.push(snapshot.clone());
            Ok(())
        }
    }

    fn sink(state: Arc<SinkState>) -> Box<dyn TurnSink> {
        Box::new(RecordingSink(state))
    }

    #[tokio::test]
    async fn slow_sink_observes_latest_snapshot_and_terminal_bypasses_throttle() {
        let state = Arc::new(SinkState {
            block_first_update: true,
            ..SinkState::default()
        });
        let (controller, emitter, receiver) = TurnController::start("session", "message");
        let coordinator = DeliveryCoordinator::for_test(Duration::from_secs(30), []);
        let task = tokio::spawn({
            let state = state.clone();
            async move {
                coordinator
                    .deliver(
                        "cli_chat",
                        "chat",
                        LivePolicy::Snapshot {
                            min_interval: Duration::from_secs(10),
                        },
                        PresentationPolicy::interactive(),
                        receiver,
                        sink(state),
                    )
                    .await
            }
        });

        state.update_started.notified().await;
        emitter
            .emit(TurnEvent::TextDelta {
                iteration: 0,
                delta: "a".into(),
            })
            .unwrap();
        emitter
            .emit(TurnEvent::TextDelta {
                iteration: 0,
                delta: "b".into(),
            })
            .unwrap();
        state.release_update.notify_waiters();
        tokio::task::yield_now().await;
        controller.complete(None);

        assert!(task.await.unwrap().is_ok());
        let terminal = state.terminal.lock().await;
        assert_eq!(terminal.len(), 1);
        assert_eq!(terminal[0].status, TurnStatus::Completed);
        assert!(
            matches!(&terminal[0].blocks[0], TurnBlock::Assistant { text, .. } if text == "ab")
        );
    }

    #[tokio::test]
    async fn hidden_reasoning_is_removed_before_sink_and_failed_update_recovers() {
        let state = Arc::new(SinkState {
            fail_updates: AtomicUsize::new(1),
            ..SinkState::default()
        });
        let (controller, emitter, receiver) = TurnController::start("session", "message");
        let coordinator = DeliveryCoordinator::for_test(Duration::from_secs(30), []);
        let task = tokio::spawn({
            let state = state.clone();
            async move {
                coordinator
                    .deliver(
                        "feishu",
                        "chat",
                        LivePolicy::Snapshot {
                            min_interval: Duration::ZERO,
                        },
                        PresentationPolicy::external(true),
                        receiver,
                        sink(state),
                    )
                    .await
            }
        });

        emitter
            .emit(TurnEvent::ReasoningDelta {
                iteration: 0,
                delta: "secret".into(),
            })
            .unwrap();
        tokio::task::yield_now().await;
        emitter
            .emit(TurnEvent::TextDelta {
                iteration: 0,
                delta: "public".into(),
            })
            .unwrap();
        tokio::task::yield_now().await;
        controller.complete(None);

        assert!(task.await.unwrap().is_ok());
        let terminal = state.terminal.lock().await;
        assert!(
            terminal[0]
                .blocks
                .iter()
                .all(|block| !matches!(block, TurnBlock::Reasoning { .. }))
        );
    }

    #[tokio::test]
    async fn final_only_skips_updates_and_retries_transient_finish() {
        let state = Arc::new(SinkState {
            fail_finish: AtomicUsize::new(2),
            ..SinkState::default()
        });
        let (controller, emitter, receiver) = TurnController::start("session", "message");
        let coordinator = DeliveryCoordinator::for_test(
            Duration::from_secs(30),
            [Duration::from_millis(1), Duration::from_millis(2)],
        );
        let task = tokio::spawn({
            let state = state.clone();
            async move {
                coordinator
                    .deliver(
                        "channel",
                        "chat",
                        LivePolicy::FinalOnly,
                        PresentationPolicy::unattended(),
                        receiver,
                        sink(state),
                    )
                    .await
            }
        });

        emitter
            .emit(TurnEvent::TextDelta {
                iteration: 0,
                delta: "done".into(),
            })
            .unwrap();
        controller.complete(None);

        assert!(task.await.unwrap().is_ok());
        assert!(state.updates.lock().await.is_empty());
        assert_eq!(state.terminal.lock().await.len(), 1);
    }

    #[tokio::test]
    async fn open_and_deliver_owns_sink_creation_and_terminal_lifecycle() {
        let state = Arc::new(SinkState::default());
        let channel = Arc::new(SinkChannel {
            state: state.clone(),
            opened: AtomicUsize::new(0),
        });
        let (controller, emitter, receiver) = TurnController::start("session", "message");
        let coordinator = DeliveryCoordinator::for_test(Duration::from_secs(30), []);
        let target = TurnTarget {
            channel: "sink-channel".into(),
            chat_id: "chat".into(),
            session_id: "session".into(),
            reply_to: None,
            metadata: HashMap::new(),
        };
        let task = tokio::spawn({
            let channel = channel.clone();
            async move {
                coordinator
                    .open_and_deliver(channel, target, PresentationPolicy::unattended(), receiver)
                    .await
            }
        });

        emitter
            .emit(TurnEvent::TextDelta {
                iteration: 0,
                delta: "done".into(),
            })
            .unwrap();
        controller.complete(None);

        assert!(task.await.unwrap().is_ok());
        assert_eq!(channel.opened.load(Ordering::SeqCst), 1);
        assert_eq!(state.terminal.lock().await.len(), 1);
    }

    #[tokio::test]
    async fn supervisor_shutdown_aborts_sink_and_waits_for_cleanup() {
        let state = Arc::new(SinkState::default());
        let (_controller, emitter, receiver) = TurnController::start("session", "message");
        emitter
            .emit(TurnEvent::TextDelta {
                iteration: 0,
                delta: "partial".into(),
            })
            .unwrap();
        let supervisor = TaskSupervisor::new();
        let coordinator = DeliveryCoordinator::for_test(Duration::from_secs(1), []);
        let result = coordinator
            .spawn_sink(
                &supervisor,
                SinkRoute {
                    channel: "channel".into(),
                    chat_id: "chat".into(),
                    live_policy: LivePolicy::FinalOnly,
                    presentation: PresentationPolicy::unattended(),
                },
                receiver,
                sink(state.clone()),
            )
            .unwrap();
        tokio::task::yield_now().await;

        supervisor.shutdown(Duration::from_secs(1)).await;

        assert!(result.await.unwrap().is_ok());
        let terminal = state.terminal.lock().await;
        assert_eq!(terminal.len(), 1);
        assert_eq!(terminal[0].status, TurnStatus::Running);
    }
}