"""Stored conversation history for the markdown frontend.

Aligns the file the user is editing against the DB-stored Anthropic-shape
message history, splicing text edits or forking on structural changes.
"""

from __future__ import annotations

import json
import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, cast

from sqlalchemy import delete
from sqlmodel import select

from beaver_gateway.backends.transcript import spoken_text
from beaver_gateway.storage.models import Conversation, ConversationMessage

if TYPE_CHECKING:
    from anthropic.types import MessageParam
    from sqlmodel.ext.asyncio.session import AsyncSession

    from beaver_gateway.frontends.markdown.parser import ParsedTurn

__all__ = [
    "ForkOutcome",
    "diff_and_fork",
    "load_conversation",
    "load_messages",
    "mint_conversation",
    "rewrite_messages",
    "set_session_id",
]


@dataclass(frozen=True, slots=True)
class ForkOutcome:
    """Result of aligning the incoming file against stored history.

    ``messages`` is the backend input; ``persist_messages`` is history to
    store before the new reply; ``divergence_index`` is where they diverged.
    """

    messages: list[MessageParam]
    persist_messages: list[dict[str, Any]]
    divergence_index: int | None
    edited: bool = False
    """An earlier assistant turn's prose was rewritten in the file."""

    @property
    def reuse_session(self) -> bool:
        return self.divergence_index is None and not self.edited


async def load_conversation(
    session: AsyncSession, *, frontend: str, external_id: str
) -> Conversation | None:
    stmt = (
        select(Conversation)
        .where(Conversation.frontend == frontend)
        .where(Conversation.external_id == external_id)
    )
    result = await session.exec(stmt)
    return result.first()


async def mint_conversation(
    session: AsyncSession, *, frontend: str, agent_name: str
) -> Conversation:
    """Create a fresh conversation row with a new uuid for external_id.

    Caller must persist the returned ``external_id`` so future requests can find it.
    """
    row = Conversation(
        frontend=frontend, external_id=str(uuid.uuid4()), agent_name=agent_name
    )
    session.add(row)
    await session.commit()
    await session.refresh(row)
    return row


async def set_session_id(
    session: AsyncSession, *, conversation_id: int, session_id: str | None
) -> None:
    conv = await session.get(Conversation, conversation_id)
    if conv is None or conv.session_id == session_id:
        return
    conv.session_id = session_id
    session.add(conv)
    await session.commit()


async def load_messages(
    session: AsyncSession, *, conversation_id: int
) -> list[dict[str, Any]]:
    """Return stored messages ordered by ``seq`` ascending.

    Each entry is a canonical Anthropic ``MessageParam`` dict, the same
    shape fed to the backend on continuation.
    """
    stmt = (
        select(ConversationMessage)
        .where(ConversationMessage.conversation_id == conversation_id)
        .order_by(ConversationMessage.seq.asc())  # ty: ignore[unresolved-attribute]
    )
    result = await session.exec(stmt)
    rows = result.all()
    return [
        {"role": r.role, "content": _sanitize_content(json.loads(r.content_json))}
        for r in rows
    ]


def _sanitize_content(content: Any) -> Any:
    """Strip wire-illegal fields from stored Anthropic content blocks.

    Drops ``tool_result.is_error: null`` — the API rejects null there
    though the field is optional.
    """
    if not isinstance(content, list):
        return content
    cleaned: list[Any] = []
    for blk in content:
        out_blk = blk
        if (
            isinstance(blk, dict)
            and blk.get("type") == "tool_result"
            and blk.get("is_error") is None
            and "is_error" in blk
        ):
            out_blk = {k: v for k, v in blk.items() if k != "is_error"}
        cleaned.append(out_blk)
    return cleaned


async def rewrite_messages(
    session: AsyncSession, *, conversation_id: int, messages: list[dict[str, Any]]
) -> None:
    """Replace the conversation's stored messages (full overwrite, no branch history).

    Deletes and flushes before inserting — SQLAlchemy's default INSERT-before-DELETE
    flush order would otherwise collide with ``uq_msg_conv_seq``.
    """
    await session.execute(  # ty: ignore[deprecated]
        delete(ConversationMessage).where(
            ConversationMessage.conversation_id == conversation_id  # ty: ignore[invalid-argument-type]
        )
    )
    await session.flush()
    for seq, m in enumerate(messages):
        session.add(
            ConversationMessage(
                conversation_id=conversation_id,
                seq=seq,
                role=str(m["role"]),
                content_json=json.dumps(m["content"], separators=(",", ":")),
            )
        )
    conv = await session.get(Conversation, conversation_id)
    if conv is not None:
        from datetime import UTC, datetime

        conv.updated_at = datetime.now(UTC)
        session.add(conv)
    await session.commit()


@dataclass(frozen=True, slots=True)
class _StoredDisplayTurn:
    """A "display turn" reconstructed from stored raw messages.

    Parser-equivalent view used to diff the file against the DB.
    """

    role: str
    messages: tuple[dict[str, Any], ...]
    spoken_text: str
    skeleton: tuple[str, ...]
    text_segment_count: int


def _group_display_turns(
    stored: list[dict[str, Any]], *, thinking_is_progress: bool
) -> list[_StoredDisplayTurn]:
    """Walk raw stored messages, group them into Obsidian-visible turns.

    A user-prompt message (no ``tool_result`` blocks) opens a new turn;
    tool-result-only messages roll into the current assistant turn.
    """
    out: list[_StoredDisplayTurn] = []
    i = 0
    while i < len(stored):
        msg = stored[i]
        role = msg["role"]
        if role == "user" and _is_user_prompt(msg.get("content")):
            out.append(
                _StoredDisplayTurn(
                    role="user",
                    messages=(msg,),
                    spoken_text=_user_prompt_text(msg.get("content")),
                    skeleton=(),
                    text_segment_count=0,
                )
            )
            i += 1
            continue
        group: list[dict[str, Any]] = []
        while i < len(stored):
            m = stored[i]
            if m["role"] == "user" and _is_user_prompt(m.get("content")):
                break
            group.append(m)
            i += 1
        spoken, skeleton, text_count = _summarize_assistant_group(
            group, thinking_is_progress=thinking_is_progress
        )
        out.append(
            _StoredDisplayTurn(
                role="assistant",
                messages=tuple(group),
                spoken_text=spoken,
                skeleton=tuple(skeleton),
                text_segment_count=text_count,
            )
        )
    return out


def _is_user_prompt(content: Any) -> bool:
    """A user message is a *prompt* unless its content carries tool_result blocks.

    Unknown content shapes are conservatively treated as a prompt.
    """
    if isinstance(content, str):
        return True
    if isinstance(content, list):
        return not any(
            isinstance(b, dict) and b.get("type") == "tool_result" for b in content
        )
    return True


def _user_prompt_text(content: Any) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        chunks = [
            str(b.get("text", ""))
            for b in content
            if isinstance(b, dict) and b.get("type") == "text"
        ]
        return "\n\n".join(c for c in chunks if c)
    return ""


def _summarize_assistant_group(
    group: list[dict[str, Any]], *, thinking_is_progress: bool
) -> tuple[str, list[str], int]:
    """Compute (spoken_text, tool_skeleton, text_segment_count) for a display group.

    Must mirror what the renderer writes and the parser reads back: text and
    progress thinking collapse into one segment, tool_use becomes a skeleton
    entry, tool_result, reasoning and empty thinking are invisible.
    """
    from beaver_gateway.frontends.markdown.parser import TextSegment, ToolSegment

    segments: list[TextSegment | ToolSegment] = []
    pending: list[str] = []

    def _flush() -> None:
        if not pending:
            return
        joined = "\n\n".join(p for p in pending if p)
        pending.clear()
        cleaned = joined.strip()
        if cleaned:
            segments.append(TextSegment(text=cleaned))

    for msg in group:
        if msg["role"] == "user":
            _flush()
            continue
        content = msg.get("content")
        if not isinstance(content, list):
            continue
        for blk in content:
            if not isinstance(blk, dict):
                continue
            text = spoken_text(blk, thinking_is_progress=thinking_is_progress)
            if text:
                pending.append(text)
            elif blk.get("type") == "tool_use":
                _flush()
                segments.append(ToolSegment(name=str(blk.get("name", ""))))
    _flush()
    spoken_chunks = [s.text for s in segments if isinstance(s, TextSegment)]
    spoken = "\n\n".join(c for c in spoken_chunks if c).strip()
    skeleton = [s.name for s in segments if isinstance(s, ToolSegment)]
    text_count = sum(1 for s in segments if isinstance(s, TextSegment))
    return spoken, skeleton, text_count


def diff_and_fork(
    *,
    stored: list[dict[str, Any]],
    incoming: list[ParsedTurn],
    thinking_is_progress: bool,
) -> ForkOutcome:
    """Align the incoming parsed file against stored history.

    ``incoming`` must end with a user turn (the new prompt); raises otherwise.
    Segment-class imports below are deferred to avoid a cycle with ``parser``.
    """
    if not incoming or incoming[-1].role != "user":
        msg = (
            "diff_and_fork expects incoming to end with a user turn "
            "(the new prompt); got "
            f"{incoming[-1].role if incoming else 'empty'}"
        )
        raise ValueError(msg)

    stored_groups = _group_display_turns(
        stored, thinking_is_progress=thinking_is_progress
    )
    new_user_turn = incoming[-1]
    prior_incoming = incoming[:-1]

    spliced_groups, divergence, edited = _walk_prefix(
        prior_incoming, stored_groups, thinking_is_progress=thinking_is_progress
    )

    if divergence is None and len(prior_incoming) < len(stored_groups):
        if _file_lags_store(stored_groups, len(prior_incoming), new_user_turn):
            spliced_groups.extend(
                list(g.messages) for g in stored_groups[len(prior_incoming) :]
            )
        else:
            divergence = len(prior_incoming)

    backend_msgs, persist_msgs = _assemble_tail(
        spliced_groups=spliced_groups,
        prior_incoming=prior_incoming,
        divergence=divergence,
        new_user_turn=new_user_turn,
    )
    return ForkOutcome(
        messages=backend_msgs,
        persist_messages=persist_msgs,
        divergence_index=divergence,
        edited=edited,
    )


def _file_lags_store(
    stored_groups: list[_StoredDisplayTurn], prior_len: int, new_user_turn: ParsedTurn
) -> bool:
    """Is the shorter incoming file a stale view rather than a deletion?

    Tell: if the DB already holds the submitted prompt at the position the
    file stops at, it's a lagging render, not a deletion — misjudging this
    forks history, breaks the fingerprint, and respawns the backend session.
    """
    if prior_len >= len(stored_groups):
        return False
    candidate = stored_groups[prior_len]
    return candidate.role == "user" and candidate.spoken_text == new_user_turn.text


def _walk_prefix(
    prior_incoming: list[ParsedTurn],
    stored_groups: list[_StoredDisplayTurn],
    *,
    thinking_is_progress: bool,
) -> tuple[list[list[dict[str, Any]]], int | None, bool]:
    """Walk incoming vs stored side-by-side until first divergence.

    An empty incoming skeleton (no tool callouts rendered) means prose
    alone decides the match; a mismatch forces a fresh backend session.
    """
    from beaver_gateway.frontends.markdown.parser import TextSegment, ToolSegment

    spliced_groups: list[list[dict[str, Any]]] = []
    edited = False
    for i, inc in enumerate(prior_incoming):
        if i >= len(stored_groups):
            return spliced_groups, i, edited
        st = stored_groups[i]
        if inc.role != st.role:
            return spliced_groups, i, edited
        if inc.role == "user":
            if inc.text != st.spoken_text:
                return spliced_groups, i, edited
            spliced_groups.append(list(st.messages))
            continue
        inc_skeleton = tuple(
            s.name for s in inc.structure if isinstance(s, ToolSegment)
        )
        inc_text_count = sum(1 for s in inc.structure if isinstance(s, TextSegment))
        if inc_skeleton and inc_skeleton != st.skeleton:
            return spliced_groups, i, edited
        if inc.text == st.spoken_text:
            spliced_groups.append(list(st.messages))
            continue
        if inc_skeleton and inc_text_count != st.text_segment_count:
            return spliced_groups, i, edited
        spliced = _splice_assistant_group(
            stored_group=st, incoming=inc, thinking_is_progress=thinking_is_progress
        )
        if spliced is None:
            return spliced_groups, i, edited
        spliced_groups.append(spliced)
        edited = True
    return spliced_groups, None, edited


def _assemble_tail(
    *,
    spliced_groups: list[list[dict[str, Any]]],
    prior_incoming: list[ParsedTurn],
    divergence: int | None,
    new_user_turn: ParsedTurn,
) -> tuple[list[MessageParam], list[dict[str, Any]]]:
    """Build the (backend, persist) lists from aligned + post-divergence tail."""
    backend_msgs: list[MessageParam] = []
    persist_msgs: list[dict[str, Any]] = []
    for spliced in spliced_groups:
        for m in spliced:
            entry: dict[str, Any] = {"role": m["role"], "content": m["content"]}
            backend_msgs.append(cast("MessageParam", entry))
            persist_msgs.append(entry)
    if divergence is not None:
        for inc in prior_incoming[divergence:]:
            if not inc.text:
                continue
            entry = {"role": inc.role, "content": inc.text}
            backend_msgs.append(cast("MessageParam", entry))
            persist_msgs.append(entry)
    backend_msgs.append({"role": "user", "content": new_user_turn.text})
    return backend_msgs, persist_msgs


def _splice_assistant_group(
    *,
    stored_group: _StoredDisplayTurn,
    incoming: ParsedTurn,
    thinking_is_progress: bool,
) -> list[dict[str, Any]] | None:
    """Apply the file's prose to a stored assistant turn.

    Prefers substituting text in place, which keeps everything the
    markdown can't express. Falls back to rebuilding from the file's
    structure when the text blocks don't line up one-to-one.
    """
    preserved = _splice_in_place(
        stored_group=stored_group,
        incoming=incoming,
        thinking_is_progress=thinking_is_progress,
    )
    if preserved is not None:
        return preserved
    return _splice_by_rebuild(stored_group=stored_group, incoming=incoming)


def _splice_in_place(
    *,
    stored_group: _StoredDisplayTurn,
    incoming: ParsedTurn,
    thinking_is_progress: bool,
) -> list[dict[str, Any]] | None:
    """Copy the stored messages, swapping only their text block contents.

    Preserves message/thinking-block boundaries that a rebuild would lose
    (and that a cache-hit continuation depends on); returns ``None`` if
    text blocks aren't 1:1 with incoming segments or the prose runs through
    signed thinking blocks, so the caller rebuilds instead.
    """
    from beaver_gateway.frontends.markdown.parser import TextSegment

    new_texts = [
        s.text for s in incoming.structure if isinstance(s, TextSegment) and s.text
    ]
    positions: list[tuple[int, int]] = []
    for mi, msg in enumerate(stored_group.messages):
        content = msg.get("content")
        if msg["role"] != "assistant" or not isinstance(content, list):
            continue
        for bi, blk in enumerate(cast("list[dict[str, Any]]", content)):
            if not spoken_text(blk, thinking_is_progress=thinking_is_progress):
                continue
            if blk["type"] == "thinking":
                return None
            positions.append((mi, bi))
    if len(positions) != len(new_texts):
        return None

    out: list[dict[str, Any]] = [
        {
            "role": m["role"],
            "content": list(m["content"])
            if isinstance(m.get("content"), list)
            else m.get("content"),
        }
        for m in stored_group.messages
    ]
    for (mi, bi), text in zip(positions, new_texts, strict=True):
        out[mi]["content"][bi] = {**out[mi]["content"][bi], "text": text}
    return out


def _splice_by_rebuild(
    *, stored_group: _StoredDisplayTurn, incoming: ParsedTurn
) -> list[dict[str, Any]] | None:
    """Rebuild an assistant display turn with new text + stored tool_use blocks.

    Matches tool_use blocks to incoming ``ToolSegment``s by position; returns
    ``None`` (caller forks) if a matching stored ``tool_result`` is missing.
    """
    from beaver_gateway.frontends.markdown.parser import TextSegment

    tool_uses, tool_results_by_id = _harvest_tool_blocks(stored_group)

    spliced: list[dict[str, Any]] = []
    current_asst: list[dict[str, Any]] = []
    next_tool = 0
    for seg in incoming.structure:
        if isinstance(seg, TextSegment):
            if seg.text:
                current_asst.append({"type": "text", "text": seg.text})
            continue
        if next_tool >= len(tool_uses):
            return None
        tu = tool_uses[next_tool]
        next_tool += 1
        current_asst.append(tu)
        spliced.append({"role": "assistant", "content": current_asst})
        current_asst = []
        tr = tool_results_by_id.get(str(tu.get("id", "")))
        if tr is None:
            return None
        spliced.append({"role": "user", "content": [tr]})
    if current_asst:
        spliced.append({"role": "assistant", "content": current_asst})
    elif not spliced:
        return None
    return spliced


def _harvest_tool_blocks(
    stored_group: _StoredDisplayTurn,
) -> tuple[list[dict[str, Any]], dict[str, dict[str, Any]]]:
    """Pull stored ``tool_use`` blocks (ordered) and ``tool_result`` blocks (by id)."""
    tool_uses: list[dict[str, Any]] = []
    tool_results_by_id: dict[str, dict[str, Any]] = {}
    for msg in stored_group.messages:
        content = msg.get("content")
        if not isinstance(content, list):
            continue
        if msg["role"] == "assistant":
            tool_uses.extend(
                blk
                for blk in content
                if isinstance(blk, dict) and blk.get("type") == "tool_use"
            )
            continue
        for blk in content:
            if not isinstance(blk, dict) or blk.get("type") != "tool_result":
                continue
            tid = blk.get("tool_use_id")
            if isinstance(tid, str):
                tool_results_by_id[tid] = blk
    return tool_uses, tool_results_by_id