#!/usr/bin/env python3
"""Exact token-budget proxy between the dsh harness and the local vLLM server.

The harness estimates context usage from characters (pi-ai: ``chars/4`` plus a
4096-token reserve) and the estimate can under-count real server tokens; the
v2 run died exactly that way (estimated 15,806 input, actual 19,903, requested
12,866 output -> 32,769 > 32,768). This proxy is the single authoritative
meter for v3:

* every request is counted with the server's own tokenizer and chat template
  through ``POST /tokenize`` (same model, same tools, same thinking settings);
* ``max_tokens`` is clamped to ``context - exact_input - reserve`` and to the
  configured per-call cap, with a documented minimum output floor;
* cumulative generated tokens per task are enforced against the budget file,
  and the plugin's continuation decision reads ``GET /budget``;
* one JSONL receipt per request records exact input, granted output, actual
  output and the cumulative counter.

The proxy fails closed: when the exact count cannot be obtained, or the input
no longer leaves the minimum output budget, it returns a structured 4xx instead
of forwarding a request that could silently violate the window.

Run with the eval venv or any Python >= 3.10 (stdlib only).
"""

from __future__ import annotations

import argparse
import gzip
import hashlib
import http.client
import json
import os
import sys
import threading
import time
import urllib.parse
from datetime import datetime, timezone
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any

EVAL_ROOT = Path(__file__).resolve().parents[1]
if str(EVAL_ROOT) not in sys.path:
    sys.path.insert(0, str(EVAL_ROOT))

from evalkit.budget import Budget, BudgetError, load_budget  # noqa: E402

HOP_BY_HOP = {
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailers",
    "transfer-encoding",
    "upgrade",
    "host",
    "content-length",
}


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


# ---- routed (student/teacher) mode -----------------------------------------
# Only explicit quota/balance evidence or rate limiting moves a teacher request
# to the next route. Auth, parameter, server and network failures stay terminal.
QUOTA_CODES = frozenset({"insufficient_quota", "quota_exceeded", "insufficient_balance",
                         "billing_hard_limit_reached", "usage_limit_exceeded", "credit_exhausted"})
RATE_CODES = frozenset({"rate_limit_exceeded", "rate_limit_error", "too_many_requests", "rate_limited"})
FALLBACK_CLASSES = frozenset({"quota", "rate_limit"})


def sha256_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def error_object(data: bytes | None) -> dict:
    candidates = [data or b""]
    candidates.extend(line[5:].strip() for line in (data or b"").splitlines() if line.startswith(b"data:"))
    for candidate in candidates:
        try:
            parsed = json.loads(candidate.decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError):
            continue
        error = parsed.get("error") if isinstance(parsed, dict) else None
        if isinstance(error, dict) and error:
            return error
    return {}


def classify_failure(status: int | None, data: bytes | None) -> str:
    error = error_object(data)
    labels = {str(error.get(key) or "").lower() for key in ("code", "type")}
    if labels & QUOTA_CODES:
        return "quota"
    if labels & RATE_CODES:
        return "rate_limit"
    if status == 402:
        return "quota"
    if status == 429:
        return "rate_limit"
    if status is None:
        return "network"
    if status in (401, 403):
        return "auth"
    if 400 <= status < 500:
        return "bad_request"
    return "server"


def retry_after_seconds(value: str | None, now: float) -> float | None:
    if not value:
        return None
    try:
        return max(0.0, float(value))
    except ValueError:
        pass
    try:
        from email.utils import parsedate_to_datetime
        return max(0.0, parsedate_to_datetime(value).timestamp() - now)
    except (TypeError, ValueError):
        return None


class CredentialUnavailable(ValueError):
    """A selected route cannot read its credential; no secret or private path is reported."""


def resolve_credential(spec: dict | None) -> str | None:
    """Read a key only for the selected route. Never log credential values or file paths."""
    spec = spec or {}
    file_env, env = spec.get("file_env"), spec.get("env")
    try:
        if file_env and os.environ.get(file_env):
            key = Path(os.environ[file_env]).read_text(encoding="utf-8").strip()
        elif env and os.environ.get(env):
            key = os.environ[env].strip()
        else:
            key = None
    except (OSError, UnicodeError):
        raise CredentialUnavailable("selected route credential file is unreadable") from None
    if not key and spec.get("required"):
        raise CredentialUnavailable("selected route credential is missing or empty")
    return key or None


def sse_summary(data: bytes) -> dict:
    """Read terminal identity from captured Chat SSE/JSON or Responses completion.

    Initial/running events cannot supply a missing terminal model or response id.
    No identity is inferred from the requested model or HTTP sidecar.
    """
    out = {"id": None, "model": None, "usage": None, "finish_reason": None, "done": False, "error": False}
    values = [line[5:].strip() for line in data.splitlines() if line.startswith(b"data:")]
    if not values:
        values = [data]  # non-streaming provider response
    for value in values:
        if value == b"[DONE]":
            out["done"] = True
            continue
        try:
            event = json.loads(value)
        except (UnicodeDecodeError, json.JSONDecodeError):
            continue
        if not isinstance(event, dict):
            continue
        kind = event.get("type")
        response = event.get("response") if isinstance(event.get("response"), dict) else event
        out["error"] = out["error"] or bool(response.get("error")) or kind in {"error", "response.failed", "response.incomplete"}
        if isinstance(response.get("usage"), dict):
            out["usage"] = response["usage"]
        finish = next((choice.get("finish_reason") for choice in event.get("choices") or []
                       if isinstance(choice, dict) and choice.get("finish_reason")), None)
        if finish or kind == "response.completed" or response.get("status") == "completed":
            out["id"], out["model"] = response.get("id"), response.get("model")
            out["finish_reason"] = finish or "completed"
            if kind == "response.completed" or response.get("status") == "completed":
                out["done"] = True
    return out


class Route:
    """One upstream with the exact provider model ID and phase-specific request fields."""

    def __init__(self, *, name: str, base_url: str, model: str, phase: str, served_models=(),
                 api_key: str | None = None, thinking: dict | None = None, reasoning_effort: str | None = None,
                 credential: dict | None = None):
        self.name, self.model, self.phase = name, model, phase
        self.base_url = base_url.rstrip("/")
        self.url = urllib.parse.urlsplit(self.base_url)
        if self.url.scheme not in {"http", "https"} or not self.url.hostname:
            raise ValueError(f"route {name} needs an http(s) base_url")
        self.served_models = tuple(served_models)
        self.api_key, self.thinking, self.reasoning_effort = api_key, thinking, reasoning_effort
        self.credential = credential

    def selected_api_key(self) -> str | None:
        return resolve_credential(self.credential) if self.credential is not None else self.api_key

    def payload(self, body: dict) -> dict:
        from evalkit.request_view import normalize_chat_reasoning
        out = dict(body)
        out["messages"] = normalize_chat_reasoning(out["messages"])
        out["model"] = self.model
        if self.phase == "teacher":
            out["thinking"] = dict(self.thinking)
            out["reasoning_effort"] = self.reasoning_effort
        else:
            out.pop("thinking", None)
            out.pop("reasoning_effort", None)
        return out

    def receipt(self) -> dict:
        return {"name": self.name, "base_url": self.base_url, "model": self.model, "phase": self.phase,
                "served_models": list(self.served_models), "thinking": self.thinking,
                "reasoning_effort": self.reasoning_effort if self.phase == "teacher" else "not_applicable"}


def load_teacher_policy(path: Path) -> dict:
    policy = json.loads(Path(path).read_text(encoding="utf-8"))
    if policy.get("schema_version") != 1 or not policy.get("policy_id"):
        raise ValueError("teacher policy needs schema_version=1 and policy_id")
    if policy.get("thinking") != {"type": "enabled"} or policy.get("reasoning_effort") != "high":
        raise ValueError("teacher policy must explicitly enable thinking with reasoning_effort=high")
    routes = policy.get("routes")
    if not isinstance(routes, list) or not routes:
        raise ValueError("teacher policy needs an ordered routes list")
    if [route.get("name") for route in routes] != ["opencode-go", "deepseekapi"]:
        raise ValueError("teacher providers must be opencode-go then deepseekapi")
    for route in routes:
        if not route.get("model") or route.get("model") in {"auto", "default"} or not route.get("served_models"):
            raise ValueError("each teacher route needs an explicit provider model ID and served_models")
    if set(policy.get("fallback_on", [])) != FALLBACK_CLASSES:
        raise ValueError("fallback_on must be exactly quota and rate_limit")
    policy["sha256"] = sha256_bytes(Path(path).read_bytes())
    return policy


class TeacherRouter:
    """Ordered teacher routes with a bounded per-route cooldown after quota/rate limits."""

    def __init__(self, policy: dict, clock=time.time):
        self.policy, self.clock = policy, clock
        self.routes = [Route(name=r["name"], base_url=r["base_url"], model=r["model"], phase="teacher",
                             served_models=r["served_models"], credential=r.get("credential", {}),
                             thinking=policy["thinking"], reasoning_effort=policy["reasoning_effort"])
                       for r in policy["routes"]]
        self.blocked_until: dict[str, float] = {}

    def candidates(self) -> list[Route]:
        now = self.clock()
        return [route for route in self.routes if self.blocked_until.get(route.name, 0.0) <= now]

    def cool_down(self, route: Route, attempt: dict) -> None:
        now = self.clock()
        seconds = attempt.get("retry_after_s")
        if seconds is None:
            seconds = (self.policy.get("rate_limit_cooldown_s", 60) if attempt["error_class"] == "rate_limit"
                       else float("inf"))  # balance/quota does not recover within one task
        self.blocked_until[route.name] = max(self.blocked_until.get(route.name, 0.0), now + seconds)
        attempt["cooldown_s"] = None if seconds == float("inf") else seconds
        attempt["cooldown_scope"] = "task" if seconds == float("inf") else "bounded"

    def cooldowns(self) -> dict:
        now = self.clock()
        return {name: (None if until == float("inf") else round(until - now, 3))
                for name, until in self.blocked_until.items() if until > now}


class Meter:
    """Authoritative per-task token accounting shared by every request thread."""

    def __init__(self, budget: Budget, task_id: str, state_path: Path | None, receipts_path: Path | None):
        self.budget = budget
        self.task_id = task_id
        self.state_path = state_path
        self.receipts_path = receipts_path
        self.lock = threading.Lock()
        self.requests = 0
        self.output_used = 0
        self.input_used = 0
        self.max_input_seen = 0
        self.max_undercount = 0
        self.exhausted = False
        self.count_cache: dict[str, int] = {}
        self._load_state()
        if self.remaining_output < self.budget.min_output_tokens:
            self.exhausted = True

    # ---- persistence -------------------------------------------------------

    def _load_state(self) -> None:
        if self.state_path is None or not self.state_path.is_file():
            return
        try:
            state = json.loads(self.state_path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            return
        self.requests = int(state.get("requests", 0))
        self.output_used = int(state.get("output_used", 0))
        self.input_used = int(state.get("input_used", 0))
        self.max_input_seen = int(state.get("max_input_seen", 0))
        self.max_undercount = int(state.get("max_undercount", 0))
        self.exhausted = bool(state.get("exhausted", False)) or (
            self.remaining_output < self.budget.min_output_tokens
        )

    def _save_state_locked(self) -> None:
        if self.state_path is None:
            return
        state = {
            "task_id": self.task_id,
            "requests": self.requests,
            "output_used": self.output_used,
            "remaining_output_tokens": self.remaining_output,
            "exhausted": self.exhausted,
            "input_used": self.input_used,
            "max_input_seen": self.max_input_seen,
            "max_undercount": self.max_undercount,
            "updated_at": utc_now(),
        }
        tmp = self.state_path.with_suffix(self.state_path.suffix + ".tmp")
        tmp.write_text(json.dumps(state, indent=1), encoding="utf-8")
        os.replace(tmp, self.state_path)

    # ---- accounting --------------------------------------------------------

    @property
    def remaining_output(self) -> int:
        return max(0, self.budget.cumulative_output_max - self.output_used)

    def admit(self, input_tokens: int) -> tuple[int, int]:
        """Return ``(granted_max_tokens, remaining_output_before)`` or raise BudgetError."""
        with self.lock:
            try:
                if self.exhausted:
                    raise BudgetError(
                        f"task budget is exhausted (output_used={self.output_used}, "
                        f"remaining={self.remaining_output})"
                    )
                if self.budget.max_requests is not None and self.requests >= self.budget.max_requests:
                    raise BudgetError(
                        f"task request limit is exhausted (requests={self.requests}, "
                        f"max_requests={self.budget.max_requests})"
                    )
                remaining = self.remaining_output
                if remaining < self.budget.min_output_tokens:
                    raise BudgetError(
                        f"task cumulative output {self.output_used} of "
                        f"{self.budget.cumulative_output_max} leaves {remaining} < "
                        f"min_output_tokens {self.budget.min_output_tokens}"
                    )
                # Clamp the per-call request against the remaining task budget too.
                requested = self.budget.per_call_output_max
                granted = self.budget.grant(input_tokens, requested)
                granted = min(granted, remaining)
                if granted < self.budget.min_output_tokens:
                    raise BudgetError(
                        f"remaining output budget {remaining} cannot cover "
                        f"min_output_tokens {self.budget.min_output_tokens}"
                    )
                return granted, remaining
            except BudgetError:
                self.exhausted = True
                self._save_state_locked()
                raise

    def record(
        self,
        *,
        index: int,
        exact_input: int | None,
        reported_input: int | None,
        requested_max: int | None,
        granted_max: int,
        finish_reason: str | None,
        output_tokens: int | None,
        charged: int,
        stream: bool,
        method: str,
        elapsed_ms: int,
    ) -> None:
        with self.lock:
            self.requests += 1
            self.input_used += exact_input if exact_input is not None else (reported_input or 0)
            self.max_input_seen = max(self.max_input_seen, exact_input or reported_input or 0)
            if exact_input is not None and reported_input is not None and reported_input > exact_input:
                self.max_undercount = max(self.max_undercount, reported_input - exact_input)
            self.output_used += charged
            if self.remaining_output < self.budget.min_output_tokens:
                self.exhausted = True
            state = {
                "task_id": self.task_id,
                "requests": self.requests,
                "output_used": self.output_used,
                "remaining_output_tokens": self.remaining_output,
                "exhausted": self.exhausted,
                **({} if self.budget.max_requests is None else {
                    "max_requests": self.budget.max_requests,
                    "request_limit_exhausted": self.requests >= self.budget.max_requests,
                }),
                "input_used": self.input_used,
                "max_input_seen": self.max_input_seen,
                "max_undercount": self.max_undercount,
            }
            self._save_state_locked()
            if self.receipts_path is not None:
                receipt = {
                    "ts": utc_now(),
                    **state,
                    "index": index,
                    "exact_input_tokens": exact_input,
                    "reported_input_tokens": reported_input,
                    "undercount": None if reported_input is None or exact_input is None else reported_input - exact_input,
                    "requested_max_tokens": requested_max,
                    "granted_max_tokens": granted_max,
                    "finish_reason": finish_reason,
                    "output_tokens": output_tokens,
                    "charged_output_tokens": charged,
                    "cumulative_output_after": self.output_used,
                    "remaining_output_after": self.remaining_output,
                    "exhausted": self.exhausted,
                    "stream": stream,
                    "count_method": method,
                    "elapsed_ms": elapsed_ms,
                }
                with self.receipts_path.open("a", encoding="utf-8") as handle:
                    handle.write(json.dumps(receipt, ensure_ascii=False) + "\n")

    def snapshot(self) -> dict:
        with self.lock:
            return {
                "task_id": self.task_id,
                "profile": self.budget.profile_name,
                "context_window": self.budget.context_window,
                "per_call_output_max": self.budget.per_call_output_max,
                "cumulative_output_max": self.budget.cumulative_output_max,
                "min_output_tokens": self.budget.min_output_tokens,
                "reserve_tokens": self.budget.reserve_tokens,
                "input_soft_cap": self.budget.input_soft_cap,
                "input_hard_cap": self.budget.input_hard_cap,
                "requests": self.requests,
                "input_tokens_total": self.input_used,
                "output_used": self.output_used,
                "remaining_output_tokens": self.remaining_output,
                "exhausted": self.exhausted,
                **({} if self.budget.max_requests is None else {
                    "max_requests": self.budget.max_requests,
                    "request_limit_exhausted": self.requests >= self.budget.max_requests,
                }),
                "max_input_seen": self.max_input_seen,
                "max_observed_undercount": self.max_undercount,
            }


class ProxyHandler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    server_version = "EvalBudgetProxy/1"

    # ---- plumbing ----------------------------------------------------------

    def log_message(self, fmt: str, *args: Any) -> None:  # noqa: A003
        if self.server.verbose:  # type: ignore[attr-defined]
            sys.stderr.write(f"[proxy] {fmt % args}\n")
            sys.stderr.flush()

    @property
    def meter(self) -> Meter:
        return self.server.meter  # type: ignore[attr-defined]

    @property
    def upstream(self) -> urllib.parse.SplitResult:
        return self.server.upstream  # type: ignore[attr-defined]

    def _send_json(self, status: int, payload: dict) -> None:
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _send_budget_error(self, status: int, code: str, message: str, extra: dict | None = None) -> None:
        payload = {
            "error": {
                "type": "budget_error",
                "code": code,
                "message": message,
                **(extra or {}),
            }
        }
        self._send_json(status, payload)

    def _read_body(self) -> bytes:
        size = int(self.headers.get("Content-Length", "0") or 0)
        return self.rfile.read(size) if size else b""

    def _forward_headers(self, route: Route | None = None) -> dict[str, str]:
        headers = {}
        for key, value in self.headers.items():
            if key.lower() in HOP_BY_HOP or key.lower() == "accept-encoding":
                continue
            headers[key] = value
        headers["Accept-Encoding"] = "identity"
        upstream_key = route.selected_api_key() if route is not None else getattr(self.server, "upstream_api_key", None)
        if upstream_key or route is not None:
            # A routed request never forwards the harness's local dummy key.
            headers = {key: value for key, value in headers.items() if key.lower() != "authorization"}
        if upstream_key:
            headers["Authorization"] = "Bearer " + upstream_key
        return headers

    def _connection(self, upstream: urllib.parse.SplitResult | None = None):
        upstream = upstream or self.upstream
        connection_type = http.client.HTTPSConnection if upstream.scheme == "https" else http.client.HTTPConnection
        return connection_type(upstream.hostname, upstream.port,
                               timeout=self.server.upstream_timeout)

    @staticmethod
    def _response_reader(response):
        encoding = (response.getheader("Content-Encoding") or "identity").lower()
        if encoding == "gzip":
            return gzip.GzipFile(fileobj=response)
        if encoding != "identity":
            raise OSError("unsupported upstream content encoding: " + encoding)
        return response

    # ---- GET ---------------------------------------------------------------

    def do_GET(self) -> None:  # noqa: N802
        if self.path.split("?")[0].rstrip("/") == "/budget":
            self._send_json(200, self.meter.snapshot())
            return
        self._passthrough("GET", b"")

    # ---- POST --------------------------------------------------------------

    def do_POST(self) -> None:  # noqa: N802
        path = urllib.parse.urlsplit(self.path).path
        controller = getattr(self.server, "public_probe", None)
        if path.rstrip("/") == "/public-probe" and controller is not None:
            # Control-plane only: the frozen public probe runs in the task container.
            try:
                body = json.loads(self._read_body() or b"{}")
                result = controller.probe("stop", feedback_allowed=body.get("feedback_allowed") is True)
            except Exception as error:  # noqa: BLE001 - a probe failure must never become feedback
                result = {"status": "infra", "feedback": None, "error": str(error)[-500:]}
            self._send_json(200, result)
            return
        if not path.endswith("/chat/completions"):
            if getattr(self.server, "takeover", None) is not None:
                self._send_json(404, {"error": {"type": "invalid_request",
                                                "message": "routed mode serves chat completions only"}})
                return
            self._passthrough("POST", self._read_body())
            return
        started = time.monotonic()
        raw = self._read_body()
        try:
            body = json.loads(raw.decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError) as error:
            self._send_json(400, {"error": {"type": "invalid_request", "message": f"invalid JSON: {error}"}})
            return
        messages = body.get("messages")
        if not isinstance(messages, list) or not messages:
            self._send_json(400, {"error": {"type": "invalid_request", "message": "messages must be a non-empty list"}})
            return
        if getattr(self.server, "takeover", None) is not None:
            self._routed_chat(body, started)
            return
        model = body.get("model") or self.server.model  # type: ignore[attr-defined]
        count_mode = getattr(self.server, "count_mode", "server-tokenize")
        exact, method = (None, "provider-usage-no-input-preflight") if count_mode == "provider-usage" else self._count_exact(model, body)
        if exact is None and count_mode != "provider-usage":
            self._send_budget_error(
                502,
                "EVAL_TOKEN_COUNT_UNAVAILABLE",
                "the budget proxy could not obtain an exact token count from the server tokenizer; refusing to forward",
            )
            return
        try:
            granted, _remaining = self.meter.admit(exact or 0)
        except BudgetError as error:
            self._send_budget_error(
                400,
                "EVAL_BUDGET_EXHAUSTED",
                str(error),
                {"exact_input_tokens": exact, "requests": self.meter.requests},
            )
            return
        requested = body.get("max_tokens")
        if not isinstance(requested, int):
            requested = None
        body["max_tokens"] = granted
        if not body.get("stream"):
            body.pop("stream_options", None)
        else:
            options = body.get("stream_options")
            if not isinstance(options, dict):
                options = {}
            options["include_usage"] = True
            body["stream_options"] = options
        payload = json.dumps(body, ensure_ascii=False).encode("utf-8")
        capture_dir = getattr(self.server, "capture_dir", None)
        if capture_dir is not None:
            self.capture_prefix = capture_dir / f"{self.meter.requests + 1:04d}"
            self.capture_prefix.with_suffix(".request.json").write_bytes(payload)
        finish, output_tokens, reported_input = self._relay(payload, bool(body.get("stream")))
        output_tokens = output_tokens if output_tokens is not None else granted
        self.meter.record(
            index=self.meter.requests + 1,
            exact_input=exact,
            reported_input=reported_input,
            requested_max=requested,
            granted_max=granted,
            finish_reason=finish,
            output_tokens=output_tokens,
            charged=output_tokens,
            stream=bool(body.get("stream")),
            method=method,
            elapsed_ms=int((time.monotonic() - started) * 1000),
        )

    # ---- routed student/teacher mode --------------------------------------

    def _routed_chat(self, body: dict, started: float) -> None:
        """One logical decision: phase plan, budget admit, ordered route attempts, receipts.

        wire/NNNN.* holds only the final (successful or terminal) attempt;
        earlier quota/rate-limit attempts go to upstream-attempts/.
        """
        server = self.server
        with server.route_lock:  # type: ignore[attr-defined]
            takeover = server.takeover  # type: ignore[attr-defined]
            index = self.meter.requests + 1
            try:
                plan = takeover.plan(index, body)
            except Exception as error:  # noqa: BLE001 - probe/tokenizer failure is a recorded terminal state
                takeover.record_terminal("takeover_evaluation_failed", request_index=index,
                                         error=type(error).__name__ + ": " + str(error)[:300])
                self._send_json(502, {"error": {"type": "route_error", "code": "EVAL_TAKEOVER_EVALUATION_FAILED",
                                                "message": "takeover boundary could not be evaluated"}})
                return
            if "stop" in plan:
                self._send_budget_error(400, "EVAL_STUDENT_PREFIX_LIMIT",
                                        "student prefix exceeds the frozen limit; trajectory stops here",
                                        {"stop": plan["stop"]})
                return
            phase = plan["phase"]
            routes = [server.student_route] if phase == "student" else server.teacher_router.candidates()  # type: ignore[attr-defined]
            if not routes:
                takeover.record_terminal("teacher_routes_exhausted", request_index=index,
                                         cooldowns=server.teacher_router.cooldowns())  # type: ignore[attr-defined]
                self._send_json(503, {"error": {"type": "route_error", "code": "EVAL_TEACHER_ROUTES_EXHAUSTED",
                                                "message": "every teacher route is cooling down"}})
                return
            exact, method = None, "provider-usage-no-input-preflight"
            if phase == "student":
                exact, method = self._count_exact(routes[0].model, body, route=routes[0])
                if exact is None:
                    takeover.record_terminal("student_token_count_unavailable", request_index=index)
                    self._send_budget_error(502, "EVAL_TOKEN_COUNT_UNAVAILABLE",
                                            "student /tokenize unavailable; refusing to forward")
                    return
            try:
                granted, _remaining = self.meter.admit(exact or 0)
            except BudgetError as error:
                self._send_budget_error(400, "EVAL_BUDGET_EXHAUSTED", str(error),
                                        {"exact_input_tokens": exact, "requests": self.meter.requests})
                return
            requested = body.get("max_tokens") if isinstance(body.get("max_tokens"), int) else None
            body["max_tokens"] = granted
            stream = bool(body.get("stream"))
            if stream:
                options = body.get("stream_options") if isinstance(body.get("stream_options"), dict) else {}
                body["stream_options"] = {**options, "include_usage": True}
            else:
                body.pop("stream_options", None)
            attempts: list[dict] = []
            final = None
            for position, route in enumerate(routes):
                payload = json.dumps(route.payload(body), ensure_ascii=False).encode("utf-8")
                can_fall_back = phase == "teacher" and position < len(routes) - 1
                final = self._attempt(route, payload, index, len(attempts) + 1, stream, can_fall_back)
                attempts.append(final["attempt"])
                if phase == "teacher" and final["attempt"]["error_class"] in FALLBACK_CLASSES:
                    server.teacher_router.cool_down(route, final["attempt"])  # type: ignore[attr-defined]
                if not final["retry"]:
                    break
            attempt = final["attempt"]
            charged = (attempt["usage"] or {}).get("completion_tokens") if attempt["usage"] else None
            if charged is None:
                charged = granted if attempt["relayed"] else 0
            self.meter.record(
                index=index, exact_input=exact, reported_input=(attempt["usage"] or {}).get("prompt_tokens"),
                requested_max=requested, granted_max=granted, finish_reason=attempt.get("finish_reason"),
                output_tokens=(attempt["usage"] or {}).get("completion_tokens"), charged=charged, stream=stream,
                method=method, elapsed_ms=int((time.monotonic() - started) * 1000))
            if attempt["outcome"] != "completed":
                reason = ("credential_unavailable" if attempt["error_class"] == "credential" else
                          "partial_response_interrupted" if attempt["outcome"] == "stream_interrupted" else
                          "teacher_routes_exhausted" if phase == "teacher" and attempt["error_class"] in FALLBACK_CLASSES else
                          f"{phase}_upstream_error")
                takeover.record_terminal(reason, request_index=index, error_class=attempt["error_class"],
                                         http_status=attempt["http_status"],
                                         cooldowns=server.teacher_router.cooldowns() if phase == "teacher" else {})  # type: ignore[attr-defined]
            route = final["route"]
            row = {"request_index": index, "phase": phase, "provider": route.name,
                   "provider_base_url": route.base_url, "requested_model": route.model,
                   "served_model": attempt["served_model"], "served_model_expected":
                       attempt["served_model"] in route.served_models,
                   "reasoning_effort": route.reasoning_effort if phase == "teacher" else "not_applicable",
                   "thinking": route.thinking if phase == "teacher" else None,
                   "response_id": attempt["response_id"], "request_sha256": attempt["request_sha256"],
                   "response_sha256": attempt["response_sha256"], "http_status": attempt["http_status"],
                   "outcome": attempt["outcome"], "finish_reason": attempt.get("finish_reason"),
                   "usage": attempt["usage"], "attempts": attempts, "started_at": attempts[0]["started_at"]}
            with server.provenance_path.open("a", encoding="utf-8", newline="\n") as stream_out:  # type: ignore[attr-defined]
                stream_out.write(json.dumps(row, ensure_ascii=False) + "\n")

    def _attempt(self, route: Route, payload: bytes, index: int, number: int, stream: bool,
                 can_fall_back: bool) -> dict:
        """Send one attempt. Nothing reaches DSH unless this attempt becomes final."""
        server = self.server
        begun, wall = time.monotonic(), utc_now()
        attempt = {"attempt": number, "provider": route.name, "requested_model": route.model,
                   "http_status": None, "error_class": None, "retry_after_s": None, "usage": None,
                   "served_model": None, "response_id": None, "finish_reason": None, "relayed": False,
                   "request_sha256": sha256_bytes(payload), "response_sha256": None, "http_sha256": None, "path": None,
                   "outcome": "error", "started_at": wall}
        wire = server.capture_dir / f"{index:04d}"  # type: ignore[attr-defined]
        side = server.attempts_dir / f"{index:04d}.a{number}"  # type: ignore[attr-defined]
        connection = self._connection(route.url)
        try:
            try:
                headers = self._forward_headers(route)
            except CredentialUnavailable:
                # This local terminal attempt made no upstream request. Keep earlier quota
                # attempts intact and label the final response as proxy-generated evidence.
                failure = {"error": {"type": "route_error", "code": "EVAL_CREDENTIAL_UNAVAILABLE",
                                     "message": "selected provider credential is unavailable", "provider": route.name}}
                data = json.dumps(failure).encode("utf-8")
                meta = {"status": None, "response_source": "proxy_credential_error", "upstream_request_sent": False}
                Path(str(wire) + ".request.json").write_bytes(payload)
                Path(str(wire) + ".http.json").write_text(json.dumps(meta), encoding="utf-8")
                Path(str(wire) + ".response.json").write_bytes(data)
                attempt.update(error_class="credential", response_sha256=sha256_bytes(data),
                               http_sha256=sha256_bytes(Path(str(wire) + ".http.json").read_bytes()),
                               elapsed_ms=int((time.monotonic() - begun) * 1000), path="wire/" + wire.name)
                self._send_json(503, failure)
                return {"route": route, "attempt": attempt, "retry": False}
            headers.update({"Content-Type": "application/json", "Content-Length": str(len(payload))})
            try:
                connection.request("POST", self._endpoint(self.path, route.url), body=payload, headers=headers)
                response = connection.getresponse()
            except (OSError, http.client.HTTPException) as error:
                attempt.update(error_class="network", error=type(error).__name__ + ": " + str(error)[:300],
                               elapsed_ms=int((time.monotonic() - begun) * 1000), path="wire/" + wire.name)
                Path(str(wire) + ".request.json").write_bytes(payload)
                meta = {"status": None, "error": attempt["error"], "response_source": "proxy_network_error"}
                Path(str(wire) + ".http.json").write_text(json.dumps(meta), encoding="utf-8")
                failure = {"error": {"type": "upstream_error", "code": "EVAL_UPSTREAM_UNREACHABLE",
                                     "message": attempt["error"]}}
                data = json.dumps(failure).encode("utf-8")
                Path(str(wire) + ".response.json").write_bytes(data)
                attempt["response_sha256"] = sha256_bytes(data)
                attempt["http_sha256"] = sha256_bytes(Path(str(wire) + ".http.json").read_bytes())
                self._send_json(502, failure)
                return {"route": route, "attempt": attempt, "retry": False}
            attempt["http_status"] = response.status
            meta = {"status": response.status, "headers": {
                key: value for key, value in response.getheaders()
                if key.lower() in {"content-type", "content-encoding", "x-request-id", "request-id",
                                   "x-magpie-request-id", "retry-after"}}, "captured_body_encoding": "identity"}
            reader = self._response_reader(response)
            prefetched = b""
            if response.status == 200 and stream:
                # Look ahead to the first SSE event: an error event there is still a clean failure.
                read = getattr(reader, "read1", reader.read)
                while b"\n\n" not in prefetched.replace(b"\r\n", b"\n"):
                    chunk = read(8192)
                    if not chunk:
                        break
                    prefetched += chunk
                first = next((line[5:].strip() for line in prefetched.splitlines() if line.startswith(b"data:")), b"")
                failed_early = bool(error_object(first)) or not first
            else:
                failed_early = response.status != 200
            if failed_early:
                if response.status == 200:
                    data = prefetched + reader.read()
                    event = next((line[5:].strip() for line in data.splitlines() if line.startswith(b"data:")), b"")
                    attempt["error_class"] = classify_failure(200, event) if error_object(event) else "network"
                else:
                    data = reader.read()
                    attempt["error_class"] = classify_failure(response.status, data)
                attempt["retry_after_s"] = retry_after_seconds(response.getheader("Retry-After"), time.time())
                attempt["response_sha256"] = sha256_bytes(data)
                attempt["elapsed_ms"] = int((time.monotonic() - begun) * 1000)
                summary = sse_summary(data)
                attempt["usage"] = summary["usage"] if isinstance(summary["usage"], dict) else None
                retry = can_fall_back and attempt["error_class"] in FALLBACK_CLASSES
                target = side if retry else wire
                attempt["path"] = str(target.relative_to(server.capture_dir.parent)).replace("\\", "/")  # type: ignore[attr-defined]
                target.parent.mkdir(parents=True, exist_ok=True)
                Path(str(target) + ".request.json").write_bytes(payload)
                Path(str(target) + ".http.json").write_text(json.dumps(meta), encoding="utf-8")
                Path(str(target) + ".response.json").write_bytes(data)
                attempt["http_sha256"] = sha256_bytes(Path(str(target) + ".http.json").read_bytes())
                if not retry:
                    self.send_response(response.status if response.status != 200 else 502)
                    for key, value in response.getheaders():
                        if key.lower() in {"content-length", "transfer-encoding", "connection", "content-encoding"}:
                            continue
                        self.send_header(key, value)
                    self.send_header("Content-Length", str(len(data)))
                    self.end_headers()
                    self.wfile.write(data)
                return {"route": route, "attempt": attempt, "retry": retry}
            attempt["path"] = "wire/" + wire.name
            Path(str(wire) + ".request.json").write_bytes(payload)
            Path(str(wire) + ".http.json").write_text(json.dumps(meta), encoding="utf-8")
            attempt["http_sha256"] = sha256_bytes(Path(str(wire) + ".http.json").read_bytes())
            self.capture_prefix = wire
            attempt["relayed"] = True
            try:
                if stream:
                    finish, output, prompt = self._relay_stream(response, prefetched, reader)
                    captured = Path(str(wire) + ".response.sse").read_bytes()
                    summary = sse_summary(captured)
                    finish = summary['finish_reason']
                else:
                    data = reader.read()
                    Path(str(wire) + ".response.json").write_bytes(data)
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Content-Length", str(len(data)))
                    self.end_headers()
                    self.wfile.write(data)
                    captured = data
                    summary = sse_summary(captured)
                    finish = summary["finish_reason"]
            except (OSError, http.client.HTTPException) as error:
                captured = Path(str(wire) + ".response.sse").read_bytes() if Path(str(wire) + ".response.sse").exists() else b""
                summary = {**sse_summary(captured), "done": False}
                attempt["error"] = type(error).__name__ + ": " + str(error)[:300]
                finish = summary.get("finish_reason")
            attempt.update(response_id=summary.get("id"), served_model=summary.get("model"),
                           usage=summary.get("usage") if isinstance(summary.get("usage"), dict) else None,
                           finish_reason=finish, response_sha256=sha256_bytes(captured),
                           elapsed_ms=int((time.monotonic() - begun) * 1000))
            if (summary.get("done") or summary.get("finish_reason")) and not summary.get("error"):
                attempt["outcome"] = "completed"
            else:
                attempt["outcome"] = "stream_interrupted"
                attempt["error_class"] = "stream_interrupted"
            return {"route": route, "attempt": attempt, "retry": False}
        finally:
            connection.close()

    # ---- token counting ----------------------------------------------------

    def _count_exact(self, model: str, body: dict, route: Route | None = None) -> tuple[int | None, str]:
        """Count the exact serialized prompt through the server tokenizer."""
        tokenize_payload: dict[str, Any] = {
            "model": model,
            "messages": body.get("messages"),
            "add_generation_prompt": True,
        }
        for key in ("tools", "tool_choice", "chat_template_kwargs"):
            if body.get(key) is not None:
                tokenize_payload[key] = body[key]
        cache_key = hashlib.sha256(
            json.dumps(tokenize_payload, sort_keys=True, ensure_ascii=False).encode("utf-8")
        ).hexdigest()
        meter = self.meter
        with meter.lock:
            cached = meter.count_cache.get(cache_key)
        if cached is not None:
            return cached, "server-tokenize-cache"
        try:
            # vLLM exposes /tokenize at the server root, not under /v1.
            parsed = self._request_upstream_json("/tokenize", tokenize_payload, root=True, route=route)
        except (OSError, http.client.HTTPException, ValueError) as error:
            sys.stderr.write(f"[proxy] /tokenize failed: {error}\n")
            return None, "unavailable"
        count = parsed.get("count") if isinstance(parsed, dict) else None
        if not isinstance(count, int) or count < 0:
            # Older servers name the field differently; accept prompt_tokens too.
            count = parsed.get("prompt_tokens") if isinstance(parsed, dict) else None
        if not isinstance(count, int) or count < 0:
            return None, "invalid-tokenize-response"
        with meter.lock:
            if len(meter.count_cache) > 512:
                meter.count_cache.clear()
            meter.count_cache[cache_key] = count
        return count, "server-tokenize"

    def _request_upstream_json(self, path: str, payload: dict, root: bool = False,
                               route: Route | None = None) -> Any:
        connection = self._connection(route.url if route else None)
        try:
            headers = {"Content-Type": "application/json"}
            api_key = self._forward_headers(route).get("Authorization")
            if api_key:
                headers["Authorization"] = api_key
            endpoint = path if root else self._endpoint(path, route.url if route else None)
            connection.request("POST", endpoint, body=json.dumps(payload).encode("utf-8"), headers=headers)
            response = connection.getresponse()
            data = self._response_reader(response).read()
            if response.status >= 400:
                raise ValueError(f"upstream {path} returned {response.status}: {data[:300]!r}")
            return json.loads(data.decode("utf-8"))
        finally:
            connection.close()

    # ---- relay -------------------------------------------------------------

    def _relay(self, payload: bytes, stream: bool) -> tuple[str | None, int | None, int | None]:
        """Forward the rewritten request and return (finish_reason, output_tokens, reported_input)."""
        connection = self._connection()
        headers = self._forward_headers()
        headers["Content-Type"] = "application/json"
        headers["Content-Length"] = str(len(payload))
        connection.request("POST", self._endpoint(self.path), body=payload, headers=headers)
        response = connection.getresponse()
        capture_prefix = getattr(self, "capture_prefix", None)
        if capture_prefix is not None:
            capture_prefix.with_suffix(".http.json").write_text(json.dumps({
                "status": response.status,
                "headers": {key: value for key, value in response.getheaders()
                            if key.lower() in {"content-type", "content-encoding", "x-request-id", "request-id", "x-magpie-request-id"}},
                "captured_body_encoding": "identity",
            }), encoding="utf-8")
        try:
            if not stream:
                data = self._response_reader(response).read()
                if capture_prefix is not None:
                    capture_prefix.with_suffix(".response.json").write_bytes(data)
                self.send_response(response.status)
                for key, value in response.getheaders():
                    if key.lower() in {"content-length", "transfer-encoding", "connection", "content-encoding"}:
                        continue
                    self.send_header(key, value)
                self.send_header("Content-Length", str(len(data)))
                self.end_headers()
                self.wfile.write(data)
                parsed = self._safe_json(data)
                usage = parsed.get("usage") if isinstance(parsed, dict) else None
                choice = (parsed.get("choices") or [{}])[0] if isinstance(parsed, dict) else {}
                return (
                    choice.get("finish_reason"),
                    usage.get("completion_tokens") if isinstance(usage, dict) else None,
                    usage.get("prompt_tokens") if isinstance(usage, dict) else None,
                )
            return self._relay_stream(response)
        finally:
            connection.close()

    def _relay_stream(self, response: http.client.HTTPResponse, prefetched: bytes = b"",
                      reader=None) -> tuple[str | None, int | None, int | None]:
        reader = reader or self._response_reader(response)
        self.send_response(response.status)
        for key, value in response.getheaders():
            if key.lower() in {"content-length", "transfer-encoding", "connection", "content-encoding"}:
                continue
            self.send_header(key, value)
        self.send_header("Transfer-Encoding", "chunked")
        self.end_headers()
        finish_reason: str | None = None
        output_tokens: int | None = None
        reported_input: int | None = None
        buffer = b""
        capture_prefix = getattr(self, "capture_prefix", None)
        capture = capture_prefix.with_suffix(".response.sse").open("wb") if capture_prefix is not None else None
        pending = [prefetched] if prefetched else []
        while True:
            chunk = pending.pop() if pending else reader.read(8192)
            if not chunk:
                break
            if capture is not None:
                capture.write(chunk)
                capture.flush()
            buffer += chunk
            self.wfile.write(f"{len(chunk):X}\r\n".encode("ascii"))
            self.wfile.write(chunk)
            self.wfile.write(b"\r\n")
            self.wfile.flush()
            while b"\n" in buffer:
                line, buffer = buffer.split(b"\n", 1)
                finish_reason, output_tokens, reported_input = self._parse_sse_line(
                    line, finish_reason, output_tokens, reported_input
                )
        # Trailing partial line.
        if buffer.strip():
            finish_reason, output_tokens, reported_input = self._parse_sse_line(
                buffer, finish_reason, output_tokens, reported_input
            )
        self.wfile.write(b"0\r\n\r\n")
        self.wfile.flush()
        if capture is not None:
            capture.close()
        return finish_reason, output_tokens, reported_input

    @staticmethod
    def _parse_sse_line(
        line: bytes, finish_reason: str | None, output_tokens: int | None, reported_input: int | None
    ) -> tuple[str | None, int | None, int | None]:
        if not line.startswith(b"data:"):
            return finish_reason, output_tokens, reported_input
        data = line[5:].strip()
        if not data or data == b"[DONE]":
            return finish_reason, output_tokens, reported_input
        try:
            parsed = json.loads(data.decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError):
            return finish_reason, output_tokens, reported_input
        if not isinstance(parsed, dict):
            return finish_reason, output_tokens, reported_input
        usage = parsed.get("usage")
        if isinstance(usage, dict):
            if isinstance(usage.get("completion_tokens"), int):
                output_tokens = usage["completion_tokens"]
            if isinstance(usage.get("prompt_tokens"), int):
                reported_input = usage["prompt_tokens"]
        choices = parsed.get("choices")
        if isinstance(choices, list) and choices and isinstance(choices[0], dict):
            reason = choices[0].get("finish_reason")
            if isinstance(reason, str):
                finish_reason = reason
        return finish_reason, output_tokens, reported_input

    @staticmethod
    def _safe_json(data: bytes) -> Any:
        try:
            return json.loads(data.decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError):
            return None

    # ---- generic passthrough ----------------------------------------------

    def _endpoint(self, path: str, upstream: urllib.parse.SplitResult | None = None) -> str:
        """Join the client/upstream path without duplicating the upstream prefix."""
        base = (upstream or self.upstream).path or ""
        if base and path != base and not path.startswith(base + "/"):
            return base + path
        return path

    def _passthrough(self, method: str, payload: bytes) -> None:
        connection = self._connection()
        try:
            headers = self._forward_headers()
            if payload:
                headers["Content-Length"] = str(len(payload))
            connection.request(method, self._endpoint(self.path), body=payload or None, headers=headers)
            response = connection.getresponse()
            data = self._response_reader(response).read()
            self.send_response(response.status)
            for key, value in response.getheaders():
                if key.lower() in {"content-length", "transfer-encoding", "connection", "content-encoding"}:
                    continue
                self.send_header(key, value)
            self.send_header("Content-Length", str(len(data)))
            self.end_headers()
            self.wfile.write(data)
        finally:
            connection.close()


def main() -> None:
    parser = argparse.ArgumentParser(description="Exact token-budget proxy for the eval dsh harness")
    parser.add_argument("--upstream", default=os.environ.get("EVAL_UPSTREAM", "http://127.0.0.1:8013/v1"))
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=int(os.environ.get("EVAL_BUDGET_PORT", "18080")))
    parser.add_argument("--task-id", default=os.environ.get("EVAL_TASK_ID", "unset"))
    parser.add_argument("--model", default=os.environ.get("EVAL_MODEL_NAME", "Qwen/Qwen3.5-4B"))
    parser.add_argument("--config", type=Path, default=None)
    parser.add_argument("--receipts", type=Path, default=None)
    parser.add_argument("--state", type=Path, default=None)
    parser.add_argument("--upstream-timeout", type=float, default=1800.0)
    parser.add_argument("--verbose", action="store_true")
    parser.add_argument("--count-mode", choices=("server-tokenize", "provider-usage"), default="server-tokenize",
                        help="provider-usage enforces output caps but cannot verify input token IDs or window preflight")
    parser.add_argument("--capture-dir", type=Path)
    args = parser.parse_args()

    budget = load_budget(args.config)
    receipts = args.receipts
    state = args.state
    if receipts is not None:
        receipts.parent.mkdir(parents=True, exist_ok=True)
    if state is not None:
        state.parent.mkdir(parents=True, exist_ok=True)
    meter = Meter(budget, args.task_id, state, receipts)

    upstream = urllib.parse.urlsplit(args.upstream.rstrip("/"))
    if upstream.scheme not in {"http"}:
        raise SystemExit("only http upstreams are supported")
    server = ThreadingHTTPServer((args.host, args.port), ProxyHandler)
    server.daemon_threads = True
    server.meter = meter  # type: ignore[attr-defined]
    server.upstream = upstream  # type: ignore[attr-defined]
    server.model = args.model  # type: ignore[attr-defined]
    server.verbose = args.verbose  # type: ignore[attr-defined]
    server.upstream_timeout = args.upstream_timeout  # type: ignore[attr-defined]
    server.count_mode = args.count_mode  # type: ignore[attr-defined]
    server.capture_dir = args.capture_dir  # type: ignore[attr-defined]
    if args.capture_dir:
        args.capture_dir.mkdir(parents=True, exist_ok=True)
    print(
        json.dumps(
            {
                "proxy": f"http://{args.host}:{args.port}",
                "upstream": args.upstream,
                "task_id": args.task_id,
                "budget": budget.as_receipt(),
            },
            ensure_ascii=False,
        ),
        flush=True,
    )
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        server.server_close()


if __name__ == "__main__":
    main()