"""Relevance-based context pruning for the provider-bound conversation.

Ported from the MIT-licensed alexrudloff/llmrouter "pinch" module, reduced to
the pure, injectable decision core that 6krrt can test offline. Where llmrouter
embeds every candidate message and scores cosine relevance, this module keeps
the same SAFE invariants without requiring an embedding model on the request
path:

  - user / assistant / system messages are ALWAYS kept verbatim,
  - only TOOL RESULTS older than the protected window are trimmed or
    summarized (never removed),
  - message order and the tool/assistant pairing are preserved so the result
    still parses as a valid conversation.

The tailoring is deliberate. Tool results are where a long agent session's
tokens actually live, they are the least likely to be needed in full by the
time a later turn is answered, and replacing one with a short placeholder is
reversible at the semantic level -- a wrong guess costs context, but it never
breaks the request. Trimming a user or assistant message, by contrast, can
change what the model is being asked, so those are never touched.

This module is pure: it takes messages and limits and returns pruned messages.
`dispatcher.py` owns reading the config and deciding when to call it.
"""

from __future__ import annotations

from config import PinchConfig

# Mirrors dispatcher.CHARS_PER_TOKEN.
CHARS_PER_TOKEN = 3


def estimate_tokens(text: str | None) -> int:
    """Crude characters-per-token estimate, consistent with the dispatcher."""
    return len(text) // CHARS_PER_TOKEN if text else 0


def extract_text(message: dict) -> str:
    """Best-effort text of a message, whether content is str or content blocks.

    ``image_url`` blocks contribute their ``url`` so image-bearing messages are
    not silently undercounted as empty. A ``tool_result`` whose ``content`` is
    itself a list of blocks is flattened recursively the same way.
    """
    content = message.get("content")
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        out: list[str] = []
        for part in content:
            if not isinstance(part, dict):
                continue
            ptype = part.get("type")
            if ptype == "text":
                text = part.get("text")
                if isinstance(text, str):
                    out.append(text)
            elif ptype == "tool_result":
                inner = part.get("content")
                if isinstance(inner, list):
                    out.append(extract_text({"content": inner}))
                elif isinstance(inner, str):
                    out.append(inner)
            elif ptype == "image_url":
                image_url = part.get("image_url") or {}
                url = image_url.get("url") if isinstance(image_url, dict) else None
                if isinstance(url, str):
                    out.append(url)
        return " ".join(out)
    return str(content) if content else ""


def _text_only(message: dict) -> str:
    """The prose text of a message, excluding image payloads.

    Unlike :func:`extract_text` -- which includes ``image_url`` urls so they
    count toward the size estimate that triggers pruning -- this returns only
    the text blocks' content. A trimmed replacement must be built from prose,
    never from raw base64, or an arbitrary byte slice of an image ends up
    sitting in a ``text`` field. ``extract_text`` still sizes the message and
    computes savings; this only shapes the replacement.
    """
    content = message.get("content")
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        out: list[str] = []
        for part in content:
            if not isinstance(part, dict):
                continue
            ptype = part.get("type")
            if ptype == "text":
                text = part.get("text")
                if isinstance(text, str):
                    out.append(text)
            elif ptype == "tool_result":
                inner = part.get("content")
                if isinstance(inner, list):
                    out.append(_text_only({"content": inner}))
                elif isinstance(inner, str):
                    out.append(inner)
        return " ".join(out)
    return str(content) if content else ""


def _tool_name(message: dict) -> str | None:
    """A tool message's name, under either the OpenAI or Anthropic spelling."""
    name = message.get("name") or message.get("tool_name") or message.get("function")
    if isinstance(name, dict):
        name = name.get("name")
    return name if isinstance(name, str) else None


def _first_user_turn_indexes(messages: list[dict]) -> list[int]:
    """Indexes of messages that start a new user turn.

    A user message starts a turn; the assistant reply and any tool results that
    follow belong to that turn until the next user message.
    """
    return [i for i, m in enumerate(messages) if m.get("role") == "user"]


def _with_text(message: dict, new_text: str) -> dict:
    """Return ``message`` with its tool-result content trimmed to ``new_text``.

    When the content is a list of blocks, the text blocks become ``new_text``
    and any ``image_url`` blocks have their payload shrunk to a stub. Trimming
    must reduce the *actual* bytes shipped upstream, and an image is usually
    the dominant contributor, so leaving it full-size while claiming a token
    saving would be a lie. The ``image_url`` block type is preserved (so the
    model still knows an image was present) but the multi-MB base64 url is
    replaced by a short marker. When there is no text block, one carrying
    ``new_text`` is appended; a message with only an image collapses to just
    the stub.
    """
    content = message.get("content")
    if isinstance(content, list):
        blocks: list[dict] = []
        replaced_text = False
        for part in content:
            if not isinstance(part, dict):
                blocks.append(part)
                continue
            ptype = part.get("type")
            if ptype == "text":
                if not replaced_text:
                    blocks.append({**part, "text": new_text})
                    replaced_text = True
                # Any subsequent text block is folded into ``new_text``.
            elif ptype == "image_url":
                # Shrink the image payload to a stub; keep the block type.
                blocks.append({**part, "image_url": {"url": "[image omitted]"}})
            else:
                blocks.append(part)
        if not replaced_text:
            blocks.append({"type": "text", "text": new_text})
        if blocks:
            return {**message, "content": blocks}
    return {**message, "content": new_text}


def prune_context(
    messages: list[dict],
    budget_tokens: int = PinchConfig.model_fields["budget_tokens"].default,
    keep_last_turns: int = PinchConfig.model_fields["keep_last_turns"].default,
    max_summarize_chars: int = PinchConfig.model_fields["max_summarize_chars"].default,
) -> tuple[list[dict], dict]:
    """Trim old tool results once a conversation exceeds ``budget_tokens``.

    Returns ``(pruned_messages, stats)``. Keeps every user/assistant/system
    message verbatim and never removes a message: tool results are replaced
    *in place* with a short summary or placeholder, so ``len(pruned)`` always
    equals ``len(messages)``. Order and role pairing are preserved, so the
    result is a valid conversation with the same shape.

    Only runs (and only mutates anything) when the estimate actually exceeds
    the budget; otherwise the original list is returned untouched.
    """
    orig_tokens = sum(estimate_tokens(extract_text(m)) for m in messages)
    if orig_tokens <= budget_tokens:
        return messages, {
            "pruned": False,
            "original_tokens": orig_tokens,
            "final_tokens": orig_tokens,
            "tokens_saved": 0,
        }

    user_indexes = _first_user_turn_indexes(messages)
    if not user_indexes:
        # No user turn at all: treat the conversation as a single ongoing turn
        # and protect the trailing ``keep_last_turns`` tool results (the ones
        # the next turn needs). ``keep_last_turns`` bounds how many are
        # protected; a tiny conversation protects all of them.
        tool_indexes = [i for i, m in enumerate(messages) if m.get("role") == "tool"]
        protected_from = (
            tool_indexes[max(len(tool_indexes) - keep_last_turns, 0)]
            if tool_indexes
            else len(messages)
        )
    else:
        num_protected_turns = min(keep_last_turns, len(user_indexes))
        if num_protected_turns <= 0:
            # Nothing to protect: everything before is a candidate for trimming.
            protected_from = len(messages)
        else:
            protected_from = user_indexes[len(user_indexes) - num_protected_turns]

    pruned: list[dict] = []
    summarized = 0

    for i, msg in enumerate(messages):
        role = msg.get("role")
        if role in ("user", "assistant", "system") or i >= protected_from:
            pruned.append(msg)
            continue
        # Only tool results are candidates here.
        if role != "tool":
            pruned.append(msg)
            continue
        content = msg.get("content")
        if isinstance(content, str):
            text = content
            prose = content
        else:
            # `text` sizes the message (includes image bytes, so an image-heavy
            # result still triggers pruning and still affects the length guard);
            # `prose` shapes the replacement (never raw base64).
            text = extract_text(msg)
            prose = _text_only(msg)
        head_len = 1500
        tail_len = 1500
        replaced = False
        if len(text) > max_summarize_chars:
            # Long tool result: keep head + tail so the shape survives, but
            # only when the elision actually saves characters. A result that is
            # not meaningfully longer than head+tail+marker would only grow (or
            # produce a negative trim), so fall through to the placeholder path.
            head = prose[:head_len]
            tail = prose[-tail_len:]
            trimmed = len(prose) - head_len - tail_len
            marker = f"\n\n[{trimmed:,} chars trimmed...]\n\n" if trimmed > 0 else ""
            elided = f"{head}{marker}{tail}"
            if trimmed > 0 and len(elided) < len(prose):
                summarized += 1
                pruned.append(_with_text(msg, elided))
                replaced = True
        if not replaced:
            # Short result (or a long one that cannot be elided to save space):
            # replace with a short placeholder, but only when it is strictly
            # shorter than the *combined* size — an image-only result may have
            # empty prose yet large real size, and shrinking it still saves
            # tokens (the image payload is stubbed by _with_text).
            name = _tool_name(msg) or "tool"
            placeholder = f"[{name}: result omitted]"
            if len(placeholder) < len(text):
                summarized += 1
                pruned.append(_with_text(msg, placeholder))
            else:
                pruned.append(msg)

    final_tokens = sum(estimate_tokens(extract_text(m)) for m in pruned)
    return pruned, {
        "pruned": True,
        "original_tokens": orig_tokens,
        "final_tokens": final_tokens,
        # Savings computed from what the payload actually shrunk by (the whole
        # list before vs after), so images + prose that really got trimmed are
        # the only thing counted — never a per-message estimate that could
        # drift from reality as #5/#11 interacted.
        "tokens_saved": max(orig_tokens - final_tokens, 0),
        "summarized": summarized,
    }