"""Shared response lifecycle helpers."""

from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, TypeVar, cast

from agno.db.base import SessionType

from mindroom.agent_storage import get_agent_session, get_team_session
from mindroom.ai_runtime import finalize_queued_notice_response_turn_async, queued_message_signal_context
from mindroom.hooks import EVENT_SESSION_STARTED, SessionHookContext, emit
from mindroom.post_response_effects import apply_post_response_effects
from mindroom.tool_system.runtime_context import resolve_tool_runtime_hook_bindings

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable

    from agno.db.base import BaseDb
    from structlog.stdlib import BoundLogger

    from mindroom.delivery_gateway import ResponseHookService, ResponseIdentity
    from mindroom.final_delivery import FinalDeliveryOutcome
    from mindroom.history.types import HistoryScope
    from mindroom.hooks import MessageEnvelope
    from mindroom.message_target import MessageTarget
    from mindroom.post_response_effects import PostResponseEffectsDeps, ResponseOutcome
    from mindroom.timing import DispatchPipelineTiming
    from mindroom.tool_system.runtime_context import ToolRuntimeContext

_LockedResponseResult = TypeVar("_LockedResponseResult")


@dataclass
class _QueuedMessageState:
    """Track queued human ingress while one response lifecycle holds the lock."""

    pending_human_message_event_ids: set[str] = field(default_factory=set)
    _active_response_turns: int = 0
    _event: asyncio.Event = field(default_factory=asyncio.Event)
    _idle_event: asyncio.Event = field(default_factory=asyncio.Event)

    def __post_init__(self) -> None:
        self._idle_event.set()

    @property
    def pending_human_messages(self) -> int:
        """Return the number of distinct queued human events."""
        return len(self.pending_human_message_event_ids)

    def begin_response_turn(self) -> bool:
        existing_turn = self._active_response_turns > 0
        self._active_response_turns += 1
        self._idle_event.clear()
        return existing_turn

    def finish_response_turn(self) -> None:
        if self._active_response_turns == 0:
            return
        self._active_response_turns -= 1
        if self._active_response_turns == 0:
            self._idle_event.set()

    def add_waiting_human_message(self, source_event_id: str) -> bool:
        previous_count = self.pending_human_messages
        self.pending_human_message_event_ids.add(source_event_id)
        self._event.set()
        return self.pending_human_messages != previous_count

    def consume_waiting_human_message(self, source_event_id: str) -> None:
        if source_event_id not in self.pending_human_message_event_ids:
            return
        self.pending_human_message_event_ids.remove(source_event_id)
        if self.pending_human_messages == 0:
            self._event.clear()

    def has_pending_human_messages(self) -> bool:
        return self.pending_human_messages > 0

    def has_active_response_turn(self) -> bool:
        return self._active_response_turns > 0

    async def wait(self) -> None:
        await self._event.wait()

    async def wait_until_idle(self) -> None:
        await self._idle_event.wait()

    def is_set(self) -> bool:
        return self._event.is_set()


@dataclass(slots=True)
class QueuedHumanNoticeReservation:
    """Owned reservation for a queued-human notice created before dispatch starts."""

    _state: _QueuedMessageState
    _source_event_id: str
    _active: bool = True

    def _release_waiting_human_message(self) -> None:
        if not self._active:
            return
        self._state.consume_waiting_human_message(self._source_event_id)
        self._active = False

    def consume(self) -> None:
        """Mark the reservation as owned by the response lifecycle."""
        self._release_waiting_human_message()

    def cancel(self) -> None:
        """Release a reservation that will not reach response lifecycle ownership."""
        self._release_waiting_human_message()


@dataclass
class ResponseLifecycleCoordinator:
    """Serialize response turns and signal active turns about queued human ingress."""

    _response_lifecycle_locks: dict[tuple[str, str | None], asyncio.Lock] = field(default_factory=dict)
    _thread_queued_signals: dict[tuple[str, str | None], _QueuedMessageState] = field(default_factory=dict)

    @staticmethod
    def _thread_key(target: MessageTarget) -> tuple[str, str | None]:
        return (target.room_id, target.resolved_thread_id)

    def _has_active_response_for_thread_key(self, thread_key: tuple[str, str | None]) -> bool:
        queued_signal = self._thread_queued_signals.get(thread_key)
        if queued_signal is not None and queued_signal.has_active_response_turn():
            return True
        lifecycle_lock = self._response_lifecycle_locks.get(thread_key)
        return lifecycle_lock.locked() if lifecycle_lock is not None else False

    def has_active_response_for_target(self, target: MessageTarget) -> bool:
        """Return whether one canonical conversation target already has an active turn."""
        return self._has_active_response_for_thread_key(self._thread_key(target))

    def active_thread_ids_for_room(self, room_id: str) -> frozenset[str | None]:
        """Return canonical thread IDs with active response lifecycles in one room."""
        known_thread_keys = set(self._thread_queued_signals) | set(self._response_lifecycle_locks)
        return frozenset(
            thread_id
            for known_room_id, thread_id in known_thread_keys
            if known_room_id == room_id and self._has_active_response_for_thread_key((known_room_id, thread_id))
        )

    async def wait_for_thread_idle(self, room_id: str, thread_id: str | None) -> None:
        """Wait until a response lifecycle lock is idle for one room/thread key."""
        thread_key = (room_id, thread_id)
        while self._has_active_response_for_thread_key(thread_key):
            queued_signal = self._thread_queued_signals.get(thread_key)
            if queued_signal is not None and queued_signal.has_active_response_turn():
                await queued_signal.wait_until_idle()
                continue
            lifecycle_lock = self._response_lifecycle_locks.get(thread_key)
            if lifecycle_lock is not None and lifecycle_lock.locked():
                async with lifecycle_lock:
                    pass
                continue
            return

    def _response_lifecycle_lock(self, target: MessageTarget) -> asyncio.Lock:
        """Return the per-target lock that serializes one response lifecycle."""
        lock_key = self._thread_key(target)
        lock = self._response_lifecycle_locks.get(lock_key)
        if lock is not None:
            return lock
        if len(self._response_lifecycle_locks) >= 100:
            for candidate, candidate_lock in list(self._response_lifecycle_locks.items()):
                if len(self._response_lifecycle_locks) < 100:
                    break
                if candidate_lock.locked():
                    continue
                self._response_lifecycle_locks.pop(candidate, None)
                self._thread_queued_signals.pop(candidate, None)
        lock = asyncio.Lock()
        self._response_lifecycle_locks[lock_key] = lock
        return lock

    def _get_or_create_queued_signal(self, target: MessageTarget) -> _QueuedMessageState:
        """Return the queued-message signal for one canonical conversation thread."""
        thread_key = self._thread_key(target)
        signal = self._thread_queued_signals.get(thread_key)
        if signal is not None:
            return signal
        signal = _QueuedMessageState()
        self._thread_queued_signals[thread_key] = signal
        return signal

    @staticmethod
    def _should_signal_queued_message(
        response_envelope: MessageEnvelope,
    ) -> bool:
        """Return whether one queued ingress should interrupt the active turn."""
        return response_envelope.origin.may_answer_interactive_prompt

    def _assert_target_matches_envelope(self, target: MessageTarget, response_envelope: MessageEnvelope) -> None:
        """Require lifecycle callers to use the envelope's canonical response target."""
        if self._thread_key(target) == self._thread_key(response_envelope.target):
            return
        msg = "Response lifecycle target must match MessageEnvelope.target"
        raise ValueError(msg)

    def reserve_waiting_human_message(
        self,
        *,
        target: MessageTarget,
        response_envelope: MessageEnvelope,
    ) -> QueuedHumanNoticeReservation | None:
        """Reserve an active-turn notice before queued dispatch owns the follow-up."""
        self._assert_target_matches_envelope(target, response_envelope)
        if not self._should_signal_queued_message(response_envelope):
            return None
        thread_key = self._thread_key(target)
        if not self._has_active_response_for_thread_key(thread_key):
            return None
        queued_signal = self._get_or_create_queued_signal(target)
        if not queued_signal.add_waiting_human_message(response_envelope.source_event_id):
            return None
        return QueuedHumanNoticeReservation(queued_signal, response_envelope.source_event_id)

    def _begin_response_turn_notice(
        self,
        *,
        lifecycle_lock: asyncio.Lock,
        queued_signal: _QueuedMessageState,
        response_envelope: MessageEnvelope,
        queued_notice_reservation: QueuedHumanNoticeReservation | None,
        signal_queued_message: bool,
    ) -> str | None:
        existing_turn = queued_signal.begin_response_turn()
        if not signal_queued_message or queued_notice_reservation is not None:
            return None
        if not (existing_turn or lifecycle_lock.locked()):
            return None
        if not self._should_signal_queued_message(response_envelope):
            return None
        if not queued_signal.add_waiting_human_message(response_envelope.source_event_id):
            return None
        return response_envelope.source_event_id

    def _consume_queued_human_notice(
        self,
        *,
        notice: str | None,
        queued_signal: _QueuedMessageState,
    ) -> None:
        if notice is None:
            return
        queued_signal.consume_waiting_human_message(notice)

    async def run_locked_response(
        self,
        *,
        target: MessageTarget,
        response_envelope: MessageEnvelope,
        queued_notice_reservation: QueuedHumanNoticeReservation | None,
        pipeline_timing: DispatchPipelineTiming | None,
        locked_operation: Callable[[MessageTarget], Awaitable[_LockedResponseResult]],
        signal_queued_message: bool = True,
    ) -> _LockedResponseResult:
        """Run one locked response operation with shared queued-message bookkeeping."""
        self._assert_target_matches_envelope(target, response_envelope)
        lifecycle_lock = self._response_lifecycle_lock(target)
        queued_signal = self._get_or_create_queued_signal(target)
        notice = self._begin_response_turn_notice(
            lifecycle_lock=lifecycle_lock,
            queued_signal=queued_signal,
            response_envelope=response_envelope,
            queued_notice_reservation=queued_notice_reservation,
            signal_queued_message=signal_queued_message,
        )
        lock_acquired = False
        reservation_consumed = False
        try:
            if pipeline_timing is not None:
                pipeline_timing.mark("lock_wait_start")
            await lifecycle_lock.acquire()
            lock_acquired = True
            if pipeline_timing is not None:
                pipeline_timing.mark("lock_acquired")
            try:
                if queued_notice_reservation is not None:
                    queued_notice_reservation.consume()
                    reservation_consumed = True
                notice = self._consume_queued_human_notice(
                    notice=notice,
                    queued_signal=queued_signal,
                )
                with queued_message_signal_context(queued_signal) as notice_context:
                    try:
                        return await locked_operation(target)
                    finally:
                        await finalize_queued_notice_response_turn_async(notice_context)
            finally:
                if lock_acquired:
                    lifecycle_lock.release()
        finally:
            if queued_notice_reservation is not None and not reservation_consumed:
                queued_notice_reservation.cancel()
            self._consume_queued_human_notice(
                notice=notice,
                queued_signal=queued_signal,
            )
            queued_signal.finish_response_turn()

    async def run_locked_target_operation(
        self,
        *,
        target: MessageTarget,
        while_waiting: Callable[[], Awaitable[None]],
        locked_operation: Callable[[], Awaitable[_LockedResponseResult]],
    ) -> _LockedResponseResult:
        """Run a non-response operation under one target's response lock."""
        lifecycle_lock = self._response_lifecycle_lock(target)
        acquire_task = asyncio.create_task(lifecycle_lock.acquire())
        lock_acquired = False
        try:
            while not acquire_task.done():
                await while_waiting()
                await asyncio.wait({acquire_task}, timeout=0.01)
            await acquire_task
            lock_acquired = True
            return await locked_operation()
        finally:
            if not acquire_task.done():
                acquire_task.cancel()
                await asyncio.gather(acquire_task, return_exceptions=True)
            elif not lock_acquired and not acquire_task.cancelled() and acquire_task.exception() is None:
                lifecycle_lock.release()
            if lock_acquired:
                lifecycle_lock.release()


@dataclass(frozen=True)
class _SessionStartedWatch:
    """Pre-computed session:started eligibility and emission arguments."""

    should_watch: bool
    tool_context: ToolRuntimeContext | None
    scope: HistoryScope
    session_id: str
    room_id: str
    thread_id: str | None
    session_type: SessionType
    correlation_id: str
    create_storage: Callable[[], BaseDb]


@dataclass(frozen=True)
class ResponseLifecycleDeps:
    """Dependencies owned by the response lifecycle boundary."""

    response_hooks: ResponseHookService
    logger: BoundLogger


def _session_exists(
    *,
    storage: BaseDb,
    session_id: str,
    session_type: SessionType,
) -> bool:
    if session_type is SessionType.TEAM:
        return get_team_session(storage, session_id) is not None
    return get_agent_session(storage, session_id) is not None


def _response_outcome_label(final_delivery_outcome: FinalDeliveryOutcome | None) -> str:
    """Return one pipeline outcome label for the canonical final delivery outcome."""
    if final_delivery_outcome is not None and final_delivery_outcome.suppressed:
        return "suppressed"
    if final_delivery_outcome is not None and final_delivery_outcome.terminal_status == "cancelled":
        return "cancelled"
    if final_delivery_outcome is not None and final_delivery_outcome.terminal_status == "error":
        return "error"
    if final_delivery_outcome is not None and final_delivery_outcome.delivery_kind is not None:
        return final_delivery_outcome.delivery_kind
    if (
        final_delivery_outcome is not None
        and final_delivery_outcome.event_id is not None
        and final_delivery_outcome.is_visible_response
    ):
        return "visible_response_preserved"
    return "no_visible_response"


class ResponseLifecycle:
    """Consolidate lifecycle helpers shared across response paths."""

    def __init__(
        self,
        deps: ResponseLifecycleDeps,
        *,
        identity: ResponseIdentity,
        pipeline_timing: DispatchPipelineTiming | None,
    ) -> None:
        self.deps = deps
        self.identity = identity
        self.pipeline_timing = pipeline_timing

    def _log_effects_failure_after_visible_delivery(
        self,
        *,
        response_event_id: str,
        error: BaseException,
    ) -> None:
        """Log one non-fatal post-response failure after visible delivery succeeded."""
        self.deps.logger.error(
            "Post-response effects failed after visible delivery",
            response_kind=self.identity.response_kind,
            response_event_id=response_event_id,
            failure_reason=str(error),
            error_type=error.__class__.__name__,
        )

    def _session_started_watch_is_needed(
        self,
        *,
        tool_context: ToolRuntimeContext | None,
        session_id: str,
        session_type: SessionType,
        create_storage: Callable[[], BaseDb],
    ) -> bool:
        if tool_context is None or not tool_context.hook_registry.has_hooks(EVENT_SESSION_STARTED):
            return False
        try:
            storage = create_storage()
            try:
                return not _session_exists(
                    storage=storage,
                    session_id=session_id,
                    session_type=session_type,
                )
            finally:
                storage.close()
        except Exception as error:
            self.deps.logger.exception(
                "Failed to probe session storage for session:started eligibility",
                session_id=session_id,
                session_type=str(session_type),
                failure_reason=str(error),
            )
            return False

    def setup_session_watch(
        self,
        *,
        tool_context: ToolRuntimeContext | None,
        session_id: str,
        session_type: SessionType,
        scope: HistoryScope,
        room_id: str,
        thread_id: str | None,
        create_storage: Callable[[], BaseDb],
    ) -> _SessionStartedWatch:
        """Pre-compute session:started eligibility for one response path."""
        return _SessionStartedWatch(
            should_watch=self._session_started_watch_is_needed(
                tool_context=tool_context,
                session_id=session_id,
                session_type=session_type,
                create_storage=create_storage,
            ),
            tool_context=tool_context,
            scope=scope,
            session_id=session_id,
            room_id=room_id,
            thread_id=thread_id,
            session_type=session_type,
            correlation_id=self.identity.correlation_id,
            create_storage=create_storage,
        )

    async def _maybe_emit_session_started(self, watch: _SessionStartedWatch) -> None:
        if watch.tool_context is None or not watch.should_watch:
            return
        storage = watch.create_storage()
        try:
            if not _session_exists(storage=storage, session_id=watch.session_id, session_type=watch.session_type):
                return
        finally:
            storage.close()

        bindings = resolve_tool_runtime_hook_bindings(watch.tool_context)
        context = SessionHookContext(
            event_name=EVENT_SESSION_STARTED,
            plugin_name="",
            settings={},
            config=watch.tool_context.config,
            runtime_paths=watch.tool_context.runtime_paths,
            logger=self.deps.logger.bind(event_name=EVENT_SESSION_STARTED, session_id=watch.session_id),
            correlation_id=watch.correlation_id,
            message_sender=bindings.message_sender,
            matrix_admin=bindings.matrix_admin,
            room_state_querier=bindings.room_state_querier,
            room_state_putter=bindings.room_state_putter,
            agent_name=watch.scope.scope_id if watch.scope.kind == "team" else watch.tool_context.agent_name,
            scope=watch.scope,
            session_id=watch.session_id,
            room_id=watch.room_id,
            thread_id=watch.thread_id,
        )
        await emit(watch.tool_context.hook_registry, EVENT_SESSION_STARTED, context)

    async def emit_session_started(self, watch: _SessionStartedWatch) -> None:
        """Emit session:started without aborting delivery on ordinary failures."""
        try:
            await self._maybe_emit_session_started(watch)
        except asyncio.CancelledError:
            raise
        except Exception as error:
            self.deps.logger.exception(
                "Failed to emit session:started",
                session_id=watch.session_id,
                room_id=watch.room_id,
                thread_id=watch.thread_id,
                failure_reason=str(error),
            )

    async def finalize(
        self,
        final_delivery_outcome: FinalDeliveryOutcome,
        *,
        build_post_response_outcome: Callable[[FinalDeliveryOutcome], ResponseOutcome],
        post_response_deps: PostResponseEffectsDeps | Callable[[], PostResponseEffectsDeps],
    ) -> FinalDeliveryOutcome:
        """Run outer lifecycle finalization and return the canonical terminal outcome."""
        response_event_id = final_delivery_outcome.final_visible_event_id
        try:
            if final_delivery_outcome.terminal_status == "completed":
                if (
                    response_event_id is not None
                    and final_delivery_outcome.final_visible_body is not None
                    and final_delivery_outcome.delivery_kind is not None
                ):
                    await self.deps.response_hooks.emit_after_response(
                        identity=self.identity,
                        response_text=final_delivery_outcome.final_visible_body,
                        response_event_id=response_event_id,
                        delivery_kind=final_delivery_outcome.delivery_kind,
                        continue_on_cancelled=True,
                    )
            else:
                await self.deps.response_hooks.emit_cancelled_response(
                    identity=self.identity,
                    visible_response_event_id=response_event_id,
                    failure_reason=final_delivery_outcome.failure_reason,
                )
        except asyncio.CancelledError as error:
            if response_event_id is None:
                raise
            self._log_effects_failure_after_visible_delivery(
                response_event_id=response_event_id,
                error=error,
            )
        except Exception as error:
            if response_event_id is None:
                raise
            self._log_effects_failure_after_visible_delivery(
                response_event_id=response_event_id,
                error=error,
            )
        await self._apply_effects_safely(
            final_delivery_outcome=final_delivery_outcome,
            post_response_outcome=lambda: build_post_response_outcome(final_delivery_outcome),
            post_response_deps=post_response_deps,
        )
        if self.pipeline_timing is not None:
            self.pipeline_timing.emit_summary(self.deps.logger, outcome=_response_outcome_label(final_delivery_outcome))
        return final_delivery_outcome

    async def _apply_effects_safely(
        self,
        *,
        final_delivery_outcome: FinalDeliveryOutcome,
        post_response_outcome: ResponseOutcome | Callable[[], ResponseOutcome],
        post_response_deps: PostResponseEffectsDeps | Callable[[], PostResponseEffectsDeps],
    ) -> None:
        """Apply post-response effects without masking failures before visible delivery."""
        response_event_id = final_delivery_outcome.final_visible_event_id
        try:
            if callable(post_response_outcome):
                post_response_outcome = cast("Callable[[], ResponseOutcome]", post_response_outcome)()
            if callable(post_response_deps):
                post_response_deps = cast("Callable[[], PostResponseEffectsDeps]", post_response_deps)()
            await apply_post_response_effects(
                final_delivery_outcome,
                post_response_outcome,
                post_response_deps,
            )
        except asyncio.CancelledError as error:
            if response_event_id is None:
                raise
            self._log_effects_failure_after_visible_delivery(
                response_event_id=response_event_id,
                error=error,
            )
        except Exception as error:
            if response_event_id is None:
                raise
            self._log_effects_failure_after_visible_delivery(
                response_event_id=response_event_id,
                error=error,
            )
