#!/usr/bin/env python3
"""Check that a bounded context packet leaves room for model output.

This is deliberately provider-neutral.  The caller supplies the model's
context window; the packet's input budget and a fixed output reserve are
checked before an API request is made.
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path
from typing import Any


OUTPUT_TOKEN_RESERVE = {
    "doc": 768,
    "feature": 1024,
    "contract": 1536,
    "milestone": 2048,
}
PROFILE_INPUT_TOKEN_LIMIT = {
    "doc": 2048,
    "feature": 4096,
    "contract": 6144,
    "milestone": 8192,
}
PROFILE_BYTE_LIMIT = {
    "doc": 8192,
    "feature": 16384,
    "contract": 24576,
    "milestone": 32768,
}
PROFILE_EVIDENCE_BYTE_LIMIT = {
    "doc": 2048,
    "feature": 6144,
    "contract": 10240,
    "milestone": 14336,
}
PROFILE_EVIDENCE_FILE_LIMIT = {
    "doc": 8,
    "feature": 16,
    "contract": 24,
    "milestone": 32,
}
DEFAULT_CONTEXT_WINDOW = 16384
DEFAULT_SAFETY_MARGIN = 256
# Without a provider tokenizer, UTF-8 bytes are the only portable hard upper
# bound. Any multi-byte-per-token heuristic can under-count provider tokens
# for punctuation-heavy or non-ASCII input and is unsafe for a request gate.
TOKEN_BYTES_PER_TOKEN = 1
PACKET_FIELDS = {
    "context_id",
    "packet_id",
    "task",
    "header",
    "baseline",
    "scope",
    "contracts",
    "evidence",
    "risks",
    "next",
    "forbidden",
    "budget",
}
FORBIDDEN_CONTEXT_FIELDS = {
    "history",
    "chat",
    "conversation",
    "messages",
    "transcript",
    "previous_packet",
    "prior_context",
}
EXCLUSION_REASON_CODES = {"E1", "E2", "E3", "E4", "E5", "E6"}


class ModelBudgetError(ValueError):
    """Raised when a request cannot fit input, output, and safety budgets."""


def _response_object(value: Any) -> dict[str, Any]:
    """Return a provider response object from a response or SSE envelope.

    The Responses API exposes usage on the response itself, while streaming
    ``response.completed`` events wrap it below a ``response`` key.  Keeping
    this normalization here prevents callers from accidentally validating the
    event envelope instead of the actual provider usage.
    """

    if not isinstance(value, dict):
        raise ModelBudgetError("provider response must be an object")
    response = value.get("response")
    if isinstance(response, dict):
        return response
    return value


def extract_response_usage(value: Any) -> dict[str, int]:
    """Extract authoritative token usage from a provider response envelope."""

    response = _response_object(value)
    usage = response.get("usage")
    if not isinstance(usage, dict):
        raise ModelBudgetError("provider response is missing usage")
    try:
        input_tokens = int(usage["input_tokens"])
        output_tokens = int(usage["output_tokens"])
        total_tokens = int(usage.get("total_tokens", input_tokens + output_tokens))
    except (KeyError, TypeError, ValueError) as exc:
        raise ModelBudgetError("provider usage has invalid token counts") from exc
    if min(input_tokens, output_tokens, total_tokens) < 0:
        raise ModelBudgetError("provider usage token counts must be non-negative")
    if total_tokens < input_tokens + output_tokens:
        raise ModelBudgetError("provider usage total is smaller than input plus output")
    return {
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "total_tokens": total_tokens,
    }


def _find_forbidden(value: Any, path: str = "") -> str | None:
    """Find accumulated conversation state, including when nested."""

    if isinstance(value, dict):
        for key, child in value.items():
            key_text = str(key)
            child_path = f"{path}.{key_text}" if path else key_text
            if key_text in FORBIDDEN_CONTEXT_FIELDS:
                return child_path
            found = _find_forbidden(child, child_path)
            if found:
                return found
    elif isinstance(value, list):
        for index, child in enumerate(value):
            found = _find_forbidden(child, f"{path}[{index}]")
            if found:
                return found
    return None


def _packet_input_bytes(packet: dict[str, Any]) -> int:
    budget = packet.get("budget")
    if not isinstance(budget, dict):
        raise ModelBudgetError("packet budget is missing")
    try:
        packet_bytes = int(budget["packet_bytes"])
        evidence_bytes = int(budget["selected_evidence_bytes"])
    except (KeyError, TypeError, ValueError) as exc:
        raise ModelBudgetError("packet budget lacks input byte totals") from exc
    if packet_bytes < 0 or evidence_bytes < 0:
        raise ModelBudgetError("packet input byte totals must be non-negative")
    return packet_bytes + evidence_bytes


def _canonical_packet_bytes(packet: dict[str, Any]) -> int:
    """Return the exact UTF-8 size used by the packet generator."""

    return len(
        json.dumps(packet, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
    )


def _validate_packet_metadata(packet: dict[str, Any], budget: dict[str, Any]) -> None:
    """Reject stale or inconsistent budget metadata before a model request.

    Small unit-test fixtures may contain only a budget object.  Full packets
    generated for a request always carry ``context_id``/``packet_id`` and an
    evidence envelope, so strict consistency checks are applied to those.
    """

    forbidden = _find_forbidden(packet)
    if forbidden:
        raise ModelBudgetError(
            "packet contains forbidden accumulated context "
            f"field: {forbidden}"
        )

    # Generated packets have a closed envelope. Reject ad-hoc top-level
    # fields that could smuggle full logs, tool output or a previous request
    # beside the audited eight segments.
    if "budget" in packet:
        unknown = sorted(set(packet) - PACKET_FIELDS)
        if unknown:
            raise ModelBudgetError(
                "packet contains unknown top-level context field: "
                f"{unknown[0]}"
            )

    if set(packet) != {"budget"}:
        actual = _canonical_packet_bytes(packet)
        declared = int(budget["packet_bytes"])
        if declared != actual:
            raise ModelBudgetError(
                "packet byte metadata does not match serialized packet "
                f"(declared={declared}, actual={actual})"
            )

    evidence = packet.get("evidence")
    if not isinstance(evidence, dict):
        return
    selected = evidence.get("selected", [])
    excluded = evidence.get("excluded", [])
    if not isinstance(selected, list) or not isinstance(excluded, list):
        raise ModelBudgetError("packet evidence must contain selected/excluded lists")
    for item in excluded:
        if not isinstance(item, dict) or item.get("reason_code") not in EXCLUSION_REASON_CODES:
            raise ModelBudgetError("excluded evidence has an unknown reason code")
    selected_bytes = 0
    for item in selected:
        if not isinstance(item, dict):
            raise ModelBudgetError("selected evidence entries must be objects")
        try:
            item_bytes = int(item["bytes"])
        except (KeyError, TypeError, ValueError) as exc:
            raise ModelBudgetError("selected evidence entry lacks bytes") from exc
        if item_bytes < 0:
            raise ModelBudgetError("selected evidence bytes must be non-negative")
        selected_bytes += item_bytes
    if selected_bytes != int(budget["selected_evidence_bytes"]):
        raise ModelBudgetError(
            "selected evidence byte metadata does not match entries "
            f"(declared={budget['selected_evidence_bytes']}, actual={selected_bytes})"
        )
    if "selected_evidence_files" in budget and int(budget["selected_evidence_files"]) != len(selected):
        raise ModelBudgetError("selected evidence file count does not match entries")
    if "excluded_evidence_files" in budget and int(budget["excluded_evidence_files"]) != len(excluded):
        raise ModelBudgetError("excluded evidence file count does not match entries")
    if "evidence_byte_limit" in budget and selected_bytes > int(budget["evidence_byte_limit"]):
        raise ModelBudgetError("selected evidence exceeds its evidence byte budget")
    if "evidence_file_limit" in budget and len(selected) > int(budget["evidence_file_limit"]):
        raise ModelBudgetError("selected evidence exceeds its file budget")


def estimate_input_tokens(packet: dict[str, Any]) -> int:
    """Return a safe provider-neutral upper bound from UTF-8 bytes.

    A tokenizer supplied by the API client may replace this estimate.  The
    byte-for-token bound is deliberately strict: it may reject a packet that
    would fit a particular provider, but it cannot silently approve a packet
    that that provider tokenizes larger than expected.
    """

    return math.ceil(_packet_input_bytes(packet) / TOKEN_BYTES_PER_TOKEN)


def validate_request(
    packet: dict[str, Any],
    *,
    context_window: int = DEFAULT_CONTEXT_WINDOW,
    max_output_tokens: int | None = None,
    safety_margin: int = DEFAULT_SAFETY_MARGIN,
    input_tokens: int | None = None,
    overhead_tokens: int = 0,
) -> dict[str, int]:
    budget = packet.get("budget")
    if not isinstance(budget, dict):
        raise ModelBudgetError("packet budget is missing")
    try:
        input_limit = int(budget["token_limit"])
        byte_limit = int(budget["byte_limit"])
        packet_bytes = int(budget["packet_bytes"])
        selected_evidence_bytes = int(budget["selected_evidence_bytes"])
    except (KeyError, TypeError, ValueError) as exc:
        raise ModelBudgetError("packet budget lacks required fields") from exc
    profile = str(budget.get("profile", ""))
    try:
        profile_output_limit = OUTPUT_TOKEN_RESERVE[profile]
        profile_input_limit = PROFILE_INPUT_TOKEN_LIMIT[profile]
        profile_byte_limit = PROFILE_BYTE_LIMIT[profile]
        profile_evidence_limit = PROFILE_EVIDENCE_BYTE_LIMIT[profile]
        profile_file_limit = PROFILE_EVIDENCE_FILE_LIMIT[profile]
    except KeyError as exc:
        raise ModelBudgetError(f"unknown packet profile: {profile}") from exc
    if max_output_tokens is None:
        max_output_tokens = profile_output_limit
    try:
        context_window = int(context_window)
        max_output_tokens = int(max_output_tokens)
        safety_margin = int(safety_margin)
        overhead_tokens = int(overhead_tokens)
        declared_evidence_limit = int(budget.get("evidence_byte_limit", profile_evidence_limit))
        declared_file_limit = int(budget.get("evidence_file_limit", profile_file_limit))
    except (TypeError, ValueError) as exc:
        raise ModelBudgetError("model budget fields must be integers") from exc
    values = (
        context_window,
        max_output_tokens,
        safety_margin,
        input_limit,
        byte_limit,
        packet_bytes,
        selected_evidence_bytes,
        overhead_tokens,
        declared_evidence_limit,
        declared_file_limit,
    )
    if any(value < 0 for value in values):
        raise ModelBudgetError("model budgets must be non-negative")
    if input_limit > profile_input_limit:
        raise ModelBudgetError(
            "packet token limit exceeds profile "
            f"(declared={input_limit}, profile={profile_input_limit})"
        )
    if byte_limit > profile_byte_limit:
        raise ModelBudgetError(
            "packet byte limit exceeds profile "
            f"(declared={byte_limit}, profile={profile_byte_limit})"
        )
    if declared_evidence_limit > profile_evidence_limit:
        raise ModelBudgetError(
            "evidence byte limit exceeds profile "
            f"(declared={declared_evidence_limit}, profile={profile_evidence_limit})"
        )
    if declared_file_limit > profile_file_limit:
        raise ModelBudgetError(
            "evidence file limit exceeds profile "
            f"(declared={declared_file_limit}, profile={profile_file_limit})"
        )
    if max_output_tokens > profile_output_limit:
        raise ModelBudgetError(
            "output token budget exceeds profile "
            f"(requested={max_output_tokens}, profile={profile_output_limit})"
        )
    try:
        _validate_packet_metadata(packet, budget)
    except ModelBudgetError:
        raise
    except (KeyError, TypeError, ValueError) as exc:
        raise ModelBudgetError("packet budget metadata is malformed") from exc
    context_bytes = packet_bytes + selected_evidence_bytes
    if packet_bytes > byte_limit or context_bytes > byte_limit:
        raise ModelBudgetError("packet exceeds its declared byte budget")
    if selected_evidence_bytes > declared_evidence_limit:
        raise ModelBudgetError("selected evidence exceeds its evidence byte budget")
    if selected_evidence_bytes > profile_evidence_limit:
        raise ModelBudgetError("selected evidence exceeds profile byte budget")
    try:
        estimated_input = estimate_input_tokens(packet) if input_tokens is None else int(input_tokens)
    except (TypeError, ValueError) as exc:
        raise ModelBudgetError("input token count must be an integer") from exc
    if estimated_input < 0:
        raise ModelBudgetError("input token count must be non-negative")
    if estimated_input > input_limit:
        raise ModelBudgetError(
            "packet input token budget exceeded "
            f"(estimated={estimated_input}, limit={input_limit})"
        )
    # Reserve the input actually sent for this request plus fixed provider
    # overhead.  Reserving the full profile here rejects small packets for no
    # reason, while the packet's own input_limit still prevents later growth.
    effective_input = estimated_input + overhead_tokens
    available = context_window - safety_margin - max_output_tokens
    if effective_input > available:
        raise ModelBudgetError(
            "model context budget exceeded "
            f"(input={effective_input}, output={max_output_tokens}, "
            f"safety={safety_margin}, window={context_window})"
        )
    return {
        "context_window": context_window,
        "input_token_limit": input_limit,
        "estimated_input_tokens": estimated_input,
        "overhead_tokens": overhead_tokens,
        "estimated_request_input_tokens": estimated_input + overhead_tokens,
        "max_output_tokens": max_output_tokens,
        "safety_margin": safety_margin,
        "reserved_total_tokens": effective_input + max_output_tokens + safety_margin,
        "available_input_tokens": max(0, available - overhead_tokens),
    }


def validate_response(
    output_tokens: int,
    *,
    max_output_tokens: int,
    finish_reason: str | None = None,
) -> dict[str, int | str | None]:
    """Validate the completion actually returned by a provider.

    The preflight check protects the request, while this postflight check
    catches SDKs that ignore or reinterpret ``max_output_tokens``.  A length
    finish is reported as truncation so callers can persist a checkpoint and
    continue with a fresh bounded request instead of replaying the old one.
    """

    try:
        actual = int(output_tokens)
        limit = int(max_output_tokens)
    except (TypeError, ValueError) as exc:
        raise ModelBudgetError("response token counts must be integers") from exc
    if actual < 0 or limit < 0:
        raise ModelBudgetError("response token counts must be non-negative")
    if actual > limit:
        raise ModelBudgetError(
            "model response exceeded output token budget "
            f"(actual={actual}, limit={limit})"
        )
    return {
        "output_tokens": actual,
        "max_output_tokens": limit,
        "finish_reason": finish_reason,
        "truncated": finish_reason in {
            "length",
            "max_tokens",
            "max_output_tokens",
            "max_completion_tokens",
        },
    }


def validate_provider_response(
    response: Any,
    *,
    max_output_tokens: int,
) -> dict[str, int | str | None | bool]:
    """Validate usage and completion state returned by a provider.

    ``response`` may be a normal Responses API object or a streaming event
    envelope.  Failed/incomplete responses are rejected as non-reusable so a
    caller cannot append partial output to the next request.  A length-limit
    completion is returned as ``truncated`` and must be resumed from a fresh
    checkpoint.
    """

    response_object = _response_object(response)
    usage = extract_response_usage(response_object)
    status = response_object.get("status")
    incomplete = response_object.get("incomplete_details")
    finish_reason = None
    if isinstance(incomplete, dict):
        reason = incomplete.get("reason")
        if reason is not None:
            finish_reason = str(reason)
    if finish_reason is None and response_object.get("finish_reason") is not None:
        finish_reason = str(response_object["finish_reason"])
    if status is None and finish_reason is None:
        raise ModelBudgetError("provider response is missing completion status")
    result = validate_response(
        usage["output_tokens"],
        max_output_tokens=max_output_tokens,
        finish_reason=finish_reason,
    )
    if status in {"failed", "cancelled"}:
        raise ModelBudgetError(
            f"provider response is not reusable (status={status})"
        )
    if status == "incomplete" and not result["truncated"]:
        raise ModelBudgetError(
            "provider response is incomplete without a length checkpoint reason"
        )
    result.update(usage)
    result["status"] = str(status) if status is not None else None
    return result


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("packet", type=Path)
    parser.add_argument("--context-window", type=int, default=DEFAULT_CONTEXT_WINDOW)
    parser.add_argument("--max-output-tokens", type=int)
    parser.add_argument("--safety-margin", type=int, default=DEFAULT_SAFETY_MARGIN)
    parser.add_argument("--input-tokens", type=int, help="exact count from the provider tokenizer")
    parser.add_argument(
        "--overhead-tokens",
        type=int,
        required=True,
        help="system/tool/history tokens outside the context packet",
    )
    args = parser.parse_args()
    try:
        with args.packet.open("r", encoding="utf-8") as handle:
            packet = json.load(handle)
        if not isinstance(packet, dict):
            raise ModelBudgetError("packet must be a JSON object")
        result = validate_request(
            packet,
            context_window=args.context_window,
            max_output_tokens=args.max_output_tokens,
            safety_margin=args.safety_margin,
            input_tokens=args.input_tokens,
            overhead_tokens=args.overhead_tokens,
        )
    except (OSError, UnicodeError, json.JSONDecodeError, ModelBudgetError) as exc:
        parser.error(str(exc))
    print(
        "model_budget=valid "
        f"input_tokens={result['input_token_limit']} "
        f"estimated_input={result['estimated_input_tokens']} "
        f"overhead={result['overhead_tokens']} "
        f"max_output={result['max_output_tokens']} "
        f"reserved_total={result['reserved_total_tokens']}"
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())