"""Bound a tool result's serialized size so the MCP client never SILENTLY
truncates it to an opaque scratchpad file — which is exactly what makes an
agent lose its place.

The client caps a single tool result. When the server returns more, the client
abridges to a file on the user's disk and the agent must guess where to resume.

Design — accuracy and not-getting-lost over byte count, and NEVER make it worse:

  1. Only RECORD-BACKED strings are abridged — a value is stubbed only when its
     immediate container carries an A360 identity (id/executionId/fileId/...),
     so every stub can hand back a precise, correct re-fetch instruction.
     Computed values (a diagnosis verdict, a rendered diff) have no identity and
     are LEFT WHOLE: the client's scratchpad copy of them is recoverable, a stub
     of them would not be.
  2. Identity / safety / navigation keys are never abridged at any depth.
  3. The budget is measured in UTF-8 BYTES (CJK text is ~3 bytes/char), because
     that is what the client ceiling actually counts.
  4. Largest-first and minimal: stub as few fields as needed to fit; if records
     overflow by COUNT rather than size, drop trailing rows with a paging hint.
  5. Honest: if even that cannot fit (e.g. one huge protected field), return
     best-effort and say so — never claim a fit that did not happen, and always
     report how many items were dropped.

Pure functions, no network, no mutation — fully unit-testable.
"""

from __future__ import annotations

import hashlib
import os
from typing import Any

import orjson

# Client single-result ceiling proxy, in UTF-8 BYTES. We stay conservatively
# under the byte-equivalent of the client token ceiling so the JSON envelope
# plus our stubs never tip the client into its own truncation.
ENV_BUDGET = "A360_RESPONSE_BUDGET_BYTES"
DEFAULT_BUDGET_BYTES = 75_000

# Only stub genuinely large bodies. Small per-row fields are handled by row
# truncation instead, so they never get a weak, identity-less re-fetch hint.
MIN_STUB_FIELD_BYTES = 2_000

# Characters of the original value kept inline so the agent can still tell what
# the field holds without a follow-up call.
PREVIEW_CHARS = 240

# A360 record identity keys. A container carrying any of these (with a non-empty
# string value) is record-backed: its large fields can be re-fetched, so they
# are stub-eligible. Order matters only for "first match wins" in the hint.
IDENTITY_KEYS = ("id", "executionId", "fileId", "deploymentId", "activityId")

# Keys whose values are identity / safety / navigation / computed-output signals.
# Never abridged, even when long, at any depth: losing them is how an agent gets
# lost, misjudges a write, or loses a non-recoverable diagnosis result, and
# computed values cannot be re-fetched by identity.
PROTECTED_KEYS = frozenset(
    {
        # identity / navigation
        "id",
        "executionId",
        "fileId",
        "deploymentId",
        "activityId",
        "name",
        "status",
        "success",
        "operation",
        "target",
        "scope",
        "bot_path",
        "snapshot_path",
        "path",
        # safety / risk signals
        "error",
        "message",
        "hint",
        "warning",
        "warnings",
        "note",
        "notes",
        "reason",
        "failure",
        "risk",
        "level",
        "factors",
        # computed, non-record-backed diagnosis outputs (cannot be re-fetched)
        "verdict",
        "diagnosis",
        "prescription",
        "diff",
        "diffs",
        "payload",
        "summary",
        "trajectory",
        "echo",
        # our own stub / marker internals
        "_fetch",
        "_sha256",
        "_full_length_bytes",
        "_abridged",
        "_abridged_fields",
        "_abridged_note",
        "dropped_items",
    }
)

# Bytes reserved for the top-level abridged marker appended after row truncation.
_MARKER_RESERVE = 400

_ABRIDGED_NOTE = (
    "Large values were abridged to fit the response budget; each stub carries "
    "_full_length_bytes, _sha256, _preview, and a _fetch hint. This is NOT the complete "
    "content — fetch any field you need in full via its _fetch instruction."
)


# --------------------------------------------------------------------------- #
# Measurement
# --------------------------------------------------------------------------- #


def get_response_budget() -> int:
    """Soft UTF-8 byte budget for a single serialized tool result."""
    raw = os.getenv(ENV_BUDGET)
    if raw:
        try:
            value = int(raw)
            if value > 0:
                return value
        except ValueError:
            pass
    return DEFAULT_BUDGET_BYTES


def byte_len(obj: Any) -> int:
    """UTF-8 byte length of *obj* serialized as compact JSON (orjson)."""
    # orjson.dumps returns raw UTF-8 bytes with no ASCII escaping and compact
    # separators — the exact wire form the client measures.
    return len(orjson.dumps(obj))


def _str_bytes(value: str) -> int:
    return len(value.encode("utf-8"))


def _sha256(value: str) -> str:
    return hashlib.sha256(value.encode("utf-8")).hexdigest()


# --------------------------------------------------------------------------- #
# Stubs and fetch hints
# --------------------------------------------------------------------------- #


def _container_identity(obj: dict[str, Any]) -> tuple[str, str]:
    """Return (identity_kind, identity_value) for the first A360 identity key
    present with a non-empty string value, or ("", "") when the container is
    not record-backed.

    Only string identities count: numeric ids are not stable re-fetch handles
    we can safely embed in a copy-pasteable hint.
    """
    for key in IDENTITY_KEYS:
        sid = obj.get(key)
        if isinstance(sid, str) and sid.strip():
            return key, sid
    return "", ""


def _build_fetch_hint(key: str | None, kind: str, value: str, tool_name: str) -> str:
    """A correct, copy-pasteable instruction for fetching the abridged field.

    Only ever called for record-backed values (identity present), so the hint is
    always actionable. Falls back to a tool-name reference when no identity is
    available (defensive — eligibility already gates on identity).
    """
    if kind and value:
        return (
            f"Field '{key}' abridged to fit the response budget. Re-run the producing "
            f"tool scoped to this record ({kind}='{value}') to get '{key}' in full."
        )
    return (
        f"Field '{key}' abridged to fit the response budget. Re-run {tool_name} to get "
        f"'{key}' in full."
    )


def _stub(
    key: str | None, value: str, kind: str, identity_value: str, tool_name: str
) -> dict[str, Any]:
    return {
        "_abridged": True,
        "_full_length_bytes": _str_bytes(value),
        "_sha256": _sha256(value),
        "_preview": value[:PREVIEW_CHARS],
        "_fetch": _build_fetch_hint(key, kind, identity_value, tool_name),
    }


# --------------------------------------------------------------------------- #
# Field stubbing (record-backed large strings)
# --------------------------------------------------------------------------- #


def _collect_eligible(obj: Any) -> list[int]:
    """Byte lengths of every stub-eligible string leaf.

    Eligible = a string whose immediate container dict has an A360 identity,
    whose key is not protected, and whose UTF-8 size is at least
    MIN_STUB_FIELD_BYTES. Protected-key subtrees are skipped entirely.
    """
    out: list[int] = []
    if isinstance(obj, dict):
        kind, identity_value = _container_identity(obj)
        for key, value in obj.items():
            if key in PROTECTED_KEYS:
                continue
            if isinstance(value, (dict, list)):
                out.extend(_collect_eligible(value))
            elif (
                isinstance(value, str)
                and kind
                and identity_value
                and _str_bytes(value) >= MIN_STUB_FIELD_BYTES
            ):
                out.append(_str_bytes(value))
    elif isinstance(obj, list):
        for item in obj:
            if isinstance(item, (dict, list)):
                out.extend(_collect_eligible(item))
    return out


def _abridge_strings(
    obj: Any, *, tool_name: str, threshold: int, stubbed: list[str], path: str = ""
) -> Any:
    """Return a NEW structure with record-backed string leaves >= *threshold*
    bytes replaced by stubs. Protected-key subtrees are copied whole."""
    if isinstance(obj, dict):
        kind, identity_value = _container_identity(obj)
        out: dict[str, Any] = {}
        for key, value in obj.items():
            child = f"{path}.{key}" if path else key
            if key in PROTECTED_KEYS:
                out[key] = value
            elif isinstance(value, (dict, list)):
                out[key] = _abridge_strings(
                    value,
                    tool_name=tool_name,
                    threshold=threshold,
                    stubbed=stubbed,
                    path=child,
                )
            elif (
                isinstance(value, str)
                and kind
                and identity_value
                and _str_bytes(value) >= threshold
            ):
                out[key] = _stub(key, value, kind, identity_value, tool_name)
                stubbed.append(child)
            else:
                out[key] = value
        return out
    if isinstance(obj, list):
        return [
            _abridge_strings(
                item,
                tool_name=tool_name,
                threshold=threshold,
                stubbed=stubbed,
                path=f"{path}[{i}]",
            )
            for i, item in enumerate(obj)
        ]
    return obj


def _with_marker(obj: Any, stubbed: list[str]) -> Any:
    """Attach the top-level abridged marker (dict results only) for measurement."""
    if isinstance(obj, dict) and stubbed:
        return {**obj, "_abridged_fields": stubbed, "_abridged_note": _ABRIDGED_NOTE}
    return obj


def _fit_by_stubbing(
    result: Any, *, tool_name: str, eligible: list[int], budget: int
) -> tuple[Any, list[str]]:
    """Stub the fewest (largest-first) record-backed fields needed to fit budget.

    Binary-searches the byte threshold; if even stubbing every eligible field
    cannot fit, returns the all-stubbed best effort.
    """
    ordered = sorted(eligible, reverse=True)
    found: tuple[Any, list[str]] | None = None
    lo, hi = 1, len(ordered)
    while lo <= hi:
        mid = (lo + hi) // 2
        threshold = ordered[mid - 1]
        stubbed: list[str] = []
        cand = _abridge_strings(result, tool_name=tool_name, threshold=threshold, stubbed=stubbed)
        if byte_len(_with_marker(cand, stubbed)) <= budget:
            found = (cand, stubbed)
            hi = mid - 1
        else:
            lo = mid + 1
    if found is not None:
        return found
    stubbed = []
    cand = _abridge_strings(result, tool_name=tool_name, threshold=ordered[-1], stubbed=stubbed)
    return cand, stubbed


# --------------------------------------------------------------------------- #
# Row truncation (lists that overflow by element COUNT)
# --------------------------------------------------------------------------- #


def _is_record_list(value: Any) -> bool:
    """True only when every element is a record-backed dict (carries identity).

    Row truncation drops trailing elements, so it is restricted to such lists:
    a dropped record row gets an actionable re-fetch hint. A computed list (no
    identity — rendered lines, diagnosis steps, audit findings) is LEFT WHOLE
    instead, so the client's scratchpad copy stays recoverable. Mirrors the
    stubbing identity gate; without it row truncation would be strictly worse
    than client truncation.
    """
    return (
        isinstance(value, list)
        and len(value) > 1
        and all(isinstance(item, dict) and _container_identity(item)[0] for item in value)
    )


def _max_candidate(
    a: tuple[tuple[Any, ...], int] | None, b: tuple[tuple[Any, ...], int] | None
) -> tuple[tuple[Any, ...], int] | None:
    if a is None:
        return b
    if b is None:
        return a
    return a if a[1] >= b[1] else b


def _largest_list(obj: Any, path: tuple[Any, ...] = ()) -> tuple[tuple[Any, ...], int] | None:
    """Locate the (path, byte_len) of the largest truncatable record list, or None."""
    best: tuple[tuple[Any, ...], int] | None = None
    if isinstance(obj, dict):
        for key, value in obj.items():
            if key in PROTECTED_KEYS:
                continue
            best = _max_candidate(best, _largest_list(value, path + (key,)))
            if _is_record_list(value):
                best = _max_candidate(best, (path + (key,), byte_len(value)))
    elif isinstance(obj, list):
        for index, item in enumerate(obj):
            best = _max_candidate(best, _largest_list(item, path + (index,)))
            if _is_record_list(item):
                best = _max_candidate(best, (path + (index,), byte_len(item)))
    return best


def _get_at(obj: Any, path: tuple[Any, ...]) -> Any:
    for step in path:
        obj = obj[step]
    return obj


def _set_at(obj: Any, path: tuple[Any, ...], value: Any) -> Any:
    """Immutably return a copy of *obj* with the node at *path* replaced."""
    if not path:
        return value
    step, rest = path[0], path[1:]
    if isinstance(obj, dict):
        return {**obj, step: _set_at(obj[step], rest, value)}
    if isinstance(obj, list):
        return [_set_at(item, rest, value) if i == step else item for i, item in enumerate(obj)]
    return value


def _row_marker(dropped: int) -> dict[str, Any]:
    return {
        "dropped_items": dropped,
        "_fetch": (
            f"{dropped} more items omitted to fit the response budget. Narrow the query "
            "or re-run the producing tool with a smaller scope to retrieve them."
        ),
    }


def _fit_by_row_truncation(bounded: Any, *, budget: int, reserve: int) -> tuple[Any, int]:
    """Drop trailing elements of the largest list until the result fits.

    Returns (new_result, dropped_count). dropped_count is 0 when no list could
    help (best effort).
    """
    located = _largest_list(bounded)
    if located is None:
        return bounded, 0
    path, _ = located
    rows = _get_at(bounded, path)
    if not isinstance(rows, list) or len(rows) <= 1:
        return bounded, 0

    # Largest keep-count whose truncated result still fits, reserving room for
    # the top-level abridged marker that enforce_response_budget appends after.
    target = max(0, budget - reserve)
    lo, hi, best_keep = 0, len(rows) - 1, 0
    while lo <= hi:
        mid = (lo + hi) // 2
        dropped = len(rows) - mid
        kept = rows[:mid] + [_row_marker(dropped)]
        if byte_len(_set_at(bounded, path, kept)) <= target:
            best_keep = mid
            lo = mid + 1
        else:
            hi = mid - 1

    dropped = len(rows) - best_keep
    kept = rows[:best_keep] + [_row_marker(dropped)]
    return _set_at(bounded, path, kept), dropped


# --------------------------------------------------------------------------- #
# Entry point
# --------------------------------------------------------------------------- #


def enforce_response_budget(
    result: Any,
    *,
    tool_name: str,
    budget: int | None = None,
) -> tuple[Any, bool]:
    """Abridge oversized values so the serialized result fits *budget* bytes.

    Returns (possibly-new result, was_abridged). Small results are returned
    unchanged. Never mutates the input and never destroys non-recoverable data.
    """
    if not isinstance(result, (dict, list)):
        return result, False

    budget = budget if budget is not None else get_response_budget()
    if byte_len(result) <= budget:
        return result, False

    stubbed: list[str] = []
    bounded: Any = result

    eligible = _collect_eligible(result)
    if eligible:
        bounded, stubbed = _fit_by_stubbing(
            result, tool_name=tool_name, eligible=eligible, budget=budget
        )

    dropped = 0
    if byte_len(_with_marker(bounded, stubbed)) > budget:
        # Reserve room for the marker enforce appends, whose _abridged_fields list
        # grows with the stub count — a fixed reserve would undercount it.
        reserve = _MARKER_RESERVE + (byte_len(stubbed) if stubbed else 0)
        bounded, dropped = _fit_by_row_truncation(bounded, budget=budget, reserve=reserve)

    if not stubbed and dropped == 0:
        # Oversized but nothing safely abridgeable (e.g. one huge protected
        # field). Leave it whole — the client's scratchpad copy is recoverable.
        return result, False

    if isinstance(bounded, dict):
        marker: dict[str, Any] = {"_abridged_note": _ABRIDGED_NOTE}
        if stubbed:
            marker["_abridged_fields"] = stubbed
        if dropped:
            marker["dropped_items"] = dropped
        still_over = byte_len({**bounded, **marker}) > budget
        if still_over:
            marker["_abridged_note"] = (
                _ABRIDGED_NOTE + " NOTE: still over budget after best-effort abridging; "
                "the client may store the remainder in an overflow file."
            )
        bounded = {**bounded, **marker}

    return bounded, True
