#!/usr/bin/env python3
"""Private FLUX.2 Klein Distilled worker used only during a GPU hot swap."""

from __future__ import annotations

import gc
import json
import os
import signal
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

HOST = os.environ.get("WORKER_HOST", "0.0.0.0")
PORT = int(os.environ.get("WORKER_PORT", "8086"))
TOKEN = os.environ.get("WORKER_TOKEN", "").strip()
MODEL_DIR = os.environ.get("FLUX_MODEL_DIR", "/models/FLUX.2-klein-4B")
OUTPUT_DIR = Path(os.environ.get("IMAGE_DIR", "/data/images")).resolve()
PIPE = None
LOAD_SECONDS = 0.0

if len(TOKEN) < 32:
    raise RuntimeError("WORKER_TOKEN is missing or too short")

# The container is intentionally disposable. After the router has received a
# completed response and saved image, an immediate process exit releases the
# CUDA context much faster than Python/PyTorch interpreter teardown.
signal.signal(signal.SIGTERM, lambda *_: os._exit(0))


def load_pipeline() -> None:
    global PIPE, LOAD_SECONDS
    if PIPE is not None:
        return
    import torch
    from diffusers import DiffusionPipeline
    started = time.monotonic()
    PIPE = DiffusionPipeline.from_pretrained(
        MODEL_DIR, torch_dtype=torch.bfloat16)
    # The full pipeline leaves too little activation headroom on a 16 GiB
    # RTX 5080. Model CPU offload keeps each active component on CUDA while
    # parking inactive components in system RAM between the four steps.
    PIPE.enable_model_cpu_offload()
    if hasattr(PIPE, "enable_vae_slicing"):
        PIPE.enable_vae_slicing()
    if hasattr(PIPE, "enable_vae_tiling"):
        PIPE.enable_vae_tiling()
    LOAD_SECONDS = time.monotonic() - started


def generate(data: dict) -> dict:
    import torch
    prompt = data.get("prompt")
    filename = data.get("filename")
    if not isinstance(prompt, str) or not prompt.strip() or len(prompt) > 8000:
        raise ValueError("invalid prompt")
    if (not isinstance(filename, str) or Path(filename).name != filename
            or not filename.endswith(".png")):
        raise ValueError("invalid filename")
    width, height = int(data.get("width", 1024)), int(data.get("height", 1024))
    if (width, height) not in {(1024, 1024), (1536, 1024), (1024, 1536),
                               (1920, 1088), (1088, 1920)}:
        raise ValueError("unsupported image size")
    steps = int(data.get("steps", 4))
    guidance = float(data.get("guidance", 1.0))
    if steps != 4 or guidance != 1.0:
        raise ValueError("distilled FLUX.2 Klein requires steps=4 and guidance=1.0")
    seed = data.get("seed")
    generator = None if seed is None else torch.Generator(device="cuda").manual_seed(int(seed))
    load_pipeline()
    started = time.monotonic()
    image = PIPE(prompt=prompt, height=height, width=width,
                 num_inference_steps=4, guidance_scale=1.0,
                 generator=generator).images[0]
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    output = OUTPUT_DIR / filename
    image.save(output)
    return {"status": "ok", "filename": filename,
            "seconds": round(time.monotonic() - started, 3),
            "load_seconds": round(LOAD_SECONDS, 3)}


class Handler(BaseHTTPRequestHandler):
    def log_message(self, fmt: str, *args: object) -> None:
        # Never log request bodies/prompts.
        print(f"[flux-worker] {self.client_address[0]} {fmt % args}", flush=True)

    def reply(self, status: int, payload: dict) -> None:
        body = json.dumps(payload, separators=(",", ":")).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self) -> None:  # noqa: N802
        if self.path == "/health":
            self.reply(200, {"status": "ok", "model_loaded": PIPE is not None})
        else:
            self.reply(404, {"error": "not found"})

    def do_POST(self) -> None:  # noqa: N802
        if self.headers.get("Authorization", "") != f"Bearer {TOKEN}":
            self.reply(401, {"error": "unauthorized"})
            return
        if self.path != "/generate":
            self.reply(404, {"error": "not found"})
            return
        try:
            length = int(self.headers.get("Content-Length", "0"))
            if length < 2 or length > 16384:
                raise ValueError("invalid request size")
            self.reply(200, generate(json.loads(self.rfile.read(length))))
        except Exception as exc:
            print(f"[flux-worker] generation failed: "
                  f"{type(exc).__name__}: {str(exc)[:1000]}", flush=True)
            self.reply(400, {"status": "error", "message": str(exc)})


try:
    ThreadingHTTPServer((HOST, PORT), Handler).serve_forever()
finally:
    if PIPE is not None:
        del PIPE
    gc.collect()