"""OpenAI-style /v1/images/generations server for Kromoss + the 8-step DMD2 LoRA.

What makes it an m64 capsule server:
- the weights (and the LoRA, which ships in the same repo) load from /weights/kromoss,
  the read-only mount of the capsule's [[artifacts]]; nothing is downloaded here;
- it starts listening on $M64_PORT only once the model is loaded; m64 polls /health
  until it answers 200;
- size, steps, guidance and shift are baked constants, the same values as [params]
  in capsule.toml; a request carries only `prompt` and `seed` (m64's front refuses
  anything else with a 400), and every answer returns them as `params_used`;
- progress goes to $M64_STATUS as JSON lines (loading, warming, ready, error), which
  `m64 status` and the /status route show.

Kept exactly from the model card: the pipeline is the repo's own
pipeline_moss_image.py via trust_remote_code; transformer and text encoder bf16, VAE fp32;
the scheduler is rebuilt from its own config (flow-matching Euler, shift 5); the LoRA is
applied UNMERGED through forward hooks, because its delta is close to bf16's rounding step
and merging would lose most of it; the initial noise comes from a CPU generator; under
CFG 1 there is no negative branch, so uncond_skip_layers is None.
"""

import base64
import io
import json
import os
import secrets
import threading
import time
import traceback
from typing import Optional

import torch
import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from safetensors.torch import load_file

SERVED_MODEL_ID = "kromoss-dmd2"
WEIGHTS = "/weights/kromoss"
LORA_NAME = "8-steps-dmd2-lora"
LORA_PATH = os.path.join(WEIGHTS, LORA_NAME + ".safetensors")
PORT = int(os.environ["M64_PORT"])
STATUS = os.environ["M64_STATUS"]

# Baked generation settings: the README's DMD2 row. Keep in sync with [params] in capsule.toml.
SIZE = "1024x1024"
NUM_INFERENCE_STEPS = 8
GUIDANCE_SCALE = 1.0
SHIFT = 5.0
WIDTH, HEIGHT = (int(x) for x in SIZE.split("x"))
MAX_SEED = 2**63 - 1

app = FastAPI()
pipe = None
lock = threading.Lock()  # one generation at a time on the single pipeline


def status(phase: str, **extra) -> None:
    with open(STATUS, "a") as f:
        f.write(json.dumps({"phase": phase, **extra}) + "\n")


def apply_lora(transformer, path: str, strength: float = 1.0) -> int:
    """Kromoss's own apply_lora: one forward hook per adapted module, never merged."""
    state = load_file(path, device="cuda")
    modules = dict(transformer.named_modules())
    n = 0
    for key in [k for k in state if k.endswith(".lora_down.weight")]:
        name = key[len("diffusion_model."):-len(".lora_down.weight")]
        down, up = state[key].float(), state[key.replace("lora_down", "lora_up")].float()
        scale = strength * state[f"diffusion_model.{name}.alpha"].item() / down.shape[0]

        def hook(module, args, out, down=down, up=up, scale=scale):
            return out + (scale * (args[0].float() @ down.T) @ up.T).to(out.dtype)

        modules[name].register_forward_hook(hook)
        n += 1
    return n


def run(prompt: str, seed: int, steps: int = NUM_INFERENCE_STEPS):
    # CPU generator: the card generates the initial noise on CPU in float32.
    return pipe(
        prompt, height=HEIGHT, width=WIDTH, num_inference_steps=steps,
        guidance_scale=GUIDANCE_SCALE, negative_prompt="", num_images_per_prompt=1,
        generator=torch.Generator("cpu").manual_seed(seed),
        uncond_skip_layers=None,  # CFG 1: no negative branch at all
    ).images


def load_model() -> None:
    global pipe
    from diffusers import DiffusionPipeline

    status("loading", progress=0.0, detail="loading Kromoss (10 GB transformer + 17 GB Qwen3.5-9B encoder)")
    print("Loading Kromoss...", flush=True)
    t0 = time.time()
    p = DiffusionPipeline.from_pretrained(
        WEIGHTS, custom_pipeline="pipeline_moss_image", trust_remote_code=True,
        torch_dtype={"default": torch.bfloat16, "vae": torch.float32},
    )
    # The card rebuilds the scheduler from its own config before use.
    p.scheduler = type(p.scheduler).from_config(p.scheduler.config)
    shift = p.scheduler.config.get("shift") if isinstance(p.scheduler.config, dict) else getattr(p.scheduler.config, "shift", None)
    if shift != SHIFT:
        raise RuntimeError(f"scheduler shift is {shift}, expected {SHIFT}")
    status("loading", progress=0.8, detail="moving to GPU, hooking the DMD2 LoRA")
    p.to("cuda")
    hooked = apply_lora(p.transformer, LORA_PATH)
    if hooked == 0:
        raise RuntimeError("DMD2 LoRA hooked no modules")
    pipe = p
    print(f"Model loaded in {time.time() - t0:.1f}s | LoRA modules hooked: {hooked} | shift {shift}", flush=True)

    status("warming", detail="one short warm-up generation")
    t0 = time.time()
    with lock:
        run("warm-up", 0, steps=2)
    print(f"Warm-up done in {time.time() - t0:.1f}s.", flush=True)


class ImageRequest(BaseModel):
    prompt: str = Field(..., min_length=1, description="Text prompt")
    seed: Optional[int] = Field(None, ge=0, le=MAX_SEED, description="RNG seed. null = random; the one used is returned.")


def baked() -> dict:
    return {"size": SIZE, "num_inference_steps": NUM_INFERENCE_STEPS,
            "guidance_scale": GUIDANCE_SCALE, "shift": SHIFT, "lora": LORA_NAME}


@app.get("/health")
def health():
    if pipe is None:
        raise HTTPException(503, "Model not loaded")
    return {"status": "ok", "model": SERVED_MODEL_ID}


@app.get("/v1/config")
def config():
    return {
        "model": SERVED_MODEL_ID,
        "params": baked(),
        "request_fields": {
            "prompt": {"required": True, "description": "Text prompt (Qwen3.5-9B encoder, up to 4096 tokens)"},
            "seed": {"required": False, "default": None, "description": "RNG seed; null = random"},
        },
    }


@app.post("/v1/images/generations")
def generate(req: ImageRequest):
    if pipe is None:
        raise HTTPException(503, "Model not loaded")
    seed = req.seed if req.seed is not None else secrets.randbelow(2**31)
    t0 = time.time()
    with lock:
        images = run(req.prompt, seed)
    elapsed = time.time() - t0
    print(f"Generated {len(images)} image(s) in {elapsed:.1f}s (seed {seed})", flush=True)
    data = []
    for img in images:
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        data.append({"b64_json": base64.b64encode(buf.getvalue()).decode()})
    return JSONResponse({
        "created": int(time.time()),
        "model": SERVED_MODEL_ID,
        "seed": seed,
        "elapsed_seconds": round(elapsed, 2),
        "data": data,
        "params_used": {**baked(), "seed": seed},
    })


def main() -> None:
    try:
        load_model()
    except Exception:
        traceback.print_exc()
        status("error", detail="model load failed; see container logs")
        raise
    status("ready")
    uvicorn.run(app, host="0.0.0.0", port=PORT)  # the gate reaches us over the capsule network


if __name__ == "__main__":
    main()