#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10,<3.11"
# dependencies = [
#   "torch==2.9.1+cu128",
#   "torchvision==0.24.1+cu128",
#   "torchaudio==2.9.1+cu128",
#   "diffusers @ git+https://github.com/huggingface/diffusers.git",
#   "transformers>=4.57.0",
#   "accelerate>=1.0.0",
#   "safetensors>=0.4.0",
#   "huggingface_hub>=0.25.0",
#   "numpy>=1.26.0",
#   "pillow>=10.0.0",
#   "fastapi>=0.115.0",
#   "uvicorn>=0.29.0",
# ]
# [tool.uv]
# extra-index-url = ["https://download.pytorch.org/whl/cu128"]
# ///

import argparse
import base64
import binascii
import io
import os
import random
import threading
import time
from typing import Any

os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")

import torch
from diffusers import Flux2KleinPipeline
from fastapi import FastAPI, HTTPException
from fastapi.responses import HTMLResponse, Response
from pydantic import BaseModel, Field
from PIL import Image


def _resolve_model_size() -> str:
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--model", choices=["4b", "9b"], default="9b")
    args, _ = parser.parse_known_args()
    return args.model


def _checkpoint_for(model_size: str) -> str:
    if model_size == "4b":
        return "black-forest-labs/FLUX.2-klein-4B"
    return "black-forest-labs/FLUX.2-klein-9B"


MODEL_SIZE = _resolve_model_size()
CHECKPOINT = _checkpoint_for(MODEL_SIZE)
USE_COMPILE = os.environ.get("FLUX2_COMPILE", "1") != "0"
WARMUP_RUNS = int(os.environ.get("FLUX2_WARMUP_RUNS", "2"))
WARMUP_STEPS = int(os.environ.get("FLUX2_WARMUP_STEPS", "4"))
WARMUP_PROMPT = "A high-resolution photo of a red fox in natural daylight"
MAX_SEED = (2**31) - 1
MAX_PIXELS = 4 * 1024 * 1024

DEFAULT_STEPS = 4
DEFAULT_GUIDANCE = 1.0

app = FastAPI(title="FLUX2 Klein Private Server")
PIPE: Flux2KleinPipeline | None = None
PIPE_LOCK = threading.Lock()


class GenerateRequest(BaseModel):
    prompt: str = Field(..., min_length=1)
    width: int = 1024
    height: int = 1024
    num_inference_steps: int = DEFAULT_STEPS
    guidance_scale: float = DEFAULT_GUIDANCE
    seed: int = 0
    randomize_seed: bool = False
    image_data_url: str | None = None


class CompatRequest(BaseModel):
    model: str | None = None
    prompt: str = Field(..., min_length=1)
    imageDataUrl: str | None = None
    size: str | None = None
    width: int | None = None
    height: int | None = None
    num_inference_steps: int | None = None
    guidance_scale: float | None = None
    seed: int | None = None
    randomize_seed: bool = False
    response_format: str = "b64_json"
    n: int = 1
    negative_prompt: str | None = None
    negativePrompt: str | None = None


def _ensure_cuda() -> None:
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for this server")


def _compile_pipe(pipe: Flux2KleinPipeline) -> bool:
    if not USE_COMPILE:
        return False
    print("Compiling transformer and VAE decoder...")
    start = time.perf_counter()
    pipe.transformer = torch.compile(
        pipe.transformer,
        mode="max-autotune",
        dynamic=False,
        fullgraph=True,
    )
    pipe.vae.decode = torch.compile(
        pipe.vae.decode,
        mode="max-autotune",
        dynamic=False,
        fullgraph=True,
    )
    print(f"Compile finished in {time.perf_counter() - start:.2f}s")
    return True


def _warmup_pipe(pipe: Flux2KleinPipeline) -> None:
    if WARMUP_RUNS <= 0:
        return
    print(f"Warmup started ({WARMUP_RUNS} runs)...")
    start = time.perf_counter()
    for run_idx in range(WARMUP_RUNS):
        generator = torch.Generator(device="cuda").manual_seed(run_idx)
        pipe(
            prompt=WARMUP_PROMPT,
            height=1024,
            width=1024,
            num_inference_steps=WARMUP_STEPS,
            guidance_scale=1.0,
            generator=generator,
        ).images[0]
    torch.cuda.synchronize()
    print(f"Warmup finished in {time.perf_counter() - start:.2f}s")


def _load_pipe(ckpt: str) -> Flux2KleinPipeline:
    print(f"Loading checkpoint: {ckpt}")
    pipe = Flux2KleinPipeline.from_pretrained(ckpt, torch_dtype=torch.bfloat16)
    pipe = pipe.to("cuda")
    pipe.transformer.fuse_qkv_projections()
    pipe.vae.fuse_qkv_projections()
    pipe.vae.to(memory_format=torch.channels_last)
    pipe.transformer.set_attention_backend("_native_flash")
    if _compile_pipe(pipe):
        _warmup_pipe(pipe)
    return pipe


def _normalize_dims(width: int, height: int) -> tuple[int, int]:
    if width <= 0 or height <= 0:
        raise HTTPException(status_code=422, detail="width and height must be > 0")
    if width % 16 != 0 or height % 16 != 0:
        raise HTTPException(
            status_code=422, detail="width and height must be multiples of 16"
        )
    if width * height > MAX_PIXELS:
        raise HTTPException(status_code=422, detail="max resolution is 4 megapixels")
    return width, height


def _parse_size(
    size: str | None, width: int | None, height: int | None
) -> tuple[int, int]:
    if width and height:
        return _normalize_dims(width, height)
    if not size:
        return _normalize_dims(1024, 1024)
    parts = size.lower().split("x")
    if len(parts) != 2:
        raise HTTPException(status_code=422, detail="size must look like 1024x1024")
    try:
        parsed_w = int(parts[0])
        parsed_h = int(parts[1])
    except ValueError as exc:
        raise HTTPException(
            status_code=422, detail="size values must be integers"
        ) from exc
    return _normalize_dims(parsed_w, parsed_h)


def _decode_image_data_url(image_data_url: str | None) -> Image.Image | None:
    if not image_data_url:
        return None
    if "," not in image_data_url:
        raise HTTPException(status_code=422, detail="imageDataUrl must be a data URL")
    _, payload = image_data_url.split(",", 1)
    try:
        raw = base64.b64decode(payload)
    except (ValueError, binascii.Error) as exc:
        raise HTTPException(
            status_code=422, detail="invalid base64 in imageDataUrl"
        ) from exc
    with Image.open(io.BytesIO(raw)) as image:
        return image.convert("RGB")


def _generate_png_bytes(
    prompt: str,
    width: int,
    height: int,
    num_inference_steps: int,
    guidance_scale: float,
    seed: int,
    image: Image.Image | None,
) -> tuple[bytes, int]:
    if PIPE is None:
        raise HTTPException(status_code=503, detail="model not loaded yet")

    generator = torch.Generator(device="cuda").manual_seed(seed)
    kwargs: dict[str, Any] = {
        "prompt": prompt,
        "height": height,
        "width": width,
        "num_inference_steps": num_inference_steps,
        "guidance_scale": guidance_scale,
        "generator": generator,
    }

    with PIPE_LOCK:
        try:
            if image is not None:
                try:
                    output = PIPE(image=image, **kwargs)
                except TypeError:
                    output = PIPE(images=[image], **kwargs)
            else:
                output = PIPE(**kwargs)
        except Exception as exc:
            raise HTTPException(
                status_code=500, detail=f"generation failed: {exc}"
            ) from exc

    result_image = output.images[0]
    buf = io.BytesIO()
    result_image.save(buf, format="PNG")
    return buf.getvalue(), seed


@app.on_event("startup")
def _startup() -> None:
    global PIPE
    _ensure_cuda()
    PIPE = _load_pipe(CHECKPOINT)
    print("Server ready")


@app.get("/", response_class=HTMLResponse)
def index() -> str:
    return """<!doctype html>
<html lang=\"en\">
<head>
  <meta charset=\"utf-8\" />
  <meta name=\"viewport\" content=\"width=device-width, initial-scale=1\" />
  <title>FLUX.2 Klein Playground</title>
  <style>
    :root { --bg:#f4f0e8; --card:#fff9ee; --text:#1f1f1f; --accent:#125b50; --muted:#6b6b6b; }
    * { box-sizing: border-box; }
    body { margin:0; font-family: ui-sans-serif, -apple-system, Segoe UI, sans-serif; background: radial-gradient(circle at top left,#fff6df,var(--bg)); color:var(--text); }
    .wrap { max-width: 980px; margin: 1.5rem auto; padding: 1rem; }
    .card { background: var(--card); border: 1px solid #eadfc8; border-radius: 14px; padding: 1rem; box-shadow: 0 8px 20px rgba(0,0,0,.06); }
    h1 { margin: 0 0 .8rem 0; font-size: 1.4rem; }
    .grid { display:grid; gap:.8rem; grid-template-columns: 1fr 1fr; }
    .full { grid-column: 1 / -1; }
    label { display:block; font-weight:600; margin-bottom:.3rem; }
    input, textarea, button { width:100%; border-radius:10px; border:1px solid #d8ccb7; padding:.65rem; font-size:.95rem; }
    textarea { min-height: 110px; resize: vertical; }
    button { background: var(--accent); color:white; border:none; font-weight:700; cursor:pointer; }
    button:disabled { opacity:.6; cursor:default; }
    .muted { color: var(--muted); font-size:.88rem; }
    .out img { width:100%; border-radius:12px; border:1px solid #e4d8c3; background:white; }
    @media (max-width: 760px) { .grid { grid-template-columns: 1fr; } }
  </style>
</head>
<body>
  <div class=\"wrap\">
    <div class=\"card\">
      <h1>FLUX.2 Klein Playground</h1>
      <p class=\"muted\">Private single-user UI at <code>/</code>. Distilled defaults: 4 steps, guidance 1.0.</p>
      <form id=\"form\" class=\"grid\">
        <div class=\"full\">
          <label for=\"prompt\">Prompt</label>
          <textarea id=\"prompt\" required>A cinematic photo portrait, natural skin texture, realistic lighting, high detail.</textarea>
        </div>
        <div>
          <label for=\"width\">Width (multiple of 16)</label>
          <input id=\"width\" type=\"number\" value=\"1024\" step=\"16\" min=\"16\" />
        </div>
        <div>
          <label for=\"height\">Height (multiple of 16)</label>
          <input id=\"height\" type=\"number\" value=\"1024\" step=\"16\" min=\"16\" />
        </div>
        <div>
          <label for=\"steps\">Inference steps</label>
          <input id=\"steps\" type=\"number\" value=\"4\" min=\"1\" max=\"100\" />
        </div>
        <div>
          <label for=\"guidance\">Guidance scale</label>
          <input id=\"guidance\" type=\"number\" value=\"1.0\" step=\"0.1\" />
        </div>
        <div>
          <label for=\"seed\">Seed</label>
          <input id=\"seed\" type=\"number\" value=\"0\" min=\"0\" />
        </div>
        <div>
          <label for=\"edit\">Edit image (optional)</label>
          <input id=\"edit\" type=\"file\" accept=\"image/*\" />
        </div>
        <div class=\"full\">
          <button id=\"submit\" type=\"submit\">Generate</button>
          <p id=\"status\" class=\"muted\">Idle</p>
        </div>
      </form>
    </div>

    <div class=\"card out\" style=\"margin-top:1rem\">
      <h1>Output</h1>
      <img id=\"preview\" alt=\"Generated image preview\" />
      <p class=\"muted\">Tip: right-click image to save.</p>
    </div>
  </div>

  <script>
    const form = document.getElementById('form');
    const statusEl = document.getElementById('status');
    const preview = document.getElementById('preview');
    const button = document.getElementById('submit');

    async function fileToDataUrl(file) {
      return await new Promise((resolve, reject) => {
        const reader = new FileReader();
        reader.onload = () => resolve(reader.result);
        reader.onerror = reject;
        reader.readAsDataURL(file);
      });
    }

    form.addEventListener('submit', async (event) => {
      event.preventDefault();
      button.disabled = true;
      statusEl.textContent = 'Generating...';
      try {
        const prompt = document.getElementById('prompt').value;
        const width = Number(document.getElementById('width').value);
        const height = Number(document.getElementById('height').value);
        const steps = Number(document.getElementById('steps').value);
        const guidance = Number(document.getElementById('guidance').value);
        const seed = Number(document.getElementById('seed').value);
        const file = document.getElementById('edit').files[0];
        const imageDataUrl = file ? await fileToDataUrl(file) : null;

        const payload = {
          model: 'black-forest-labs/FLUX.2-klein-9B',
          prompt,
          size: `${width}x${height}`,
          num_inference_steps: steps,
          guidance_scale: guidance,
          seed,
          response_format: 'b64_json',
          imageDataUrl,
        };

        const resp = await fetch('/v1/images/generations', {
          method: 'POST',
          headers: { 'Content-Type': 'application/json' },
          body: JSON.stringify(payload),
        });

        if (!resp.ok) {
          const text = await resp.text();
          throw new Error(text || `HTTP ${resp.status}`);
        }

        const data = await resp.json();
        const b64 = data?.data?.[0]?.b64_json;
        if (!b64) throw new Error('No image returned');
        preview.src = `data:image/png;base64,${b64}`;
        statusEl.textContent = `Done. Seed: ${data.seed ?? seed}`;
      } catch (error) {
        statusEl.textContent = `Error: ${error.message}`;
      } finally {
        button.disabled = false;
      }
    });
  </script>
</body>
</html>
"""


@app.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok", "model": CHECKPOINT}


@app.post("/generate")
def generate(req: GenerateRequest) -> Response:
    width, height = _normalize_dims(req.width, req.height)
    seed = req.seed if not req.randomize_seed else random.randint(0, MAX_SEED)
    image = _decode_image_data_url(req.image_data_url)
    png_bytes, used_seed = _generate_png_bytes(
        prompt=req.prompt,
        width=width,
        height=height,
        num_inference_steps=req.num_inference_steps,
        guidance_scale=req.guidance_scale,
        seed=seed,
        image=image,
    )
    return Response(
        content=png_bytes,
        media_type="image/png",
        headers={"X-Seed": str(used_seed)},
    )


@app.post("/v1/images/generations")
def compat_generate(req: CompatRequest) -> dict[str, Any]:
    width, height = _parse_size(req.size, req.width, req.height)
    seed_value = req.seed if req.seed is not None else 0
    seed = seed_value if not req.randomize_seed else random.randint(0, MAX_SEED)
    image = _decode_image_data_url(req.imageDataUrl)

    png_bytes, used_seed = _generate_png_bytes(
        prompt=req.prompt,
        width=width,
        height=height,
        num_inference_steps=req.num_inference_steps or DEFAULT_STEPS,
        guidance_scale=req.guidance_scale or DEFAULT_GUIDANCE,
        seed=seed,
        image=image,
    )
    b64_png = base64.b64encode(png_bytes).decode("ascii")
    return {
        "created": int(time.time()),
        "model": CHECKPOINT,
        "seed": used_seed,
        "data": [{"b64_json": b64_png}],
    }


if __name__ == "__main__":
    import uvicorn

    parser = argparse.ArgumentParser(description="Private FLUX.2 Klein FastAPI server")
    parser.add_argument("--model", choices=["4b", "9b"], default=MODEL_SIZE)
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=6006)
    args = parser.parse_args()

    MODEL_SIZE = args.model
    CHECKPOINT = _checkpoint_for(MODEL_SIZE)
    uvicorn.run(app, host=args.host, port=args.port)