#!/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 subprocess
import sys
import threading
import time
from contextlib import redirect_stdout
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")
OFFLOAD_MODE = os.environ.get("HERMES_LOCAL_IMAGE_OFFLOAD_MODE", "sequential")
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"))
RENDER_TIMEOUT_SECONDS = float(
    os.environ.get("HERMES_LOCAL_IMAGE_RENDER_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_in_process(payload: dict[str, Any]) -> dict[str, Any]:
    """Load and run FLUX inside the disposable CUDA worker process."""
    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,
        )
        if OFFLOAD_MODE == "sequential":
            # The shared 10 GiB card keeps the desktop and idle Wolf daemon
            # resident. Layer-level offload preserves 1024px output quality
            # without requiring those low-memory services to be torn down.
            pipe.enable_sequential_cpu_offload()
        elif OFFLOAD_MODE == "model":
            pipe.enable_model_cpu_offload()
        else:
            raise ValueError("HERMES_LOCAL_IMAGE_OFFLOAD_MODE must be sequential or model")
        _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 _render(payload: dict[str, Any]) -> dict[str, Any]:
    """Render in a child process so process exit releases every CUDA allocation."""
    _set_state(phase="render-worker", last_model=MODEL_ID)
    completed = subprocess.run(
        [sys.executable, str(Path(__file__).resolve()), "--render-worker"],
        input=json.dumps(payload, separators=(",", ":")).encode("utf-8"),
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        timeout=RENDER_TIMEOUT_SECONDS,
        check=False,
    )
    if completed.returncode != 0:
        detail = completed.stderr.decode("utf-8", errors="replace").strip()
        if len(detail) > 4000:
            detail = detail[-4000:]
        raise RuntimeError(detail or f"FLUX worker exited {completed.returncode}")
    try:
        result = json.loads(completed.stdout)
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise RuntimeError("FLUX worker returned an invalid response") from exc
    if not isinstance(result, dict) or not result.get("success"):
        raise RuntimeError("FLUX worker returned an unsuccessful response")
    return result


def _run_render_worker() -> int:
    """Read one bounded request and emit one JSON response for the parent."""
    raw = sys.stdin.buffer.read(MAX_BODY_BYTES + 1)
    if not raw or len(raw) > MAX_BODY_BYTES:
        print("invalid render worker request size", file=sys.stderr)
        return 2
    try:
        payload = json.loads(raw)
        if not isinstance(payload, dict):
            raise ValueError("JSON object required")
        # Keep library progress output off the machine-readable stdout channel.
        with redirect_stdout(sys.stderr):
            result = _render_in_process(payload)
        sys.stdout.write(json.dumps(result, separators=(",", ":")))
        sys.stdout.flush()
        return 0
    except Exception as exc:
        print(f"{type(exc).__name__}: {exc}", file=sys.stderr, flush=True)
        return 1


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__":
    if sys.argv[1:] == ["--render-worker"]:
        raise SystemExit(_run_render_worker())
    _recover_stale_image_lease()
    ThreadingHTTPServer((HOST, PORT), Handler).serve_forever()