#!/usr/bin/env python3
"""Credential-isolating Codex Responses proxy for Hermes chat tenants."""

from __future__ import annotations

import base64
import binascii
import hmac
import json
import os
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any, Iterable

import httpx

from routing_catalog import resolve_route


HOST = os.environ.get("HERMES_CODEX_BROKER_LISTEN_HOST", "0.0.0.0")
PORT = int(os.environ.get("HERMES_CODEX_BROKER_LISTEN_PORT", "9003"))
TOKEN = os.environ.get("HERMES_IMAGE_BROKER_KEY", "").strip()
UPSTREAM = os.environ.get(
    "HERMES_CODEX_BROKER_UPSTREAM",
    "https://chatgpt.com/backend-api/codex",
).rstrip("/")
MAX_BODY_BYTES = int(os.environ.get("HERMES_CODEX_BROKER_MAX_BODY", str(64 << 20)))
READ_TIMEOUT_SECONDS = float(os.environ.get("HERMES_CODEX_BROKER_READ_TIMEOUT", "900"))
FALLBACK_ALLOWED_MODELS = {
    value.strip()
    for value in os.environ.get(
        "HERMES_CODEX_BROKER_MODELS",
        "gpt-5.6-luna,gpt-5.6-sol,gpt-5.6-terra",
    ).split(",")
    if value.strip()
}
ROUTED_MODEL_PREFIX = "route/codex/"


def _real_model(model: str) -> str:
    """Translate a Switchyard effort-qualified target into a Codex model."""
    if model.startswith(ROUTED_MODEL_PREFIX):
        model = resolve_route(model)
    return model


def _authorized(header: str | None) -> bool:
    """Authenticate a tenant without exposing the relay secret."""
    if not TOKEN or not header or not header.startswith("Bearer "):
        return False
    return hmac.compare_digest(header[7:].strip(), TOKEN)


def _access_token() -> str:
    """Read the current owner token; the Codex CLI remains refresh owner."""
    codex_home = Path(
        os.environ.get("CODEX_HOME", str(Path.home() / ".codex"))
    ).expanduser()
    try:
        payload = json.loads((codex_home / "auth.json").read_text(encoding="utf-8"))
    except (OSError, ValueError) as exc:
        raise RuntimeError("owner Codex authentication is unavailable") from exc
    if not isinstance(payload, dict):
        raise RuntimeError("owner Codex authentication is invalid")
    tokens = payload.get("tokens") or {}
    token = tokens.get("access_token") if isinstance(tokens, dict) else None
    if not isinstance(token, str) or not token.strip():
        raise RuntimeError("owner Codex access token is unavailable")
    token = token.strip()
    try:
        encoded = token.split(".")[1]
        encoded += "=" * (-len(encoded) % 4)
        expires_at = json.loads(base64.urlsafe_b64decode(encoded)).get("exp", 0)
    except (IndexError, ValueError, TypeError, json.JSONDecodeError, binascii.Error):
        expires_at = 0
    if expires_at and time.time() >= float(expires_at):
        raise RuntimeError("owner Codex access token is expired")
    return token


def _upstream_headers(token: str) -> dict[str, str]:
    """Build the first-party headers expected by the Codex backend."""
    from agent.auxiliary_client import _codex_cloudflare_headers

    headers = _codex_cloudflare_headers(token)
    headers.update(
        {
            "Accept": "text/event-stream",
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        }
    )
    return headers


def _validate_payload(payload: Any) -> dict[str, Any]:
    """Allow only bounded Responses requests for the approved model catalog."""
    if not isinstance(payload, dict):
        raise ValueError("JSON object required")
    model = payload.get("model")
    if not isinstance(model, str):
        raise ValueError("unsupported Codex model")
    model = _real_model(model)
    if not model.startswith("gpt-"):
        raise ValueError("unsupported Codex model")
    payload["model"] = model
    response_input = payload.get("input")
    if isinstance(response_input, str):
        if not response_input.strip():
            raise ValueError("non-empty Responses input required")
        # Switchyard accepts OpenAI Chat Completions requests and translates
        # their final text to a scalar Responses input.  The first-party Codex
        # endpoint is stricter than the public Responses API and only accepts
        # a list of typed input items.
        payload["input"] = [
            {
                "type": "message",
                "role": "user",
                "content": [{"type": "input_text", "text": response_input}],
            }
        ]
    elif isinstance(response_input, dict):
        payload["input"] = [response_input]
    elif not isinstance(response_input, list) or not response_input:
        raise ValueError("non-empty Responses input list required")
    # Tenant conversations must not enter the owner's server-side history.
    payload["store"] = False
    payload["stream"] = True
    return payload


def _completed_response(lines: Iterable[str]) -> dict[str, Any]:
    """Collapse a Codex SSE stream for a non-streaming Responses caller."""
    terminal_response: dict[str, Any] | None = None
    output_items: dict[int, dict[str, Any]] = {}
    upstream_error = ""
    for line in lines:
        if not line.startswith("data:"):
            continue
        value = line[5:].strip()
        if not value or value == "[DONE]":
            continue
        try:
            event = json.loads(value)
        except (TypeError, ValueError, json.JSONDecodeError):
            continue
        if not isinstance(event, dict):
            continue
        event_type = str(event.get("type") or "")
        response = event.get("response")
        if event_type == "response.output_item.done":
            item = event.get("item")
            output_index = event.get("output_index")
            if isinstance(item, dict) and isinstance(output_index, int):
                output_items[output_index] = item
        if event_type in {
            "response.completed",
            "response.failed",
            "response.incomplete",
        } and isinstance(response, dict):
            terminal_response = response
        elif event_type == "error":
            error = event.get("error")
            if isinstance(error, dict):
                upstream_error = str(error.get("message") or error.get("type") or "")
            else:
                upstream_error = str(error or "")
    if terminal_response is not None:
        # The subscription Codex endpoint streams complete output items but
        # currently leaves the terminal response's output array empty.  Public
        # Responses clients, including Switchyard, expect those items there.
        if output_items and not terminal_response.get("output"):
            terminal_response["output"] = [
                output_items[index] for index in sorted(output_items)
            ]
        return terminal_response
    raise RuntimeError(upstream_error or "Codex stream ended without a terminal response")


class Handler(BaseHTTPRequestHandler):
    """Authenticated streaming proxy; request bodies and tokens are never logged."""

    server_version = "HermesCodexBroker/1"

    def _json(self, status: int, value: dict[str, Any]) -> 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 _check_auth(self) -> bool:
        if _authorized(self.headers.get("Authorization")):
            return True
        self._json(401, {"error": {"message": "unauthorized", "type": "auth_error"}})
        return False

    def do_GET(self) -> None:  # noqa: N802 - BaseHTTPRequestHandler API
        if not self._check_auth():
            return
        if self.path == "/health":
            try:
                _access_token()
            except RuntimeError as exc:
                self._json(503, {"ok": False, "error": str(exc)})
                return
            self._json(200, {"ok": True, "provider": "openai-codex"})
            return
        if self.path in {"/models", "/v1/models"}:
            self._json(
                200,
                {
                    "object": "list",
                    "data": [
                        {"id": model, "object": "model", "owned_by": "openai-codex"}
                        for model in sorted(FALLBACK_ALLOWED_MODELS)
                    ],
                },
            )
            return
        self._json(404, {"error": {"message": "not found", "type": "not_found"}})

    def do_POST(self) -> None:  # noqa: N802 - BaseHTTPRequestHandler API
        if self.path not in {"/responses", "/v1/responses"}:
            self._json(404, {"error": {"message": "not found", "type": "not_found"}})
            return
        if not self._check_auth():
            return
        response_started = False
        try:
            length = int(self.headers.get("Content-Length", "0"))
        except ValueError:
            length = 0
        if length <= 0 or length > MAX_BODY_BYTES:
            self._json(413, {"error": {"message": "invalid request size", "type": "invalid_request_error"}})
            return

        try:
            payload = json.loads(self.rfile.read(length))
            requested_stream = payload.get("stream") is True
            payload = _validate_payload(payload)
            token = _access_token()
            timeout = httpx.Timeout(
                READ_TIMEOUT_SECONDS,
                connect=30.0,
                read=READ_TIMEOUT_SECONDS,
                write=60.0,
                pool=30.0,
            )
            with httpx.Client(timeout=timeout, headers=_upstream_headers(token)) as client:
                with client.stream("POST", f"{UPSTREAM}/responses", json=payload) as response:
                    if response.status_code >= 400:
                        body = response.read()
                        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(body)))
                        self.send_header("Cache-Control", "no-store")
                        self.end_headers()
                        self.wfile.write(body)
                        return

                    if not requested_stream:
                        self._json(200, _completed_response(response.iter_lines()))
                        return

                    # HTTP/1.0 close-delimited streaming avoids buffering a
                    # potentially long tool-calling turn in the broker.
                    self.send_response(200)
                    self.send_header("Content-Type", "text/event-stream")
                    self.send_header("Cache-Control", "no-store")
                    self.send_header("X-Accel-Buffering", "no")
                    self.end_headers()
                    response_started = True
                    for chunk in response.iter_raw():
                        if chunk:
                            self.wfile.write(chunk)
                            self.wfile.flush()
        except ValueError as exc:
            self._json(400, {"error": {"message": str(exc), "type": "invalid_request_error"}})
        except (BrokenPipeError, ConnectionResetError):
            return
        except Exception as exc:
            if response_started:
                return
            self._json(
                502,
                {
                    "error": {
                        "message": f"Codex broker failed: {type(exc).__name__}: {exc}",
                        "type": "upstream_error",
                    }
                },
            )

    def log_message(self, format: str, *args: Any) -> None:
        """Log only method/path/status metadata, never bodies or headers."""
        print(f"codex-broker {self.address_string()} {format % args}", flush=True)


def main() -> None:
    """Serve until Kubernetes terminates the sidecar."""
    if not TOKEN:
        raise SystemExit("HERMES_IMAGE_BROKER_KEY is required")
    ThreadingHTTPServer((HOST, PORT), Handler).serve_forever()


if __name__ == "__main__":
    main()