import os
import base64
import threading
from io import BytesIO
from contextlib import asynccontextmanager

import torch
from fastapi import FastAPI
from pydantic import BaseModel

BASE = os.path.dirname(os.path.abspath(__file__))
MODEL = os.environ.get(
    "MODEL",
    os.path.join(BASE, "rv6_diffusers"),
)

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

pipe = None
load_error = None
_load_thread = None


def _load():
    global pipe, load_error
    from diffusers import StableDiffusionPipeline
    try:
        print(f"[load] model={MODEL}", flush=True)
        print(f"[load] torch {torch.__version__}, threads="
              f"{torch.get_num_threads()}", flush=True)
        p = StableDiffusionPipeline.from_pretrained(
            MODEL,
            torch_dtype=torch.float16,
            safety_checker=None,
            requires_safety_checker=False,
            disable_mmap=True,
        )
        p = p.to("cpu")
        if hasattr(p, "enable_attention_slicing"):
            p.enable_attention_slicing()
        for meth in ("enable_vae_tiling",):
            if hasattr(p, meth):
                getattr(p, meth)()
        p.set_progress_bar_config(disable=True)
        pipe = p
        print("[load] PIPELINE READY", flush=True)
    except Exception as e:  # noqa: BLE001
        import traceback
        traceback.print_exc()
        load_error = f"{type(e).__name__}: {e}"
        print(f"[load] FAILED: {load_error}", flush=True)


@asynccontextmanager
async def lifespan(_app):
    global _load_thread
    _load_thread = threading.Thread(target=_load, daemon=True)
    _load_thread.start()
    yield


app = FastAPI(title="onlyfake realistic-vision (RV6)", lifespan=lifespan)


class GenRequest(BaseModel):
    prompt: str
    negative_prompt: str = DEFAULT_NEGATIVE
    num_inference_steps: int = 12
    guidance_scale: float = 7.0
    width: int = 384
    height: int = 384
    seed: int = -1


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


@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-v6", "endpoints": ["/health", "/generate", "/docs"]}