#!/usr/bin/env python3
"""Lease-aware FLUX image generation for the shared titan-24 GPU."""

from __future__ import annotations

import base64
import gc
import io
import json
import os
import ssl
import threading
from datetime import UTC, datetime
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any
from urllib.error import HTTPError
from urllib.request import Request, urlopen


HOST = os.environ.get("HERMES_LOCAL_IMAGE_HOST", "0.0.0.0")
PORT = int(os.environ.get("HERMES_LOCAL_IMAGE_PORT", "9004"))
MODEL_ID = os.environ.get(
    "HERMES_LOCAL_IMAGE_MODEL", "black-forest-labs/FLUX.2-klein-4B"
)
MODEL_REVISION = os.environ.get(
    "HERMES_LOCAL_IMAGE_REVISION", "e7b7dc27f91deacad38e78976d1f2b499d76a294"
)
MODEL_CACHE = os.environ.get("HF_HOME", "/models/huggingface")
LEASE_NAMESPACE = os.environ.get("LEASE_NAMESPACE", "hermes")
LEASE_NAME = os.environ.get("LEASE_NAME", "titan-24-gpu-owner")
LEASE_IDLE_OWNER = os.environ.get("LEASE_IDLE_OWNER", "hermes")
LEASE_IMAGE_OWNER = os.environ.get("LEASE_IMAGE_OWNER", "hermes-image")
MAX_BODY_BYTES = int(os.environ.get("HERMES_LOCAL_IMAGE_MAX_BODY", str(96 << 20)))
QUEUE_TIMEOUT_SECONDS = float(os.environ.get("HERMES_LOCAL_IMAGE_QUEUE_TIMEOUT", "1200"))
API_HOST = os.environ.get("KUBERNETES_SERVICE_HOST", "kubernetes.default.svc")
API_PORT = os.environ.get("KUBERNETES_SERVICE_PORT_HTTPS", "443")
TOKEN_PATH = Path("/var/run/secrets/kubernetes.io/serviceaccount/token")
CA_PATH = Path("/var/run/secrets/kubernetes.io/serviceaccount/ca.crt")
LEASE_URL = (
    f"https://{API_HOST}:{API_PORT}/apis/coordination.k8s.io/v1/"
    f"namespaces/{LEASE_NAMESPACE}/leases/{LEASE_NAME}"
)
ASPECTS = {
    "square": (1024, 1024),
    "landscape": (768, 1280),
    "portrait": (1280, 768),
}

_generation_lock = threading.Lock()
_state_lock = threading.Lock()
_state: dict[str, Any] = {
    "busy": False,
    "phase": "idle",
    "last_error": None,
    "last_model": None,
}


def _set_state(**values: Any) -> None:
    """Update health state without exposing prompts or image content."""
    with _state_lock:
        _state.update(values)


def _state_snapshot() -> dict[str, Any]:
    """Return a stable copy for health responses."""
    with _state_lock:
        return dict(_state)


def _kube_request(method: str, body: dict[str, Any] | None = None) -> dict[str, Any]:
    """Call the Lease API using the pod service account."""
    token = TOKEN_PATH.read_text(encoding="utf-8").strip()
    raw = None if body is None else json.dumps(body).encode("utf-8")
    headers = {"Authorization": f"Bearer {token}"}
    if raw is not None:
        headers["Content-Type"] = "application/merge-patch+json"
    request = Request(LEASE_URL, data=raw, headers=headers, method=method)
    context = ssl.create_default_context(cafile=str(CA_PATH))
    with urlopen(request, timeout=5, context=context) as response:
        return json.load(response)


def _lease_owner() -> str:
    """Read the current GPU owner, failing closed on API errors."""
    payload = _kube_request("GET")
    return str((payload.get("spec") or {}).get("holderIdentity") or "unavailable")


def _change_owner(expected: str, desired: str) -> bool:
    """Change owner only if the resource version and holder still match."""
    current = _kube_request("GET")
    if str((current.get("spec") or {}).get("holderIdentity")) != expected:
        return False
    resource_version = str((current.get("metadata") or {}).get("resourceVersion") or "")
    if not resource_version:
        return False
    patch = {
        "metadata": {"resourceVersion": resource_version},
        "spec": {
            "holderIdentity": desired,
            "renewTime": datetime.now(UTC).isoformat().replace("+00:00", "Z"),
        },
    }
    try:
        _kube_request("PATCH", patch)
    except HTTPError as exc:
        if exc.code == 409:
            return False
        raise
    return True


def _recover_stale_image_lease() -> None:
    """Release this deployment's lease after a prior renderer crash."""
    try:
        if _lease_owner() == LEASE_IMAGE_OWNER:
            _change_owner(LEASE_IMAGE_OWNER, LEASE_IDLE_OWNER)
    except Exception as exc:
        _set_state(last_error=f"lease recovery failed: {type(exc).__name__}: {exc}")


def _decode_image(value: str):
    """Decode a bounded image data URL for FLUX editing."""
    from PIL import Image

    if not value.startswith("data:image/") or ";base64," not in value:
        raise ValueError("local reference images must be data URLs")
    raw = base64.b64decode(value.split(",", 1)[1], validate=True)
    if not raw or len(raw) > 25 << 20:
        raise ValueError("reference image must be between 1 byte and 25 MiB")
    return Image.open(io.BytesIO(raw)).convert("RGB")


def _render(payload: dict[str, Any]) -> dict[str, Any]:
    """Load FLUX for one request, render, then fully release its GPU memory."""
    import torch
    from diffusers import Flux2KleinPipeline

    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 ASPECTS:
        raise ValueError("aspect_ratio must be landscape, square, or portrait")
    height, width = ASPECTS[aspect]
    inputs = []
    if payload.get("image_url"):
        inputs.append(_decode_image(str(payload["image_url"])))
    for item in payload.get("reference_image_urls") or []:
        inputs.append(_decode_image(str(item)))
    inputs = inputs[:4]

    _set_state(phase="loading-flux", last_model=MODEL_ID)
    pipe = None
    try:
        pipe = Flux2KleinPipeline.from_pretrained(
            MODEL_ID,
            revision=MODEL_REVISION,
            torch_dtype=torch.bfloat16,
            cache_dir=MODEL_CACHE,
        )
        pipe.enable_model_cpu_offload()
        _set_state(phase="rendering")
        arguments: dict[str, Any] = {
            "prompt": prompt,
            "height": height,
            "width": width,
            "guidance_scale": 1.0,
            "num_inference_steps": 4,
        }
        if inputs:
            arguments["image"] = inputs
        image = pipe(**arguments).images[0]
        output = io.BytesIO()
        image.save(output, format="PNG", optimize=True)
        return {
            "success": True,
            "image_b64": base64.b64encode(output.getvalue()).decode("ascii"),
            "mime_type": "image/png",
            "model": "flux-2-klein-4b-local",
            "quality": "local-high",
            "size": f"{width}x{height}",
            "input_image_count": len(inputs),
            "route": "local",
        }
    finally:
        del pipe
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()


def _generate(payload: dict[str, Any]) -> dict[str, Any]:
    """Serialize renders and perform an exclusive lease-backed GPU handoff."""
    acquired = _generation_lock.acquire(timeout=QUEUE_TIMEOUT_SECONDS)
    if not acquired:
        raise TimeoutError("local image queue is busy; retry shortly")
    claimed = False
    _set_state(busy=True, phase="claiming-gpu", last_error=None)
    try:
        claimed = _change_owner(LEASE_IDLE_OWNER, LEASE_IMAGE_OWNER)
        if not claimed:
            owner = _lease_owner()
            raise RuntimeError(f"local image GPU unavailable while titan-24 owner is {owner}")
        result = _render(payload)
        _set_state(phase="releasing-flux")
        if _lease_owner() == LEASE_IMAGE_OWNER:
            if not _change_owner(LEASE_IMAGE_OWNER, LEASE_IDLE_OWNER):
                raise RuntimeError("GPU ownership changed while releasing the image lane")
        return result
    except Exception as exc:
        _set_state(last_error=f"{type(exc).__name__}: {exc}")
        raise
    finally:
        # If Wolf has moved the lease to wolf-draining, never overwrite that
        # transition. Ariadne waits for busy=false before using the GPU.
        if claimed:
            try:
                if _lease_owner() == LEASE_IMAGE_OWNER:
                    _change_owner(LEASE_IMAGE_OWNER, LEASE_IDLE_OWNER)
            except Exception as exc:
                _set_state(last_error=f"restore failed: {type(exc).__name__}: {exc}")
        _set_state(busy=False, phase="idle")
        _generation_lock.release()


class Handler(BaseHTTPRequestHandler):
    """Internal JSON interface used only by the owner image broker."""

    server_version = "HermesLocalImage/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 do_GET(self) -> None:  # noqa: N802
        if self.path != "/health":
            self._json(404, {"success": False, "error": "not found"})
            return
        state = _state_snapshot()
        try:
            owner = _lease_owner()
        except Exception:
            owner = "unavailable"
        self._json(
            200,
            {
                "success": True,
                "available": owner == LEASE_IDLE_OWNER and not state["busy"],
                "gpu_owner": owner,
                "model": "flux-2-klein-4b-local",
                **state,
            },
        )

    def do_POST(self) -> None:  # noqa: N802
        if self.path != "/v1/images/generations":
            self._json(404, {"success": False, "error": "not found"})
            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, 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:
            self._json(
                503,
                {"success": False, "error": f"local image generation failed: {type(exc).__name__}: {exc}"},
            )

    def log_message(self, format_string: str, *args: Any) -> None:
        print(f"local-image {self.address_string()} {format_string % args}", flush=True)


if __name__ == "__main__":
    _recover_stale_image_lease()
    ThreadingHTTPServer((HOST, PORT), Handler).serve_forever()