import os
import base64
import torch
from io import BytesIO
from fastapi import FastAPI
from pydantic import BaseModel
from diffusers import StableDiffusionPipeline
from contextlib import asynccontextmanager
from huggingface_hub import hf_hub_download

# Onlyfake/realistic-vision-v4 LFS weights are broken (135B pointers).
# Use official equivalent weights.
MODEL_REPO = os.environ.get("MODEL_REPO", "SG161222/Realistic_Vision_V4.0_noVAE")
MODEL_FILE = os.environ.get("MODEL_FILE", "Realistic_Vision_V4.0.safetensors")
pipe = None

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"
)


@asynccontextmanager
async def lifespan(app: FastAPI):
    global pipe
    print(f"Downloading {MODEL_REPO}/{MODEL_FILE} ...", flush=True)
    ckpt = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
    size_gb = os.path.getsize(ckpt) / (1024**3)
    print(f"Checkpoint ready: {ckpt} ({size_gb:.2f} GB)", flush=True)
    print("Loading pipeline on CPU (this can take several minutes)...", flush=True)
    pipe = StableDiffusionPipeline.from_single_file(
        ckpt,
        torch_dtype=torch.float32,
        safety_checker=None,
        requires_safety_checker=False,
        low_cpu_mem_usage=True,
    )
    pipe = pipe.to("cpu")
    pipe.enable_attention_slicing("max")
    print("Model loaded.", flush=True)
    yield


app = FastAPI(title="Realistic Vision V4", lifespan=lifespan)


class GenRequest(BaseModel):
    prompt: str
    negative_prompt: str = DEFAULT_NEGATIVE
    num_inference_steps: int = 20
    guidance_scale: float = 7.5
    width: int = 512
    height: int = 512
    seed: int = -1


@app.get("/health")
def health():
    return {
        "status": "ok" if pipe is not None else "loading",
        "model": f"{MODEL_REPO}/{MODEL_FILE}",
        "device": "cpu",
        "note": "Onlyfake weights broken; using SG161222 official equivalent",
    }


@app.post("/generate")
def generate(req: GenRequest):
    if pipe is None:
        return {"error": "model not loaded"}
    generator = None
    if req.seed >= 0:
        generator = torch.Generator().manual_seed(req.seed)
    image = 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,
        generator=generator,
    ).images[0]
    buf = BytesIO()
    image.save(buf, format="JPEG", quality=90)
    return {"image": base64.b64encode(buf.getvalue()).decode()}


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