"""Transcript Policy: Provider-specific message sanitization before persistence.

Ensures tool call/result pairing integrity and provider format compliance.
"""

from __future__ import annotations

import re
from typing import Any

from loguru import logger


# ── Synthetic result for orphaned tool calls ──────────────────────────

_SYNTHETIC_RESULT = "[no result received]"


def sanitize_transcript(
    messages: list[dict[str, Any]],
    provider_type: str | None = None,
) -> list[dict[str, Any]]:
    """Sanitize a message transcript for persistence or re-submission.

    1. Fix tool call/result pairing (universal)
    2. Apply provider-specific rules
    3. Drop malformed messages
    """
    result = _fix_tool_pairing(messages)
    result = _drop_malformed(result)

    if provider_type:
        p = provider_type.lower()
        if p in ("gemini", "google"):
            result = _sanitize_gemini(result)
        elif p == "anthropic":
            result = _sanitize_anthropic(result)
        elif p == "mistral":
            result = _sanitize_mistral(result)

    return result


# ── Universal: tool call/result pairing ──────────────────────────────

def _fix_tool_pairing(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Ensure every tool_call has a matching tool result.

    - Collect all tool_call_ids from assistant messages
    - Collect all tool_call_ids from tool result messages
    - For any orphaned tool_call (no matching result), append a synthetic result
    """
    # Collect declared tool calls and existing results
    declared_ids: list[tuple[str, str]] = []  # (tool_call_id, tool_name)
    result_ids: set[str] = set()

    for m in messages:
        if m.get("role") == "assistant":
            for tc in m.get("tool_calls", []):
                tc_id = tc.get("id") or tc.get("tool_call_id")
                tc_name = tc.get("function", {}).get("name", "") or tc.get("name", "")
                if tc_id:
                    declared_ids.append((tc_id, tc_name))
        elif m.get("role") == "tool":
            tc_id = m.get("tool_call_id")
            if tc_id:
                result_ids.add(tc_id)

    # Find orphaned calls
    orphaned = [(tc_id, name) for tc_id, name in declared_ids if tc_id not in result_ids]
    if not orphaned:
        return messages

    logger.debug("Fixing {} orphaned tool calls", len(orphaned))
    result = list(messages)
    for tc_id, name in orphaned:
        result.append({
            "role": "tool",
            "tool_call_id": tc_id,
            "name": name,
            "content": _SYNTHETIC_RESULT,
        })

    return result


def _drop_malformed(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Drop tool_calls that lack both 'input' and 'arguments'."""
    result = []
    for m in messages:
        if m.get("role") == "assistant" and m.get("tool_calls"):
            clean_calls = []
            for tc in m["tool_calls"]:
                fn = tc.get("function", {})
                has_args = fn.get("arguments") is not None or tc.get("input") is not None
                if has_args or fn.get("name"):
                    clean_calls.append(tc)
                else:
                    logger.debug("Dropping malformed tool_call: {}", tc)
            m = {**m, "tool_calls": clean_calls}
            if not clean_calls and not m.get("content"):
                continue
        result.append(m)
    return result


# ── Gemini-specific ──────────────────────────────────────────────────

_GEMINI_ID_RE = re.compile(r"[^a-zA-Z0-9]")


def _sanitize_gemini(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Gemini requires:
    - tool_call_ids as alphanumeric only
    - strict user/assistant alternation
    """
    # Clean tool_call_ids
    id_map: dict[str, str] = {}
    result = []

    for m in messages:
        m = dict(m)
        if m.get("role") == "assistant" and m.get("tool_calls"):
            new_calls = []
            for tc in m["tool_calls"]:
                tc = dict(tc)
                old_id = tc.get("id") or tc.get("tool_call_id")
                if old_id:
                    clean_id = _GEMINI_ID_RE.sub("", old_id)
                    if not clean_id:
                        clean_id = "call" + str(len(id_map))
                    id_map[old_id] = clean_id
                    if "id" in tc:
                        tc["id"] = clean_id
                    if "tool_call_id" in tc:
                        tc["tool_call_id"] = clean_id
                new_calls.append(tc)
            m["tool_calls"] = new_calls

        if m.get("role") == "tool":
            old_id = m.get("tool_call_id")
            if old_id and old_id in id_map:
                m["tool_call_id"] = id_map[old_id]

        result.append(m)

    # Enforce strict alternation (merge consecutive same-role messages)
    return _enforce_alternation(result)


def _enforce_alternation(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Merge consecutive messages of the same role (user or assistant)."""
    if not messages:
        return messages

    result: list[dict[str, Any]] = [messages[0]]
    for m in messages[1:]:
        role = m.get("role")
        prev_role = result[-1].get("role")

        # tool messages are part of the assistant turn — don't merge
        if role in ("tool", "system") or prev_role in ("tool", "system"):
            result.append(m)
            continue

        if role == prev_role:
            # Merge content
            prev_content = result[-1].get("content", "")
            new_content = m.get("content", "")
            if isinstance(prev_content, str) and isinstance(new_content, str):
                result[-1] = {**result[-1], "content": f"{prev_content}\n\n{new_content}"}
            else:
                result.append(m)
        else:
            result.append(m)

    return result


# ── Anthropic-specific ───────────────────────────────────────────────

def _sanitize_anthropic(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Anthropic requires:
    - No consecutive user messages (merge them)
    - Tool results paired with tool calls
    """
    return _enforce_alternation(messages)


# ── Mistral-specific ─────────────────────────────────────────────────

def _sanitize_mistral(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Mistral requires tool_call_ids to be max 9 characters."""
    id_map: dict[str, str] = {}
    counter = 0
    result = []

    for m in messages:
        m = dict(m)
        if m.get("role") == "assistant" and m.get("tool_calls"):
            new_calls = []
            for tc in m["tool_calls"]:
                tc = dict(tc)
                for key in ("id", "tool_call_id"):
                    old_id = tc.get(key)
                    if old_id and len(old_id) > 9:
                        if old_id not in id_map:
                            id_map[old_id] = f"tc{counter:06d}"
                            counter += 1
                        tc[key] = id_map[old_id]
                new_calls.append(tc)
            m["tool_calls"] = new_calls

        if m.get("role") == "tool":
            old_id = m.get("tool_call_id")
            if old_id and old_id in id_map:
                m["tool_call_id"] = id_map[old_id]

        result.append(m)

    return result
