#!/usr/bin/env python3
"""Private GPT Image broker for isolated Hermes chat tenants."""

from __future__ import annotations

import base64
import hmac
import importlib.util
import json
import os
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any


HOST = os.environ.get("HERMES_IMAGE_BROKER_LISTEN_HOST", "0.0.0.0")
# Kubernetes reserves <SERVICE>_PORT for service-link discovery and injects a
# tcp://... value.  Keep listener configuration out of that namespace.
PORT = int(os.environ.get("HERMES_IMAGE_BROKER_LISTEN_PORT", "9002"))
TOKEN = os.environ.get("HERMES_IMAGE_BROKER_KEY", "").strip()
PROVIDER_PATH = Path(
    os.environ.get(
        "HERMES_IMAGE_PROVIDER_PATH",
        "/opt/hermes/plugins/image_gen/openai-codex/__init__.py",
    )
)
DEFAULT_MODEL = os.environ.get(
    "HERMES_IMAGE_BROKER_DEFAULT_MODEL", "gpt-image-2-high"
)
MAX_BODY_BYTES = int(os.environ.get("HERMES_IMAGE_BROKER_MAX_BODY", str(96 << 20)))
QUEUE_TIMEOUT_SECONDS = float(
    os.environ.get("HERMES_IMAGE_BROKER_QUEUE_TIMEOUT", "900")
)
ALLOWED_MODELS = {
    "gpt-image-2-low",
    "gpt-image-2-medium",
    "gpt-image-2-high",
}
ALLOWED_ASPECTS = {"landscape", "square", "portrait"}
_generation_lock = threading.Lock()


def _load_provider() -> Any:
    """Load the pinned Hermes provider without duplicating its API client."""
    spec = importlib.util.spec_from_file_location(
        "atlas_openai_codex_image_provider", PROVIDER_PATH
    )
    if spec is None or spec.loader is None:
        raise RuntimeError(f"cannot load image provider from {PROVIDER_PATH}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module.OpenAICodexImageGenProvider()


_PROVIDER: Any | None = None


def _provider() -> Any:
    """Load the provider once after the sidecar environment is ready."""
    global _PROVIDER
    if _PROVIDER is None:
        _PROVIDER = _load_provider()
    return _PROVIDER


def _authorized(header: str | None) -> bool:
    """Compare the bearer token without leaking timing information."""
    if not TOKEN or not header or not header.startswith("Bearer "):
        return False
    return hmac.compare_digest(header[7:].strip(), TOKEN)


def _generate(payload: dict[str, Any]) -> dict[str, Any]:
    """Generate once, return bytes, and remove the broker-side cache copy."""
    prompt = str(payload.get("prompt") or "").strip()
    if not prompt:
        raise ValueError("prompt is required")
    aspect = str(payload.get("aspect_ratio") or "landscape").lower()
    if aspect not in ALLOWED_ASPECTS:
        raise ValueError("aspect_ratio must be landscape, square, or portrait")
    model = str(payload.get("model") or DEFAULT_MODEL)
    if model not in ALLOWED_MODELS:
        raise ValueError("unsupported image quality tier")
    image_url = payload.get("image_url")
    references = payload.get("reference_image_urls")
    if image_url is not None and not isinstance(image_url, str):
        raise ValueError("image_url must be a string")
    if references is not None and not (
        isinstance(references, list)
        and all(isinstance(item, str) for item in references)
    ):
        raise ValueError("reference_image_urls must be a list of strings")

    acquired = _generation_lock.acquire(timeout=QUEUE_TIMEOUT_SECONDS)
    if not acquired:
        raise TimeoutError("image generation queue is busy; retry shortly")
    old_model = os.environ.get("OPENAI_IMAGE_MODEL")
    try:
        os.environ["OPENAI_IMAGE_MODEL"] = model
        result = _provider().generate(
            prompt,
            aspect,
            image_url=image_url,
            reference_image_urls=references,
        )
    finally:
        if old_model is None:
            os.environ.pop("OPENAI_IMAGE_MODEL", None)
        else:
            os.environ["OPENAI_IMAGE_MODEL"] = old_model
        _generation_lock.release()

    if not result.get("success"):
        return result
    image_path = Path(str(result["image"]))
    try:
        raw = image_path.read_bytes()
    finally:
        image_path.unlink(missing_ok=True)
    response = {key: value for key, value in result.items() if key != "image"}
    response.update(
        {
            "image_b64": base64.b64encode(raw).decode("ascii"),
            "mime_type": "image/png",
        }
    )
    return response


class Handler(BaseHTTPRequestHandler):
    """Small authenticated JSON API; prompts and images are never logged."""

    server_version = "HermesImageBroker/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, {"success": False, "error": "unauthorized"})
        return False

    def do_GET(self) -> None:  # noqa: N802 - BaseHTTPRequestHandler API
        if self.path != "/health":
            self._json(404, {"success": False, "error": "not found"})
            return
        if not self._check_auth():
            return
        self._json(
            200,
            {
                "success": True,
                "provider": "openai-codex",
                "available": bool(_provider().is_available()),
                "busy": _generation_lock.locked(),
            },
        )

    def do_POST(self) -> None:  # noqa: N802 - BaseHTTPRequestHandler API
        if self.path != "/v1/images/generations":
            self._json(404, {"success": False, "error": "not found"})
            return
        if not self._check_auth():
            return
        try:
            length = int(self.headers.get("Content-Length", "0"))
        except ValueError:
            length = 0
        if length <= 0 or length > MAX_BODY_BYTES:
            self._json(413, {"success": False, "error": "invalid request size"})
            return
        try:
            payload = json.loads(self.rfile.read(length))
            if not isinstance(payload, dict):
                raise ValueError("JSON object required")
            result = _generate(payload)
            self._json(200 if result.get("success") else 502, result)
        except ValueError as exc:
            self._json(400, {"success": False, "error": str(exc)})
        except TimeoutError as exc:
            self._json(503, {"success": False, "error": str(exc)})
        except Exception as exc:  # keep provider details useful but omit prompts
            self._json(
                502,
                {
                    "success": False,
                    "error": f"image provider failed: {type(exc).__name__}: {exc}",
                },
            )

    def log_message(self, format: str, *args: Any) -> None:
        """Log only method/path/status metadata, never request bodies."""
        print(f"image-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")
    server = ThreadingHTTPServer((HOST, PORT), Handler)
    server.serve_forever()


if __name__ == "__main__":
    main()