#!/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
import time
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")
# Kubernetes reserves <SERVICE>_PORT for service-link discovery. Keep the
# listener setting out of that namespace so it is always numeric.
PORT = int(os.environ.get("HERMES_LOCAL_IMAGE_LISTEN_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"))
GPU_ACTIVITY_URL = os.environ.get(
    "HERMES_LOCAL_IMAGE_GPU_ACTIVITY_URL",
    "http://nvidia-process-exporter.monitoring.svc.cluster.local:9401/metrics",
)
GPU_ACTIVITY_NODE = os.environ.get("HERMES_LOCAL_IMAGE_GPU_ACTIVITY_NODE", "titan-24")
GPU_ACTIVE_SM_PERCENT = float(
    os.environ.get("HERMES_LOCAL_IMAGE_GPU_ACTIVE_SM_PERCENT", "8")
)
GPU_MAX_EXTERNAL_MEMORY_BYTES = int(
    os.environ.get("HERMES_LOCAL_IMAGE_GPU_MAX_EXTERNAL_MEMORY_BYTES", str(3 << 30))
)
GPU_ACTIVITY_SAMPLES = max(
    1, int(os.environ.get("HERMES_LOCAL_IMAGE_GPU_ACTIVITY_SAMPLES", "3"))
)
GPU_ACTIVITY_SAMPLE_INTERVAL = max(
    0.0, float(os.environ.get("HERMES_LOCAL_IMAGE_GPU_ACTIVITY_SAMPLE_INTERVAL", "1"))
)
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,
    "interactive_active": None,
    "external_gpu_memory_bytes": None,
    "external_gpu_sm_percent": None,
    "gpu_guard_reason": None,
}


def _prometheus_labels(raw: str) -> dict[str, str]:
    """Parse the simple quoted labels emitted by nvidia-process-exporter."""
    labels: dict[str, str] = {}
    for item in raw.split(","):
        key, separator, value = item.partition("=")
        if not separator:
            continue
        labels[key.strip()] = value.strip().strip('"')
    return labels


def _parse_gpu_activity(metrics: str) -> dict[str, Any]:
    """Summarize non-Hermes GPU processes on the configured image node."""
    memory_bytes = 0
    max_sm_percent = 0.0
    processes: set[str] = set()
    for line in metrics.splitlines():
        if not line or line.startswith("#") or "{" not in line or "}" not in line:
            continue
        metric, remainder = line.split("{", 1)
        raw_labels, separator, raw_value = remainder.partition("}")
        if not separator or metric not in {
            "nvidia_process_gpu_memory_used_bytes",
            "nvidia_process_gpu_sm_util_percent",
        }:
            continue
        labels = _prometheus_labels(raw_labels)
        if labels.get("node") != GPU_ACTIVITY_NODE or labels.get("namespace") == "hermes":
            continue
        try:
            value = float(raw_value.strip().split()[0])
        except (ValueError, IndexError):
            continue
        process = labels.get("process") or "unknown"
        processes.add(f"{labels.get('namespace', 'unknown')}/{process}")
        if metric == "nvidia_process_gpu_memory_used_bytes":
            memory_bytes += max(0, int(value))
        else:
            max_sm_percent = max(max_sm_percent, value)
    active = (
        max_sm_percent >= GPU_ACTIVE_SM_PERCENT
        or memory_bytes >= GPU_MAX_EXTERNAL_MEMORY_BYTES
    )
    reasons = []
    if max_sm_percent >= GPU_ACTIVE_SM_PERCENT:
        reasons.append(f"external SM utilization is {max_sm_percent:g}%")
    if memory_bytes >= GPU_MAX_EXTERNAL_MEMORY_BYTES:
        reasons.append(f"external GPU memory is {memory_bytes} bytes")
    return {
        "interactive_active": active,
        "external_gpu_memory_bytes": memory_bytes,
        "external_gpu_sm_percent": max_sm_percent,
        "external_gpu_processes": sorted(processes),
        "gpu_guard_reason": "; ".join(reasons) or None,
    }


def _gpu_activity_snapshot() -> dict[str, Any]:
    """Fetch live per-process GPU attribution, failing closed if unavailable."""
    request = Request(GPU_ACTIVITY_URL, headers={"Accept": "text/plain"})
    with urlopen(request, timeout=5) as response:
        metrics = response.read(2 << 20).decode("utf-8", errors="replace")
    return _parse_gpu_activity(metrics)


def _ensure_interactive_lane_idle() -> dict[str, Any]:
    """Reject FLUX while desktop or Wolf activity indicates interactive use."""
    worst: dict[str, Any] | None = None
    for sample in range(GPU_ACTIVITY_SAMPLES):
        snapshot = _gpu_activity_snapshot()
        if worst is None or (
            snapshot["external_gpu_sm_percent"], snapshot["external_gpu_memory_bytes"]
        ) > (worst["external_gpu_sm_percent"], worst["external_gpu_memory_bytes"]):
            worst = snapshot
        if snapshot["interactive_active"]:
            break
        if sample + 1 < GPU_ACTIVITY_SAMPLES:
            time.sleep(GPU_ACTIVITY_SAMPLE_INTERVAL)
    assert worst is not None
    _set_state(**{key: value for key, value in worst.items() if key != "external_gpu_processes"})
    if worst["interactive_active"]:
        raise RuntimeError(
            "titan-24 desktop/Wolf lane is active; reserve or release it from the Atlas GPU checkout"
            f" ({worst['gpu_guard_reason']})"
        )
    return worst


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:
        if _lease_owner() != LEASE_IDLE_OWNER:
            owner = _lease_owner()
            raise RuntimeError(f"local image GPU unavailable while titan-24 owner is {owner}")
        _set_state(phase="checking-interactive-lane")
        _ensure_interactive_lane_idle()
        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"
        try:
            activity = _gpu_activity_snapshot()
            _set_state(**{key: value for key, value in activity.items() if key != "external_gpu_processes"})
            state = _state_snapshot()
        except Exception as exc:
            activity = {"interactive_active": True, "gpu_guard_reason": f"GPU guard unavailable: {exc}"}
            _set_state(**activity)
            state = _state_snapshot()
        self._json(
            200,
            {
                "success": True,
                "available": (
                    owner == LEASE_IDLE_OWNER
                    and not state["busy"]
                    and not activity["interactive_active"]
                ),
                "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()