"""OpenAI-style /v1/images/generations server for SANA 1.5 1.6B (m64 capsule).

Ported from an earlier standalone SANA server. Differences from the original:
- weights load from /weights/sana (an m64 artifact mount), port from M64_PORT;
- size, steps and guidance 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.
"""

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

SERVED_MODEL_ID = "sana-1.5-1.6b"
WEIGHTS = "/weights/sana"
PORT = int(os.environ["M64_PORT"])
STATUS = os.environ["M64_STATUS"]

# Baked generation settings. Keep in sync with [params] in capsule.toml.
SIZE = "1024x1024"
NUM_INFERENCE_STEPS = 20
GUIDANCE_SCALE = 4.5
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 load_model() -> None:
    global pipe
    from diffusers import SanaPipeline

    status("loading", progress=0.0, detail="loading SANA 1.5 1.6B (bf16)")
    print("Loading SANA 1.5 1.6B...", flush=True)
    t0 = time.time()
    pipe = SanaPipeline.from_pretrained(WEIGHTS, torch_dtype=torch.bfloat16)
    pipe.to("cuda")
    pipe.text_encoder.to(torch.bfloat16)
    print(f"Model loaded in {time.time() - t0:.1f}s.", flush=True)

    status("warming", detail="one short warm-up generation")
    t0 = time.time()
    with lock:
        pipe(prompt="warm-up", width=WIDTH, height=HEIGHT, num_inference_steps=2,
             guidance_scale=GUIDANCE_SCALE, generator=torch.Generator("cuda").manual_seed(0))
    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}


@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"},
            "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:
        gen = torch.Generator("cuda").manual_seed(seed)
        images = pipe(
            prompt=req.prompt,
            width=WIDTH, height=HEIGHT,
            num_inference_steps=NUM_INFERENCE_STEPS,
            guidance_scale=GUIDANCE_SCALE,
            num_images_per_prompt=1,
            generator=gen,
        )[0]
    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()