"""Running turns: one worker per conversation over its queue, the backend call."""

from __future__ import annotations

import asyncio
import base64
import contextlib
import inspect
import logging
from datetime import UTC, datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast
from uuid import uuid4

from claude_agent_sdk import (
    AssistantMessage,
    RateLimitEvent,
    ResultMessage,
    StreamEvent,
    ToolResultBlock,
    ToolUseBlock,
    UserMessage,
)

from beaver_gateway.backends.capture import TurnCapture
from beaver_gateway.backends.transcript import reply_text, text_of
from beaver_gateway.conversations import injects
from beaver_gateway.conversations.seeds import Seeds
from beaver_gateway.conversations.state import Runner, aware, iso
from beaver_gateway.conversations.texts import UserSaid
from beaver_gateway.frontends.accumulate import StreamAccumulator
from beaver_gateway.storage.models import Conversation, InjectQueueItem, RateLimit

if TYPE_CHECKING:
    from collections.abc import AsyncIterator, Callable, Sequence

    from beaver_gateway.events.stream import MessageStreamEvent

__all__ = ["Turns", "content_of", "image_block"]

_log = logging.getLogger(__name__)


class Turns(Seeds):
    async def turn(
        self,
        conv: Conversation,
        *,
        messages: Sequence[Any],
        origin: str,
        capture: TurnCapture | None = None,
        session_id: str | None = None,
        use_session: bool = True,
        tools: bool = True,
        turn_id: str | None = None,
        item_origin: str | None = None,
        shown: str | None = None,
        attach: bool = False,
    ) -> AsyncIterator[MessageStreamEvent]:
        """Run one turn under the conversation's lock; the only path to the backend.

        ``shown`` is what frontends display as the turn's prompt (the user's
        words, not the envelope). ``attach`` sends nothing: the turn is one
        the live session started by itself, a subagent reporting back.
        """
        row_id = cast("int", conv.id)
        runner = self._runner(row_id)
        backend = self._backend(conv.agent_name)
        turn_id = turn_id or f"turn_{uuid4().hex[:12]}"
        capture = capture or TurnCapture()
        resume = session_id if session_id is not None else conv.session_id
        async with runner.lock:
            runner.turn_id = turn_id
            runner.origin = origin
            runner.text = shown if shown is not None else _prompt_preview(messages)
            runner.started_at = datetime.now(UTC)
            runner.tools = {}
            await self._mark_running(conv, turn_id)
            before = await self.before_turn(conv)
            self._bus.publish(
                "turn.start",
                conversation_id=conv.external_id,
                turn_id=turn_id,
                origin=origin,
                item_origin=item_origin,
                text=runner.text,
            )
            stop = "error"
            cut = False
            try:
                observer = self._observer(conv, runner, turn_id, origin)
                events = (
                    backend.attach(
                        key=conv.external_id,
                        capture=capture,
                        observer=observer,
                        turn_id=turn_id,
                    )
                    if attach
                    else backend.complete(
                        agent=self._claude_agent(conv.agent_name),
                        messages=messages,
                        conversation_id=conv.external_id,
                        session_id=resume if use_session else None,
                        reseed=not use_session,
                        capture=capture,
                        kind=conv.kind,
                        pinned=conv.kind == "master",
                        tools=tools,
                        observer=observer,
                        turn_id=turn_id,
                    )
                )
                async for event in events:
                    yield event
                stop = "interrupted" if capture.interrupted else "end_turn"
            except asyncio.CancelledError:
                cut = True
                raise
            finally:
                runner.turn_id = None
                await self._mark_done(conv, capture, cut=cut)
                if stop != "error":
                    try:
                        await self.after_turn(conv, before)
                    except Exception:  # noqa: BLE001
                        _log.exception("after-turn hook on %s failed", conv.external_id)
                self._bus.publish(
                    "turn.end",
                    conversation_id=conv.external_id,
                    turn_id=turn_id,
                    origin=origin,
                    item_origin=item_origin,
                    stop=stop,
                    usage=_usage_dict(capture),
                )

    async def run_text_turn(
        self,
        conv: Conversation,
        text: str,
        *,
        origin: str,
        tools: bool = True,
        turn_id: str | None = None,
        item_origin: str | None = None,
        attachments: Sequence[dict[str, Any]] | None = None,
        shown: str | None = None,
        attach: bool = False,
    ) -> tuple[str, TurnCapture]:
        capture = TurnCapture()
        acc = StreamAccumulator()
        agent = self._claude_agent(conv.agent_name)
        async for event in self.turn(
            conv,
            messages=[{"role": "user", "content": content_of(text, attachments)}],
            origin=origin,
            capture=capture,
            tools=tools,
            turn_id=turn_id,
            item_origin=item_origin,
            shown=shown,
            attach=attach,
        ):
            acc.feed(event)
        message = acc.finalize(model=agent.model)
        return reply_text(
            message.content, thinking_is_progress=agent.thinking_is_progress
        ), capture

    async def before_turn(self, conv: Conversation) -> str | None:  # noqa: ARG002
        return None

    async def after_turn(self, conv: Conversation, before: str | None) -> None:  # noqa: ARG002
        return

    def turn_origin(self, conv: Conversation) -> str | None:
        runner = self._runners.get(cast("int", conv.id))
        return runner.origin if runner is not None and runner.turn_id else None

    async def _woken(self, key: str, note: str) -> None:
        """The live session of ``key`` started a turn by itself; run it next."""
        conv = await self.get(key)
        if conv is None:
            _log.warning("self-started turn on unknown conversation %s: %s", key, note)
            return
        _log.info("conversation %s woke by itself: %s", key, note)
        self._ensure_worker(cast("int", conv.id))

    async def _run_pending(self, conv: Conversation) -> None:
        turn_id = f"turn_{uuid4().hex[:12]}"
        note = self._backend(conv.agent_name).pending_note(conv.external_id) or ""
        try:
            text, _ = await self.run_text_turn(
                conv,
                note,
                origin="task",
                turn_id=turn_id,
                item_origin="task",
                shown=note,
                attach=True,
            )
        except Exception:  # noqa: BLE001
            _log.exception(
                "self-started turn %s on %s failed", turn_id, conv.external_id
            )
            return
        if not text.strip():
            return
        self._bus.publish(
            "reply",
            conversation_id=conv.external_id,
            turn_id=turn_id,
            item=None,
            item_origin="task",
            source="task",
            prompt=self._texts.task_prompt.format(note=note),
            user_text=None,
            text=text,
        )

    def _ensure_worker(self, row_id: int) -> None:
        runner = self._runner(row_id)
        runner.wake.set()
        if runner.task is None or runner.task.done():
            runner.task = asyncio.create_task(self._worker(row_id))
            self._track(runner.task)

    async def _worker(self, row_id: int) -> None:
        runner = self._runner(row_id)
        while True:
            conv = await self.get_row(row_id)
            if conv is not None and self._has_pending(conv):
                await self._run_pending(conv)
                continue
            items = await self._queue.pending(row_id)
            batch, wait = self._pick(items)
            if batch is None:
                runner.wake.clear()
                if wait is None:
                    return
                with contextlib.suppress(TimeoutError):
                    await asyncio.wait_for(runner.wake.wait(), timeout=wait)
                continue
            conv = await self.get_row(row_id)
            if conv is None:
                await self._queue.finish(batch, status="failed")
                return
            await self._run_batch(conv, batch)

    def _pick(
        self, items: list[InjectQueueItem]
    ) -> tuple[list[InjectQueueItem] | None, float | None]:
        if not items:
            return None, None
        head = items[0]
        if head.priority == "urgent":
            return [head], None
        tail = [i for i in items if i is not head and i.priority in ("wake", "normal")]
        if head.priority in ("user", "wake"):
            return [head, *tail], None
        age = (datetime.now(UTC) - aware(head.created_at)).total_seconds()
        if age >= self._normal_window:
            return [head, *tail], None
        return None, max(self._normal_window - age, 1.0)

    async def _run_batch(
        self, conv: Conversation, batch: list[InjectQueueItem]
    ) -> None:
        head = batch[0]
        turn_id = f"turn_{uuid4().hex[:12]}"
        await self._queue.start(batch, turn_id)
        if head.priority == "user":
            origin = "user"
            prompt = head.text
            await self._note_user(conv, head.text)
            envelope = self._envelope_for(conv, head.text)
            if envelope:
                prompt += "\n\n" + envelope
            if len(batch) > 1:
                prompt += "\n\n" + self._bundle(batch[1:])
        else:
            origin = "inject"
            prompt = "\n\n".join(
                f"{self._texts.inject_header(injects.context_of(i))}\n{i.text}"
                for i in batch
            )
        seed = await self.pending_seed(conv)
        if seed:
            prompt = f"{seed}\n\n{prompt}"
        try:
            text, capture = await self.run_text_turn(
                conv,
                prompt,
                origin=origin,
                turn_id=turn_id,
                item_origin=head.origin,
                attachments=head.attachments if origin == "user" else None,
                shown=head.text if origin == "user" else None,
            )
        except Exception:  # noqa: BLE001
            _log.exception("turn %s on %s failed", turn_id, conv.external_id)
            await self._queue.finish(batch, status="failed")
            return
        await self._queue.finish(
            batch, status="interrupted" if capture.interrupted else "done"
        )
        if origin == "user":
            self._bus.publish(
                "reply",
                conversation_id=conv.external_id,
                turn_id=turn_id,
                item=head.id,
                item_origin=head.origin,
                source="queue",
                prompt=prompt,
                user_text=head.text,
                text=text,
            )

    def _has_pending(self, conv: Conversation) -> bool:
        try:
            return self._backend(conv.agent_name).has_pending(conv.external_id)
        except LookupError:
            return False

    def _bundle(self, items: Sequence[InjectQueueItem]) -> str:
        lines = [self._texts.bundle_header.format(since=iso(items[0].created_at))]
        lines.extend(f"- [{i.origin}] {i.text}" for i in items)
        return "\n".join(lines)

    def _envelope_for(self, conv: Conversation, text: str = "") -> str | None:
        if self._envelope is None:
            return None
        if conv.kind == "master":
            return self._envelope.build(text=text, kind="master", agent=conv.agent_name)
        if conv.kind == "branch":
            return self._envelope.recall_only(
                text=text, kind="branch", agent=conv.agent_name
            )
        return None

    async def _note_user(self, conv: Conversation, text: str) -> None:
        if self._user_sink is None:
            return
        message = UserSaid(
            conversation_id=conv.external_id,
            kind=conv.kind,
            title=conv.title,
            text=text,
            at=datetime.now(UTC),
            agent=conv.agent_name,
        )
        try:
            result = self._user_sink(message)
            if inspect.isawaitable(result):
                await result
        except Exception:  # noqa: BLE001
            _log.exception("user sink failed for %s", conv.external_id)

    def _observer(
        self, conv: Conversation, runner: Runner, turn_id: str, origin: str
    ) -> Callable[[Any], None]:
        conversation_id = conv.external_id

        def observe(message: Any) -> None:
            parent = getattr(message, "parent_tool_use_id", None)
            if isinstance(message, RateLimitEvent):
                self._observe_rate_limit(conv, message)
            elif isinstance(message, StreamEvent):
                self._bus.publish(
                    "stream",
                    conversation_id=conversation_id,
                    turn_id=turn_id,
                    origin=origin,
                    parent_tool_use_id=parent,
                    event=message.event,
                )
            elif isinstance(message, AssistantMessage):
                for block in message.content:
                    if isinstance(block, ToolUseBlock):
                        event = self._bus.publish(
                            "tool",
                            conversation_id=conversation_id,
                            turn_id=turn_id,
                            origin=origin,
                            parent_tool_use_id=parent,
                            tool_use_id=block.id,
                            name=block.name,
                            input=block.input,
                        )
                        runner.tools[block.id] = {
                            "tool_use_id": block.id,
                            "name": block.name,
                            "input": block.input,
                            "parent_tool_use_id": parent,
                            "started_at": event["ts"],
                            "ended_at": None,
                            "is_error": None,
                            "content": None,
                        }
            elif isinstance(message, UserMessage):
                blocks = message.content if isinstance(message.content, list) else ()
                for block in blocks:
                    if isinstance(block, ToolResultBlock):
                        event = self._bus.publish(
                            "tool.result",
                            conversation_id=conversation_id,
                            turn_id=turn_id,
                            origin=origin,
                            parent_tool_use_id=parent,
                            tool_use_id=block.tool_use_id,
                            is_error=bool(block.is_error),
                            content=_result_preview(block.content),
                        )
                        node = runner.tools.get(block.tool_use_id)
                        if node is not None:
                            node["ended_at"] = event["ts"]
                            node["is_error"] = event["is_error"]
                            node["content"] = event["content"]
            elif isinstance(message, ResultMessage) and parent is None:
                self._bus.publish(
                    "result",
                    conversation_id=conversation_id,
                    turn_id=turn_id,
                    origin=origin,
                    subtype=message.subtype,
                    is_error=message.is_error,
                    num_turns=message.num_turns,
                )

        return observe

    def _observe_rate_limit(self, conv: Conversation, message: RateLimitEvent) -> None:
        info = message.rate_limit_info
        row = RateLimit(
            window=info.rate_limit_type or "unknown",
            status=info.status,
            utilization=info.utilization,
            resets_at=_from_unix(info.resets_at),
            overage_status=info.overage_status,
            overage_resets_at=_from_unix(info.overage_resets_at),
            agent_name=conv.agent_name,
            session_id=message.session_id,
            raw=dict(info.raw),
        )
        self._bus.publish(
            "rate_limit",
            conversation_id=conv.external_id,
            window=row.window,
            status=row.status,
            utilization=row.utilization,
            resets_at=iso(row.resets_at),
            overage_status=row.overage_status,
        )
        self._track(asyncio.create_task(self._record_rate_limit(row)))

    async def _record_rate_limit(self, row: RateLimit) -> None:
        try:
            async with self._db.session() as session:
                session.add(row)
                await session.commit()
        except Exception:  # noqa: BLE001
            _log.exception("rate limit write failed")

    async def _mark_running(self, conv: Conversation, turn_id: str) -> None:
        async def apply(row: Conversation) -> None:
            row.running_turn = turn_id
            row.last_activity_at = datetime.now(UTC)

        await self._update(conv, apply)

    async def _mark_done(
        self, conv: Conversation, capture: TurnCapture, *, cut: bool = False
    ) -> None:
        async def apply(row: Conversation) -> None:
            if not cut:
                row.running_turn = None
            row.last_activity_at = datetime.now(UTC)
            if capture.session_id is not None:
                row.session_id = capture.session_id

        await self._update(conv, apply)


IMAGE_MAX_BYTES = 5 * 1024 * 1024


def content_of(
    text: str, attachments: Sequence[dict[str, Any]] | None
) -> str | list[dict[str, Any]]:
    """Plain text, or text plus the images that ride with it."""
    blocks = [b for a in attachments or () if (b := image_block(a)) is not None]
    if not blocks:
        return text
    return [{"type": "text", "text": text}, *blocks]


def image_block(attachment: dict[str, Any]) -> dict[str, Any] | None:
    media_type = str(attachment.get("media_type") or "")
    path = Path(str(attachment.get("path") or ""))
    if not media_type.startswith("image/") or not path.is_file():
        return None
    if path.stat().st_size > IMAGE_MAX_BYTES:
        return None
    data = base64.b64encode(path.read_bytes()).decode("ascii")
    return {
        "type": "image",
        "source": {"type": "base64", "media_type": media_type, "data": data},
    }


def _prompt_preview(messages: Sequence[Any], limit: int = 400) -> str | None:
    if not messages:
        return None
    text = text_of(messages[-1].get("content"))
    return text[:limit] if text else None


def _from_unix(value: int | None) -> datetime | None:
    return datetime.fromtimestamp(value, tz=UTC) if value is not None else None


def _result_preview(
    content: str | list[dict[str, Any]] | None, limit: int = 400
) -> str:
    if content is None:
        return ""
    text = (
        content
        if isinstance(content, str)
        else "\n".join(
            str(part.get("text", ""))
            for part in content
            if isinstance(part, dict) and part.get("type") == "text"
        )
    )
    return text if len(text) <= limit else text[:limit] + "…"


def _usage_dict(capture: TurnCapture) -> dict[str, Any] | None:
    usage = capture.usage
    if usage is None:
        return None
    return {
        "input": usage.input_tokens,
        "output": usage.output_tokens,
        "cache_read": usage.cache_read_tokens,
        "cache_creation": usage.cache_creation_tokens,
        "cost_usd": usage.cost_usd,
        "duration_ms": usage.duration_ms,
    }