"""
cortex/hippo/transcript_repair.py — Tool use/result pairing repair for assembled context.

Ported from lossless-claw:
  extensions/lossless-claw/src/transcript-repair.ts

Called on every assembled message list before returning to ensure Anthropic
(and other providers) don't reject transcripts where assistant tool calls are
not immediately followed by matching tool results.
"""

from __future__ import annotations

import time
from typing import Any


# -- Constants --

TOOL_CALL_TYPES: frozenset[str] = frozenset([
    "toolCall",
    "toolUse",
    "tool_use",
    "tool-use",
    "functionCall",
    "function_call",
])

OPENAI_FUNCTION_CALL_TYPES: frozenset[str] = frozenset([
    "functionCall",
    "function_call",
])


# -- Extraction helpers --

def _extract_tool_call_id(block: dict[str, Any]) -> str | None:
    """Extract tool call id from a content block (id or call_id)."""
    bid = block.get("id")
    if isinstance(bid, str) and bid:
        return bid
    call_id = block.get("call_id")
    if isinstance(call_id, str) and call_id:
        return call_id
    return None


def _normalize_assistant_reasoning_blocks(message: dict[str, Any]) -> dict[str, Any]:
    """
    Reorder reasoning/thinking blocks to appear before tool calls.

    Only repairs the specific OpenAI shape: a single function call that has
    one or more reasoning blocks after it. Multi-call turns are left untouched.
    """
    content = message.get("content")
    if not isinstance(content, list):
        return message

    saw_tool_call = False
    reasoning_after_tool_call = False
    function_call_count = 0

    for block in content:
        if not isinstance(block, dict):
            return message

        btype = block.get("type")

        if btype in ("reasoning", "thinking"):
            if saw_tool_call:
                reasoning_after_tool_call = True
            continue

        if isinstance(btype, str) and btype in TOOL_CALL_TYPES:
            saw_tool_call = True
            if btype in OPENAI_FUNCTION_CALL_TYPES:
                function_call_count += 1
            continue

        # Any other block type — return unmodified
        return message

    # Only repair single-function-call turns with trailing reasoning
    if not reasoning_after_tool_call or function_call_count != 1:
        return message

    reasoning = [b for b in content if b.get("type") in ("reasoning", "thinking")]
    tool_calls = [b for b in content if isinstance(b.get("type"), str) and b.get("type") in TOOL_CALL_TYPES]

    return {**message, "content": reasoning + tool_calls}


def _extract_tool_calls_from_assistant(msg: dict[str, Any]) -> list[dict[str, Any]]:
    """Return list of {id, name} dicts for each tool call in an assistant message."""
    content = msg.get("content")
    if not isinstance(content, list):
        return []

    tool_calls: list[dict[str, Any]] = []
    for block in content:
        if not isinstance(block, dict):
            continue
        btype = block.get("type")
        if not isinstance(btype, str) or btype not in TOOL_CALL_TYPES:
            continue
        bid = _extract_tool_call_id(block)
        if not bid:
            continue
        name = block.get("name")
        tool_calls.append({
            "id": bid,
            "name": name if isinstance(name, str) else None,
        })
    return tool_calls


def _extract_tool_result_id(msg: dict[str, Any]) -> str | None:
    """Extract tool result id from a toolResult message (top-level field)."""
    tool_call_id = msg.get("toolCallId")
    if isinstance(tool_call_id, str) and tool_call_id:
        return tool_call_id
    tool_use_id = msg.get("toolUseId")
    if isinstance(tool_use_id, str) and tool_use_id:
        return tool_use_id
    return None


def _make_missing_tool_result(tool_call_id: str, tool_name: str | None = None) -> dict[str, Any]:
    """Create a synthetic error tool result for a missing tool call response."""
    return {
        "role": "toolResult",
        "toolCallId": tool_call_id,
        "toolName": tool_name if tool_name is not None else "unknown",
        "content": [
            {
                "type": "text",
                "text": "[lossless-claw] missing tool result in session history; inserted synthetic error result for transcript repair.",
            }
        ],
        "isError": True,
        "timestamp": int(time.time() * 1000),
    }


# -- Main repair function --

def sanitize_tool_use_result_pairing(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """
    Repair tool use/result pairing in an assembled message transcript.

    Anthropic (and Cloud Code Assist) reject transcripts where assistant tool
    calls are not immediately followed by matching tool results. This function:
    - Moves matching toolResult messages directly after their assistant toolCall turn
    - Inserts synthetic error toolResults for missing IDs
    - Drops duplicate toolResults for the same ID
    - Drops orphaned toolResults with no matching tool call

    Matches the TypeScript logic in lossless-claw/src/transcript-repair.ts exactly.
    """
    out: list[dict[str, Any]] = []
    seen_tool_result_ids: set[str] = set()
    dropped_duplicate_count = 0
    dropped_orphan_count = 0
    moved = False
    changed = False

    def push_tool_result(msg: dict[str, Any]) -> None:
        nonlocal dropped_duplicate_count, changed
        rid = _extract_tool_result_id(msg)
        if rid and rid in seen_tool_result_ids:
            dropped_duplicate_count += 1
            changed = True
            return
        if rid:
            seen_tool_result_ids.add(rid)
        out.append(msg)

    i = 0
    while i < len(messages):
        msg = messages[i]

        if not isinstance(msg, dict):
            out.append(msg)
            i += 1
            continue

        role = msg.get("role")

        if role != "assistant":
            if role != "toolResult":
                out.append(msg)
            else:
                # Orphaned toolResult — no preceding assistant tool call
                dropped_orphan_count += 1
                changed = True
            i += 1
            continue

        # Normalize reasoning block order for OpenAI shape
        normalized = _normalize_assistant_reasoning_blocks(msg)
        if normalized is not msg:
            changed = True

        # Skip tool call extraction for aborted/errored assistant messages
        stop_reason = normalized.get("stopReason")
        if stop_reason in ("error", "aborted"):
            out.append(normalized)
            i += 1
            continue

        tool_calls = _extract_tool_calls_from_assistant(normalized)
        if not tool_calls:
            out.append(normalized)
            i += 1
            continue

        tool_call_ids: set[str] = {tc["id"] for tc in tool_calls}

        span_results_by_id: dict[str, dict[str, Any]] = {}
        remainder: list[dict[str, Any]] = []

        j = i + 1
        while j < len(messages):
            nxt = messages[j]

            if not isinstance(nxt, dict):
                remainder.append(nxt)
                j += 1
                continue

            next_role = nxt.get("role")

            if next_role == "assistant":
                break

            if next_role == "toolResult":
                rid = _extract_tool_result_id(nxt)
                if rid and rid in tool_call_ids:
                    if rid in seen_tool_result_ids:
                        dropped_duplicate_count += 1
                        changed = True
                        j += 1
                        continue
                    if rid not in span_results_by_id:
                        span_results_by_id[rid] = nxt
                    else:
                        # Within-span duplicate: silent drop, but mark as changed
                        dropped_duplicate_count += 1
                        changed = True
                    j += 1
                    continue
                # toolResult not matching this turn → orphan
                dropped_orphan_count += 1
                changed = True
                j += 1
                continue

            # Non-toolResult, non-assistant messages
            remainder.append(nxt)
            j += 1

        out.append(normalized)

        if span_results_by_id and remainder:
            moved = True
            changed = True

        for call in tool_calls:
            existing = span_results_by_id.get(call["id"])
            if existing is not None:
                push_tool_result(existing)
            else:
                missing = _make_missing_tool_result(call["id"], call.get("name"))
                changed = True
                push_tool_result(missing)

        for rem in remainder:
            out.append(rem)

        i = j - 1
        i += 1

    changed_or_moved = changed or moved
    return out if changed_or_moved else messages
