"""Orphan tool result prevention for crash recovery.

When an agent crashes mid-tool-execution, the Anthropic Messages API
conversation history may contain ``tool_use`` blocks from the assistant
that have no matching ``tool_result`` block in the subsequent user turn.
This leaves the transcript in an inconsistent state: the API requires
every ``tool_use`` block to be answered before the next model turn can
proceed.

This module detects such *orphaned* ``tool_use`` blocks and generates
synthetic ``tool_result`` entries so the transcript can be safely resumed.
The synthetic results are clearly annotated so downstream code can
distinguish them from real tool output.

Typical usage in crash recovery::

    from bernstein.core.orphan_tool_result import repair_transcript

    # Load the persisted conversation history
    messages = load_conversation_history(session_id)

    # Repair any orphans introduced by the crash
    result = repair_transcript(messages)
    if result.orphan_count > 0:
        logger.warning(
            "Repaired %d orphaned tool_use block(s) for session %s: %s",
            result.orphan_count,
            session_id,
            result.orphan_ids,
        )

    # Pass the repaired transcript to the resumed agent
    resume_agent(result.messages)
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Any, cast

logger = logging.getLogger(__name__)

# Content injected into synthetic tool_result blocks so observers can
# identify them as recovery artefacts rather than real tool output.
SYNTHETIC_RESULT_MARKER: str = "[synthetic_orphan_recovery]"
_DEFAULT_SYNTHETIC_CONTENT: str = (
    f"{SYNTHETIC_RESULT_MARKER} "
    "Tool result unavailable — agent was interrupted during execution. "
    "This is a synthetic placeholder generated by orphan recovery."
)


# ---------------------------------------------------------------------------
# Public data types
# ---------------------------------------------------------------------------


@dataclass
class OrphanRepairResult:
    """Outcome of repairing orphaned ``tool_use`` blocks in a transcript.

    Args:
        messages: The repaired conversation history.  Identical to the
            input when no orphans were found.
        orphan_count: Number of orphaned ``tool_use`` blocks that received
            a synthetic ``tool_result``.
        orphan_ids: Ordered list of ``tool_use`` IDs that were repaired.
    """

    messages: list[dict[str, Any]]
    orphan_count: int
    orphan_ids: list[str] = field(default_factory=list[str])


# ---------------------------------------------------------------------------
# Detection
# ---------------------------------------------------------------------------


def find_orphaned_tool_uses(
    messages: list[dict[str, Any]],
) -> list[tuple[int, str, str]]:
    """Identify ``tool_use`` blocks that have no matching ``tool_result``.

    Walks the conversation in order.  For each ``assistant`` message,
    records every ``tool_use`` block's ``id`` and ``name``.  For each
    ``user`` message, marks resolved any ``tool_use_id`` that appears in a
    ``tool_result`` block.

    Args:
        messages: Conversation history in Anthropic Messages API format —
            a list of ``{"role": ..., "content": [...]}`` dicts.

    Returns:
        A list of ``(message_index, tool_use_id, tool_name)`` tuples for
        every unresolved ``tool_use`` block, in the order they appear in
        the transcript.  An empty list means the transcript is consistent.
    """
    # Ordered list of (msg_index, tool_use_id, tool_name)
    pending: list[tuple[int, str, str]] = []
    resolved_ids: set[str] = set()

    for idx, msg in enumerate(messages):
        role = msg.get("role", "")
        content = msg.get("content", [])

        # ``content`` may be a plain string for simple text-only messages
        if isinstance(content, str):
            continue

        if role == "assistant":
            for block in content:
                if not isinstance(block, dict):
                    continue
                b = cast("dict[str, Any]", block)
                if b.get("type") == "tool_use":
                    tool_id = cast("str", b.get("id", ""))
                    tool_name = cast("str", b.get("name", "unknown_tool"))
                    if tool_id:
                        pending.append((idx, tool_id, tool_name))

        elif role == "user":
            for block in content:
                if not isinstance(block, dict):
                    continue
                b = cast("dict[str, Any]", block)
                if b.get("type") == "tool_result":
                    tid = cast("str", b.get("tool_use_id", ""))
                    if tid:
                        resolved_ids.add(tid)

    return [(idx, tid, name) for idx, tid, name in pending if tid not in resolved_ids]


# ---------------------------------------------------------------------------
# Repair
# ---------------------------------------------------------------------------


def repair_transcript(
    messages: list[dict[str, Any]],
    *,
    synthetic_content: str | None = None,
) -> OrphanRepairResult:
    """Insert synthetic ``tool_result`` blocks for orphaned ``tool_use`` IDs.

    For each orphaned ``tool_use`` block the function inserts a new
    ``user`` message immediately after the ``assistant`` message that
    contains the orphan.  When multiple orphans share the same assistant
    message they are collected into a single synthetic ``user`` message to
    keep the transcript compact.

    The transcript is not modified in-place; a shallow-copy list is
    returned so callers can compare the before/after versions.

    Args:
        messages: Conversation history in Anthropic Messages API format.
        synthetic_content: Text to use as the ``content`` of each
            synthetic ``tool_result`` block.  Defaults to
            :data:`_DEFAULT_SYNTHETIC_CONTENT`.

    Returns:
        :class:`OrphanRepairResult` with the repaired ``messages`` list
        and metadata about what was changed.

    Edge cases:
        - If ``messages`` is empty or has no orphans, the original list
          (shallow copy) is returned unchanged with ``orphan_count == 0``.
        - Duplicate ``tool_use`` IDs within the same message are each
          given their own ``tool_result`` entry.
        - Partial-stream ``tool_use`` blocks whose ``id`` field is the
          empty string are skipped (they cannot be matched by the API).
    """
    content_text = synthetic_content if synthetic_content is not None else _DEFAULT_SYNTHETIC_CONTENT

    orphans = find_orphaned_tool_uses(messages)
    if not orphans:
        return OrphanRepairResult(
            messages=list(messages),
            orphan_count=0,
            orphan_ids=[],
        )

    # Group orphans by the assistant message index they belong to so we
    # can insert a single synthetic user message per assistant turn.
    # Use dict to preserve insertion order (Python 3.7+).
    by_msg_idx: dict[int, list[tuple[str, str]]] = {}
    for msg_idx, tool_id, tool_name in orphans:
        by_msg_idx.setdefault(msg_idx, []).append((tool_id, tool_name))

    orphan_ids = [tid for _, tid, _ in orphans]

    # Build the repaired message list by inserting synthetic user messages
    # after each affected assistant turn.  We iterate in reverse order of
    # insertion point so earlier indices stay valid as we splice.
    repaired = list(messages)
    for msg_idx in sorted(by_msg_idx.keys(), reverse=True):
        orphaned_blocks = by_msg_idx[msg_idx]
        synthetic_results: list[dict[str, Any]] = []
        for tool_id, tool_name in orphaned_blocks:
            logger.info(
                "Orphan recovery: inserting synthetic tool_result for tool_use id=%r name=%r (message index %d)",
                tool_id,
                tool_name,
                msg_idx,
            )
            synthetic_results.append(
                {
                    "type": "tool_result",
                    "tool_use_id": tool_id,
                    "content": content_text,
                    "is_error": False,
                    "_synthetic": True,  # non-API field for internal tracing
                }
            )

        synthetic_user_msg: dict[str, Any] = {
            "role": "user",
            "content": synthetic_results,
            "_orphan_recovery": True,  # non-API field for internal tracing
        }
        # Insert immediately after the assistant message
        repaired.insert(msg_idx + 1, synthetic_user_msg)

    return OrphanRepairResult(
        messages=repaired,
        orphan_count=len(orphan_ids),
        orphan_ids=orphan_ids,
    )
