"""The conversation rows: create, find, bind to windows, flags, status, history."""

from __future__ import annotations

import json
import logging
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any, cast
from uuid import uuid4

from sqlalchemy import String, func
from sqlalchemy import cast as sa_cast
from sqlmodel import col, select

from beaver_gateway.backends.transcript import (
    messages_from_entries,
    render_messages,
    text_of,
)
from beaver_gateway.conversations.kinds import KINDS, Kind
from beaver_gateway.conversations.state import State, iso
from beaver_gateway.conversations.texts import split_service
from beaver_gateway.frontends.markdown.history import load_messages
from beaver_gateway.storage.models import (
    Conversation,
    ConversationBinding,
    ConversationMessage,
    RateLimit,
    TranscriptEntry,
    Usage,
)

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable, Iterable

    from beaver_gateway.frontends.base import Frontend

__all__ = [
    "MASTER_ALIAS",
    "PARENT_ALIAS",
    "STATUSES",
    "TITLE_MAX",
    "Rows",
    "context_of",
    "implied_title",
]

_log = logging.getLogger(__name__)

MASTER_ALIAS = "master"
PARENT_ALIAS = "parent"
STATUSES = ("open", "merged", "closed", "archived")
TITLE_MAX = 80


class Rows(State):
    async def create(
        self,
        *,
        kind: Kind,
        agent: str,
        parent: Conversation | None = None,
        title: str | None = None,
        origin: str = "api",
        session_id: str | None = None,
        flags: dict[str, Any] | None = None,
    ) -> Conversation:
        if kind not in KINDS:
            msg = f"unknown conversation kind {kind!r}"
            raise ValueError(msg)
        if not self._claude_agent(agent).serves(kind):
            msg = f"agent {agent!r} does not serve kind {kind!r}"
            raise ValueError(msg)
        now = datetime.now(UTC)
        row = Conversation(
            frontend=origin,
            external_id=str(uuid4()),
            agent_name=agent,
            kind=kind,
            parent_id=parent.id if parent is not None else None,
            title=title,
            session_id=session_id,
            flags=dict(flags or {}),
            last_activity_at=now,
        )
        async with self._db.session() as session:
            session.add(row)
            await session.commit()
            await session.refresh(row)
        self._bus.publish("conversation.created", **self.public(row))
        return row

    async def get(self, public_id: str) -> Conversation | None:
        async with self._db.session() as session:
            result = await session.exec(
                select(Conversation).where(Conversation.external_id == public_id)
            )
            return result.first()

    async def get_row(self, row_id: int) -> Conversation | None:
        async with self._db.session() as session:
            return await session.get(Conversation, row_id)

    async def resolve(
        self, key: str, *, origin: Conversation | None = None
    ) -> Conversation | None:
        """By public id, or ``master`` / ``parent`` relative to ``origin``."""
        key = key.strip()
        if key == MASTER_ALIAS:
            return await self.open_master()
        if key == PARENT_ALIAS:
            if origin is None or origin.parent_id is None:
                return None
            return await self.get_row(origin.parent_id)
        return await self.get(key)

    async def open_master(self) -> Conversation | None:
        masters = await self.find(kind="master", status="open", limit=1)
        return masters[0] if masters else None

    async def find(
        self,
        *,
        status: str | None = None,
        kind: str | None = None,
        parent: Conversation | None = None,
        limit: int = 200,
    ) -> list[Conversation]:
        stmt = select(Conversation).order_by(col(Conversation.id).desc()).limit(limit)
        if status is not None:
            stmt = stmt.where(Conversation.status == status)
        if kind is not None:
            stmt = stmt.where(Conversation.kind == kind)
        if parent is not None:
            stmt = stmt.where(Conversation.parent_id == parent.id)
        async with self._db.session() as session:
            return list((await session.exec(stmt)).all())

    async def search(self, query: str, *, limit: int = 20) -> list[Conversation]:
        needle = f"%{query}%"
        by_title = select(Conversation.id).where(col(Conversation.title).ilike(needle))
        by_first = select(ConversationMessage.conversation_id).where(
            ConversationMessage.seq == 0,
            col(ConversationMessage.content_json).ilike(needle),
        )
        by_entry = (
            select(Conversation.id)
            .join(
                TranscriptEntry,
                col(TranscriptEntry.session_id) == col(Conversation.session_id),
            )
            .where(sa_cast(col(TranscriptEntry.entry), String).ilike(needle))
        )
        stmt = (
            select(Conversation)
            .where(col(Conversation.id).in_(by_title.union(by_first, by_entry)))
            .order_by(col(Conversation.id).desc())
            .limit(limit)
        )
        async with self._db.session() as session:
            return list((await session.exec(stmt)).all())

    async def bindings(self, conv: Conversation) -> list[ConversationBinding]:
        async with self._db.session() as session:
            result = await session.exec(
                select(ConversationBinding)
                .where(ConversationBinding.conversation_id == conv.id)
                .order_by(col(ConversationBinding.id))
            )
            return list(result.all())

    async def bind(
        self,
        conv: Conversation,
        *,
        frontend: str,
        external_id: str,
        visible: bool = True,
    ) -> ConversationBinding:
        if conv.kind not in self.frontend(frontend).kinds:
            msg = f"frontend {frontend!r} does not show kind {conv.kind!r}"
            raise ValueError(msg)
        async with self._db.session() as session:
            existing = list(
                (
                    await session.exec(
                        select(ConversationBinding).where(
                            ConversationBinding.conversation_id == conv.id,
                            ConversationBinding.frontend == frontend,
                        )
                    )
                ).all()
            )
            row = next((b for b in existing if b.external_id == external_id), None)
            if visible:
                for other in existing:
                    if other is not row and other.visible:
                        other.visible = False
                        session.add(other)
                same_window = await session.exec(
                    select(ConversationBinding).where(
                        ConversationBinding.frontend == frontend,
                        ConversationBinding.external_id == external_id,
                        ConversationBinding.conversation_id != conv.id,
                        col(ConversationBinding.visible).is_(True),
                    )
                )
                for other in same_window.all():
                    other.visible = False
                    session.add(other)
            if row is None:
                row = ConversationBinding(
                    conversation_id=cast("int", conv.id),
                    frontend=frontend,
                    external_id=external_id,
                    visible=visible,
                )
            else:
                row.visible = visible
            session.add(row)
            await session.commit()
            await session.refresh(row)
        self._bus.publish(
            "conversation.bound",
            conversation_id=conv.external_id,
            frontend=frontend,
            external_id=external_id,
            visible=visible,
        )
        return row

    async def find_bound(
        self, *, frontend: str, external_id: str
    ) -> Conversation | None:
        async with self._db.session() as session:
            result = await session.exec(
                select(Conversation)
                .join(
                    ConversationBinding,
                    col(ConversationBinding.conversation_id) == col(Conversation.id),
                )
                .where(
                    ConversationBinding.frontend == frontend,
                    ConversationBinding.external_id == external_id,
                    col(ConversationBinding.visible).is_(True),
                )
                .order_by(col(Conversation.id).desc())
            )
            return result.first()

    async def last_binding(
        self, *, frontend: str, kind: str
    ) -> ConversationBinding | None:
        """The window ``frontend`` last used for ``kind``; outlives a rotation."""
        async with self._db.session() as session:
            result = await session.exec(
                select(ConversationBinding)
                .join(
                    Conversation,
                    col(Conversation.id) == col(ConversationBinding.conversation_id),
                )
                .where(
                    ConversationBinding.frontend == frontend, Conversation.kind == kind
                )
                .order_by(col(ConversationBinding.id).desc())
            )
            return result.first()

    async def window_of(self, conv: Conversation) -> str | None:
        for binding in await self.bindings(conv):
            if binding.visible:
                return binding.external_id
        return None

    async def set_flags(
        self, conv: Conversation, flags: dict[str, Any]
    ) -> Conversation:
        async def apply(row: Conversation) -> None:
            row.flags = {**row.flags, **flags}

        return await self._update(conv, apply)

    async def set_status(self, conv: Conversation, status: str) -> Conversation:
        if status not in STATUSES:
            msg = f"unknown status {status!r}"
            raise ValueError(msg)

        async def apply(row: Conversation) -> None:
            row.status = status

        return await self._update(conv, apply)

    async def set_title(self, conv: Conversation, title: str) -> Conversation:
        async def apply(row: Conversation) -> None:
            row.title = title

        return await self._update(conv, apply)

    async def touch_user(self, conv: Conversation) -> Conversation:
        async def apply(row: Conversation) -> None:
            row.last_user_activity_at = datetime.now(UTC)

        return await self._update(conv, apply)

    async def reparent(self, conv: Conversation, parent: Conversation) -> Conversation:
        async def apply(row: Conversation) -> None:
            row.parent_id = parent.id

        return await self._update(conv, apply)

    async def _update(
        self, conv: Conversation, apply: Callable[[Conversation], Awaitable[None]]
    ) -> Conversation:
        async with self._db.session() as session:
            row = await session.get(Conversation, conv.id)
            if row is None:
                msg = f"conversation {conv.external_id} vanished"
                raise LookupError(msg)
            await apply(row)
            row.updated_at = datetime.now(UTC)
            session.add(row)
            await session.commit()
            await session.refresh(row)
        self._bus.publish("conversation.updated", **self.public(row))
        return row

    def public(self, conv: Conversation) -> dict[str, Any]:
        return {
            "id": conv.external_id,
            "kind": conv.kind,
            "agent": conv.agent_name,
            "title": conv.title,
            "status": conv.status,
            "parent_row": conv.parent_id,
            "session_id": conv.session_id,
            "running_turn": conv.running_turn,
            "pending_question": conv.pending_question,
            "flags": conv.flags,
            "origin": conv.frontend,
            "created_at": iso(conv.created_at),
            "last_user_activity_at": iso(conv.last_user_activity_at),
            "last_activity_at": iso(conv.last_activity_at),
        }

    async def describe(self, conv: Conversation) -> dict[str, Any]:
        out = self.public(conv)
        out["title"] = await self.implied_title(conv)
        out["context_tokens"] = await self.context_tokens(conv)
        parent = await self.get_row(conv.parent_id) if conv.parent_id else None
        out["parent"] = parent.external_id if parent is not None else None
        out["bindings"] = [
            {"frontend": b.frontend, "external_id": b.external_id, "visible": b.visible}
            for b in await self.bindings(conv)
        ]
        live = self._pool.get(conv.external_id)
        agent = self._agents.get(conv.agent_name)
        out["thinking_is_progress"] = agent is not None and agent.thinking_is_progress
        out["live"] = live is not None
        out["busy"] = live.busy if live is not None else False
        runner = self._runners.get(cast("int", conv.id))
        out["turn"] = runner.snapshot() if runner is not None else None
        pending = self.pending_question(conv.external_id)
        out["question"] = (
            {"id": pending[0], "questions": pending[1]} if pending else None
        )
        return out

    def pending_question(self, key: str) -> tuple[str, list[dict[str, Any]]] | None:
        for question_id, pending in self._questions.items():
            if pending.conversation_id == key and not pending.answer.done():
                return question_id, pending.questions
        return None

    async def rate_limits(self, *, limit: int = 100) -> list[RateLimit]:
        async with self._db.session() as session:
            result = await session.exec(
                select(RateLimit).order_by(col(RateLimit.id).desc()).limit(limit)
            )
            return list(result.all())

    async def context_tokens(self, conv: Conversation) -> int:
        async with self._db.session() as session:
            row = (
                await session.exec(
                    select(Usage)
                    .where(Usage.conversation_id == conv.external_id)
                    .order_by(col(Usage.id).desc())
                    .limit(1)
                )
            ).first()
        return context_of(row)

    async def context_tokens_by_id(self, ids: Iterable[str]) -> dict[str, int]:
        wanted = list(ids)
        if not wanted:
            return {}
        latest = (
            select(func.max(Usage.id))
            .where(col(Usage.conversation_id).in_(wanted))
            .group_by(Usage.conversation_id)
        )
        async with self._db.session() as session:
            rows = (
                await session.exec(select(Usage).where(col(Usage.id).in_(latest)))
            ).all()
        return {
            cast("str", row.conversation_id): context_of(row)
            for row in rows
            if row.conversation_id
        }

    async def usage_tokens(self, since: datetime) -> int:
        async with self._db.session() as session:
            rows = (
                await session.exec(
                    select(Usage).where(
                        col(Usage.ts) >= since.astimezone(UTC).replace(tzinfo=None)
                    )
                )
            ).all()
        return sum(
            r.input_tokens + r.output_tokens + r.cache_creation_tokens for r in rows
        )

    async def busy(self, conv: Conversation) -> bool:
        row = await self.get_row(cast("int", conv.id)) or conv
        if row.running_turn or row.pending_question:
            return True
        live = self._pool.get(row.external_id)
        if live is not None and live.busy:
            return True
        pending = await self._queue.pending(cast("int", row.id))
        return any(i.priority in ("user", "urgent", "wake") for i in pending)

    @property
    def frontends(self) -> list[Frontend]:
        return list(self._frontends)

    def frontend(self, name: str) -> Frontend:
        for fe in self._frontends:
            if fe.name == name:
                return fe
        msg = f"unknown frontend {name!r}"
        raise ValueError(msg)

    def default_agent(self, kind: Kind) -> str | None:
        for fe in self._frontends:
            if kind in fe.kinds and (agent := fe.agent_for(kind)):
                return agent
        return None

    async def materialize(self, conv: Conversation) -> ConversationBinding | None:
        for fe in self._frontends:
            if conv.kind not in fe.kinds:
                continue
            binding = await fe.materialize(conv)
            if binding is not None:
                return binding
        return None

    async def mark_closed(self, conv: Conversation) -> bool:
        marked = False
        for fe in self._frontends:
            if conv.kind in fe.kinds:
                try:
                    marked = await fe.mark_closed(conv) or marked
                except Exception:  # noqa: BLE001
                    _log.exception("%s could not mark %s", fe.name, conv.external_id)
        return marked

    async def read(self, conv: Conversation, *, window: int | None = None) -> str:
        return render_messages(
            await self.history(conv),
            thinking_is_progress=self._claude_agent(
                conv.agent_name
            ).thinking_is_progress,
            window=window,
        )

    async def history(self, conv: Conversation) -> list[dict[str, Any]]:
        if conv.session_id is None:
            async with self._db.session() as session:
                return await load_messages(
                    session, conversation_id=cast("int", conv.id)
                )
        messages = messages_from_entries(
            cast("Any", await self.entries(conv)), stamped=True
        )
        return [self._unwrapped(message) for message in messages]

    def _unwrapped(self, message: dict[str, Any]) -> dict[str, Any]:
        if message["role"] != "user":
            return message
        content = message["content"]
        if isinstance(content, str):
            text, service = split_service(content, self._texts)
            return (
                message
                if service is None
                else {**message, "content": text, "service": service}
            )
        for index, block in enumerate(content):
            if block.get("type") != "text":
                continue
            text, service = split_service(str(block.get("text") or ""), self._texts)
            if service is not None:
                blocks = [*content]
                blocks[index] = {**block, "text": text}
                return {**message, "content": blocks, "service": service}
        return message

    async def entries(self, conv: Conversation, *, subpath: str = "") -> list[Any]:
        if conv.session_id is None:
            return []
        key = {**self._store_key(conv), "subpath": subpath}
        return list(await self._store.load(cast("Any", key)) or [])

    async def subpaths(self, conv: Conversation) -> list[str]:
        if conv.session_id is None:
            return []
        return list(await self._store.list_subkeys(cast("Any", self._store_key(conv))))

    async def first_user_texts(self, ids: Iterable[int]) -> dict[int, str]:
        wanted = list(ids)
        if not wanted:
            return {}
        async with self._db.session() as session:
            rows = (
                await session.exec(
                    select(ConversationMessage).where(
                        col(ConversationMessage.conversation_id).in_(wanted),
                        ConversationMessage.seq == 0,
                        ConversationMessage.role == "user",
                    )
                )
            ).all()
        return {
            r.conversation_id: text_of(json.loads(r.content_json)).strip() for r in rows
        }

    async def implied_title(self, conv: Conversation) -> str | None:
        if conv.title:
            return conv.title
        text = (await self.first_user_texts([cast("int", conv.id)])).get(
            cast("int", conv.id)
        )
        return implied_title(text)

    async def chat_name(self, conv: Conversation) -> str:
        """What a ``[[wikilink]]`` to the chat says: the file's stem when it has one."""
        window = await self.window_of(conv)
        if window and window.endswith(".md"):
            return window.rsplit("/", 1)[-1][: -len(".md")]
        return await self.implied_title(conv) or conv.external_id

    async def adopt(self, *, kind: Kind, first_user_text: str) -> Conversation | None:
        """The one unbound, session-less conversation whose history starts here."""
        text = first_user_text.strip()
        if not text:
            return None
        bound = select(ConversationBinding.conversation_id).where(
            col(ConversationBinding.visible).is_(True)
        )
        async with self._db.session() as session:
            rows = (
                await session.exec(
                    select(Conversation).where(
                        Conversation.kind == kind,
                        Conversation.status == "open",
                        col(Conversation.session_id).is_(None),
                        col(Conversation.id).not_in(bound),
                    )
                )
            ).all()
        firsts = await self.first_user_texts(cast("int", r.id) for r in rows)
        hits = [r for r in rows if firsts.get(cast("int", r.id)) == text]
        return hits[0] if len(hits) == 1 else None


def implied_title(text: str | None) -> str | None:
    if not text:
        return None
    line = text.strip().splitlines()[0].strip()
    return line if len(line) <= TITLE_MAX else line[: TITLE_MAX - 1] + "…"


def context_of(row: Usage | None) -> int:
    """The last API call's input, or the per-call average for older rows."""
    if row is None:
        return 0
    if row.context_tokens:
        return row.context_tokens
    total = row.input_tokens + row.cache_read_tokens + row.cache_creation_tokens
    return round(total / max(row.num_turns or 1, 1))