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

Ported from an earlier standalone server, run there with MOSS_LORA=dmd2.
Differences from the original:
- weights (and the LoRA, which ships in the same repo) load from /weights/kromoss, an
  m64 artifact mount; port from M64_PORT;
- only the DMD2 variant: size, steps, guidance and shift are baked constants (the capsule's
  [params]); a request carries only `prompt` and `seed` (the m64 front rejects
  anything else with a 400);
- progress is reported as JSON lines appended to M64_STATUS.

Kept exactly from the original (and 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()