#!/usr/bin/env python3
"""Bound Switchyard classifier context before forwarding it to local Ollama."""

from __future__ import annotations

import copy
import json
import os
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any, Final

import httpx


HOST: Final = os.environ.get("HERMES_CLASSIFIER_BROKER_HOST", "0.0.0.0")
PORT: Final = int(os.environ.get("HERMES_CLASSIFIER_BROKER_PORT", "9008"))
UPSTREAM: Final = os.environ.get(
    "HERMES_CLASSIFIER_BROKER_UPSTREAM",
    "http://ollama.ai.svc.cluster.local:11434",
).rstrip("/")
MAX_BODY_BYTES: Final = int(
    os.environ.get("HERMES_CLASSIFIER_BROKER_MAX_BODY", str(32 << 20))
)
MAX_SYSTEM_CHARS: Final = int(
    os.environ.get("HERMES_CLASSIFIER_MAX_SYSTEM_CHARS", "6500")
)
MAX_CONTEXT_CHARS: Final = int(
    os.environ.get("HERMES_CLASSIFIER_MAX_CONTEXT_CHARS", "7000")
)
READ_TIMEOUT_SECONDS: Final = float(
    os.environ.get("HERMES_CLASSIFIER_BROKER_READ_TIMEOUT", "60")
)
ALLOWED_PATHS: Final = {"/v1/chat/completions", "/v1/models"}


def _bounded_text(value: str, limit: int) -> str:
    """Keep both ends of text because intent and current status often sit apart."""
    if len(value) <= limit:
        return value
    if limit < 80:
        return value[:limit]
    marker = "\n...[classifier context compacted]...\n"
    remaining = limit - len(marker)
    head = remaining // 2
    return f"{value[:head]}{marker}{value[-(remaining - head):]}"


def _compact_content(content: Any, limit: int) -> Any:
    """Remove binary/multimodal payloads and bound text sent to the judge."""
    if isinstance(content, str):
        return _bounded_text(content, limit)
    if not isinstance(content, list):
        return content

    compacted: list[Any] = []
    remaining = limit
    for block in content:
        if not isinstance(block, dict):
            continue
        kind = str(block.get("type") or "")
        if kind in {"text", "input_text", "output_text"}:
            key = "text"
            text = str(block.get(key) or "")
            if not text or remaining <= 0:
                continue
            text = _bounded_text(text, remaining)
            compacted.append({**block, key: text})
            remaining -= len(text)
        elif kind in {"image", "image_url", "input_image"}:
            marker = "[image attachment available to the selected worker]"
            if remaining >= len(marker):
                compacted.append({"type": "text", "text": marker})
                remaining -= len(marker)
    return compacted


def _compact_message(message: dict[str, Any], limit: int) -> dict[str, Any]:
    """Copy routing-relevant message metadata while bounding large values."""
    result = copy.deepcopy(message)
    if "content" in result:
        result["content"] = _compact_content(result["content"], limit)
    tool_calls = result.get("tool_calls")
    if isinstance(tool_calls, list):
        kept: list[dict[str, Any]] = []
        for call in tool_calls[-4:]:
            if not isinstance(call, dict):
                continue
            item = copy.deepcopy(call)
            function = item.get("function")
            if isinstance(function, dict) and isinstance(function.get("arguments"), str):
                function["arguments"] = _bounded_text(function["arguments"], 600)
            kept.append(item)
        result["tool_calls"] = kept
    return result


def _compact_messages(messages: list[Any]) -> list[Any]:
    """Keep the routing contract, opening task, and latest decision context."""
    valid = [message for message in messages if isinstance(message, dict)]
    system_indices = [
        index
        for index, message in enumerate(valid)
        if message.get("role") in {"system", "developer"}
    ]
    non_system = [index for index in range(len(valid)) if index not in system_indices]
    user_indices = [index for index in non_system if valid[index].get("role") == "user"]

    opening_user = user_indices[0] if user_indices else None
    latest_user = user_indices[-1] if user_indices else None
    selected = set(system_indices)
    selected.update(non_system[-4:])
    if opening_user is not None:
        selected.add(opening_user)
    if latest_user is not None:
        selected.add(latest_user)

    other_indices = [
        index
        for index in selected
        if index not in system_indices and index not in {opening_user, latest_user}
    ]
    latest_budget = min(3500, MAX_CONTEXT_CHARS)
    opening_budget = min(1200, max(0, MAX_CONTEXT_CHARS - latest_budget))
    other_budget = max(0, MAX_CONTEXT_CHARS - latest_budget - opening_budget)
    other_limit = min(1200, other_budget // max(1, len(other_indices)))
    system_limit = MAX_SYSTEM_CHARS // max(1, len(system_indices))

    result: list[Any] = []
    for index, message in enumerate(valid):
        if index not in selected:
            continue
        role = str(message.get("role") or "")
        if role in {"system", "developer"}:
            compacted = _compact_message(message, system_limit)
        elif index == latest_user:
            compacted = _compact_message(message, latest_budget)
        elif index == opening_user:
            compacted = _compact_message(message, opening_budget)
        else:
            if other_limit <= 0:
                continue
            compacted = _compact_message(message, other_limit)
        result.append(compacted)
    return result


def compact_payload(payload: dict[str, Any]) -> dict[str, Any]:
    """Return the same OpenAI request with only classifier input compacted."""
    result = copy.deepcopy(payload)
    messages = result.get("messages")
    if isinstance(messages, list):
        result["messages"] = _compact_messages(messages)
    # The judge never needs tools or binary inputs. Switchyard supplies a
    # response schema separately, and that contract must remain untouched.
    result.pop("tools", None)
    return result


class Handler(BaseHTTPRequestHandler):
    """Proxy only Switchyard's local judge calls with strict input bounds."""

    server_version = "HermesClassifierBroker/1"

    def log_message(self, format: str, *args: object) -> None:
        return

    def _json(self, status: int, value: dict[str, object]) -> None:
        body = json.dumps(value, separators=(",", ":")).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self) -> None:  # noqa: N802
        if self.path == "/health":
            self._json(200, {"ok": True, "upstream": "ollama"})
            return
        if self.path not in ALLOWED_PATHS:
            self._json(404, {"error": "not found"})
            return
        self._proxy(None)

    def do_POST(self) -> None:  # noqa: N802
        if self.path not in ALLOWED_PATHS:
            self._json(404, {"error": "not found"})
            return
        try:
            length = int(self.headers.get("Content-Length", "0"))
        except ValueError:
            length = -1
        if length <= 0 or length > MAX_BODY_BYTES:
            self._json(413, {"error": "request too large"})
            return
        try:
            payload = json.loads(self.rfile.read(length))
            if not isinstance(payload, dict):
                raise ValueError("request must be a JSON object")
            body = json.dumps(
                compact_payload(payload), separators=(",", ":")
            ).encode("utf-8")
        except (ValueError, TypeError, json.JSONDecodeError) as exc:
            self._json(400, {"error": str(exc)})
            return
        print(
            f"classifier-broker request_bytes={length} compacted_bytes={len(body)}",
            flush=True,
        )
        self._proxy(body)

    def _proxy(self, body: bytes | None) -> None:
        try:
            timeout = httpx.Timeout(10.0, read=READ_TIMEOUT_SECONDS)
            with httpx.Client(timeout=timeout) as client:
                response = client.request(
                    self.command,
                    f"{UPSTREAM}{self.path}",
                    headers={"Content-Type": "application/json"},
                    content=body,
                )
            self.send_response(response.status_code)
            self.send_header(
                "Content-Type", response.headers.get("Content-Type", "application/json")
            )
            self.send_header("Content-Length", str(len(response.content)))
            self.send_header("Cache-Control", "no-store")
            self.end_headers()
            self.wfile.write(response.content)
        except (httpx.HTTPError, OSError) as exc:
            self._json(503, {"error": f"classifier unavailable: {exc}"})


def main() -> None:
    server = ThreadingHTTPServer((HOST, PORT), Handler)
    server.serve_forever()


if __name__ == "__main__":
    main()
