import os
import base64
from io import BytesIO

import torch
from fastapi import FastAPI
from pydantic import BaseModel, Field
from diffusers import StableDiffusionPipeline

MODEL = os.environ.get(
    "MODEL",
    os.path.join(os.path.dirname(os.path.abspath(__file__)), "realistic_vision_v4_fp16.safetensors"),
)

DEFAULT_NEGATIVE = (
    "(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, "
    "cartoon, drawing, anime:1.4), text, close up, cropped, out of frame, "
    "worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, "
    "mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn "
    "face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, "
    "extra limbs, cloned face, disfigured, gross proportions, malformed limbs, "
    "missing arms, missing legs, extra arms, extra legs, fused fingers, "
    "too many fingers, long neck"
)


def pick_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


device = pick_device()
FUEL_DTYPE = torch.float16 if device.type in ("cuda", "mps") else torch.float32
pipe = None


def load_pipeline():
    global pipe
    print(f"Loading model from {MODEL} on {device} ...", flush=True)
    pipe = StableDiffusionPipeline.from_single_file(
        MODEL,
        torch_dtype=FUEL_DTYPE,
        safety_checker=None,
        requires_safety_checker=False,
        low_cpu_mem_usage=True,
    )
    pipe = pipe.to(device)
    if device.type == "cuda":
        pipe.enable_attention_slicing("auto")
    print("Model loaded.", flush=True)


app = FastAPI(title="onlyfake realistic-vision-v4 (FP16)")


class GenRequest(BaseModel):
    prompt: str
    negative_prompt: str = Field(default=DEFAULT_NEGATIVE)
    num_inference_steps: int = Field(default=25, ge=1, le=200)
    guidance_scale: float = Field(default=7.5, ge=0.0)
    width: int = Field(default=512, ge=64, le=1024)
    height: int = Field(default=512, ge=64, le=1024)
    seed: int = Field(default=-1)
    num_images_per_prompt: int = Field(default=1, ge=1, le=8)


def _to_b64(img) -> str:
    buf = BytesIO()
    img.save(buf, format="JPEG", quality=90)
    return base64.b64encode(buf.getvalue()).decode()


@app.on_event("startup")
def _startup():
    load_pipeline()


@app.get("/health")
def health():
    return {
        "status": "ok" if pipe is not None else "loading",
        "model": os.path.basename(MODEL),
        "device": str(device),
    }


@app.get("/")
def root():
    return {
        "service": "onlyfake realistic-vision-v4 (FP16)",
        "endpoints": ["/health", "/generate"],
    }


@app.post("/generate")
def generate(req: GenRequest):
    if pipe is None:
        return {"error": "model not loaded"}

    gen = None
    if req.seed >= 0:
        gen = torch.Generator(device=device).manual_seed(req.seed)

    with torch.no_grad():
        result = pipe(
            prompt=req.prompt,
            negative_prompt=req.negative_prompt,
            num_inference_steps=req.num_inference_steps,
            guidance_scale=req.guidance_scale,
            width=req.width,
            height=req.height,
            num_images_per_prompt=req.num_images_per_prompt,
            generator=gen,
        )

    return {"images": [_to_b64(img) for img in result.images]}