import io
import os
import random
import threading
import time
from queue import Queue, Empty

import torch
from fastapi import FastAPI, Response
from fastapi.responses import JSONResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from diffusers import AutoPipelineForText2Image
from PIL import Image

MODEL_ID = os.getenv("MODEL_ID", "stabilityai/sd-turbo")

W = int(os.getenv("W", "384"))
H = int(os.getenv("H", "384"))

STEPS = int(os.getenv("STEPS", "3"))
GUIDANCE = float(os.getenv("GUIDANCE", "0.0"))

USE_FP32_ON_MPS = os.getenv("USE_FP32_ON_MPS", "1") == "1"
QUEUE_MAX = int(os.getenv("QUEUE_MAX", "2"))

# Epoch distribution
# 50%: Impressionism+ → Contemporary
# 30%: Early Renaissance → Romanticism
# 20%: Pre-Gothic / Medieval / Gothic / Icon-like
P_MODERN = float(os.getenv("P_MODERN", "0.50"))
P_CLASSIC = float(os.getenv("P_CLASSIC", "0.30"))
P_EARLY = float(os.getenv("P_EARLY", "0.20"))

# Content balance
CLASSIC_PEOPLE_WEIGHT = float(os.getenv("CLASSIC_PEOPLE_WEIGHT", "0.60"))  # in classic 30%, more portraits/figures than landscapes
MODERN_PEOPLE_WEIGHT = float(os.getenv("MODERN_PEOPLE_WEIGHT", "0.35"))    # in modern 50%, fewer faces (helps with hands)

# Strong "no frame" enforcement
NEGATIVE = (
    # Kill frames / borders / museum context
    "frame, picture frame, painting frame, ornate frame, gold frame, "
    "border, canvas edge, cropped canvas, mat, passepartout, "
    "gallery wall, museum wall, hanging painting, framed artwork, "
    "wood frame, gilded frame, edge of painting, "
    # Kill text/logos
    "text, watermark, logo, signature, letters, "
    # Safety
    "nsfw, nude, naked, porn, gore, violence, "
    # Quality
    "blur, blurry, out of focus, lowres, jpeg artifacts, "
    # Anatomy (but do NOT force hands in positive prompt)
    "bad anatomy, bad proportions, bad face, deformed face, "
    "bad hands, malformed hands, deformed hands, "
    "extra fingers, missing fingers, fused fingers, extra limbs, "
    "hands in foreground, close-up hands, cropped hands, "
    # Avoid photo/CGI look
    "photorealistic, hyperrealistic, cgi, 3d render, plastic skin, anime, cartoon"
)

# ---- Prompt building blocks ----

BASE_PAINTING_QUALITY = (
    "fine art painting, museum-quality artwork, painterly, expressive brushwork, "
    "coherent artistic style, unified composition, natural color harmony, "
    "sharp focus, high detail, canvas texture subtle, "
    "no frame, no border, no canvas edge, "
    "no photographic realism"
)

COMPOSITION_CALM = (
    "balanced composition, calm pose, medium shot, hands not emphasized, "
    "hands partially obscured by clothing or out of frame, "
    "no dramatic hand gestures, hands not in foreground"
)

LIGHTING_SOFT = "soft natural light, gentle contrast, pleasing tonal range"
LIGHTING_DRAMATIC = "dramatic chiaroscuro, deep shadows, warm highlights"

# ---- 20% Early (pre-gothic / medieval / gothic / icons) ----
EARLY_POOL = [
    "early medieval illuminated manuscript style, flat composition, symbolic forms, tempera, muted pigments",
    "byzantine icon painting style, gold leaf tones, sacred atmosphere, stylized features, tempera on wood",
    "gothic panel painting style, elongated forms, ornate patterns, flat background, tempera",
    "romanese mural painting style, fresco texture, simplified figures, symbolic composition",
    "medieval devotional painting style, stylized drapery, flat shapes, decorative borders implied (but no frame)",
]

# ---- 30% Classic (early renaissance → romanticism) ----
CLASSIC_PEOPLE = [
    "early renaissance oil painting portrait, sfumato, subtle realism, classical balance",
    "high renaissance portrait painting, refined anatomy, calm expression, old master",
    "baroque oil painting figure scene, rich pigments, theatrical lighting",
    "dutch golden age interior scene painting, soft window light, oil on canvas",
    "romanticism portrait painting, warm skin tones, painterly texture, emotional mood",
]

CLASSIC_LANDSCAPES = [
    "renaissance landscape painting, atmospheric perspective, classical composition",
    "baroque landscape painting, dramatic sky, warm highlights, painterly",
    "romantic landscape painting, luminous clouds, distant horizon, oil on canvas",
    "classical pastoral landscape painting, soft light, calm mood, painterly",
]

# ---- 50% Modern+ (impressionism → postimpressionism → modern → contemporary) ----
MODERN_PEOPLE = [
    "impressionist portrait painting, visible brush strokes, light and color, soft edges",
    "post-impressionist portrait painting, structured brushwork, rich color, painterly",
    "fauvism portrait painting, bold color harmony, simplified shapes, expressive",
    "expressionist figure painting, energetic brushwork, emotional color, painterly",
    "modern figurative painting, simplified forms, contemporary palette, painterly",
]

MODERN_LANDSCAPES = [
    "impressionist landscape painting, plein air, shimmering light, visible brush strokes",
    "post-impressionist landscape painting, vibrant color, structured strokes, painterly",
    "fauvism landscape painting, bold color fields, simplified forms, expressive",
    "modern abstract landscape-inspired painting, color fields, texture, painterly",
    "contemporary painting, abstract forms, subtle glitch-like texture, mixed media feel (still painterly)",
    "minimal color-field painting, soft gradients, subtle texture, contemporary art",
]

def weighted_choice(groups):
    r = random.random()
    acc = 0.0
    for p, name in groups:
        acc += p
        if r <= acc:
            return name
    return groups[-1][1]

def pick_epoch_group():
    total = P_MODERN + P_CLASSIC + P_EARLY
    if total <= 0:
        return "modern"
    pm = P_MODERN / total
    pc = P_CLASSIC / total
    pe = P_EARLY / total
    return weighted_choice([(pm, "modern"), (pc, "classic"), (pe, "early")])

def pick_prompt():
    epoch = pick_epoch_group()

    if epoch == "early":
        style = random.choice(EARLY_POOL)
        return f"{style}, fine art painting, no frame, no border, painterly"

    if epoch == "classic":
        style = random.choice(CLASSIC_PEOPLE) if random.random() < CLASSIC_PEOPLE_WEIGHT else random.choice(CLASSIC_LANDSCAPES)
        return f"{style}, oil painting, museum quality, no frame, no border, calm pose"

    style = random.choice(MODERN_PEOPLE) if random.random() < MODERN_PEOPLE_WEIGHT else random.choice(MODERN_LANDSCAPES)
    return f"{style}, painterly, expressive brushwork, no frame, no border"

    # modern (50%)
    # Reduce hands/faces issues by favoring landscapes/abstract more often
    if random.random() < MODERN_PEOPLE_WEIGHT:
        style = random.choice(MODERN_PEOPLE)
        comp = COMPOSITION_CALM
    else:
        style = random.choice(MODERN_LANDSCAPES)
        comp = "balanced composition, no frame, no border"
    lighting = random.choice([LIGHTING_SOFT, "natural daylight, atmospheric light", "soft studio light"])
    return f"{style}, {BASE_PAINTING_QUALITY}, {lighting}, {comp}"

# Device selection
if torch.backends.mps.is_available():
    DEVICE = "mps"
elif torch.cuda.is_available():
    DEVICE = "cuda"
else:
    DEVICE = "cpu"

if DEVICE == "mps":
    DTYPE = torch.float32 if USE_FP32_ON_MPS else torch.float16
elif DEVICE == "cuda":
    DTYPE = torch.float16
else:
    DTYPE = torch.float32

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

pipe = None
pipe_lock = threading.Lock()

q = Queue(maxsize=QUEUE_MAX)

latest_id = 0
last_error = ""
last_gen_ms = None
generated_total = 0
last_prompt = ""

def load_pipeline():
    global pipe
    pipe = AutoPipelineForText2Image.from_pretrained(
        MODEL_ID,
        torch_dtype=DTYPE,
        safety_checker=None,
        feature_extractor=None,
    ).to(DEVICE)
    try:
        pipe.set_progress_bar_config(disable=True)
    except Exception:
        pass

def render_png():
    global last_prompt
    prompt = pick_prompt()
    last_prompt = prompt

    t0 = time.perf_counter()
    with pipe_lock, torch.inference_mode():
        out = pipe(
            prompt=prompt,
            negative_prompt=NEGATIVE,
            width=W,
            height=H,
            num_inference_steps=STEPS,
            guidance_scale=GUIDANCE,
        )
    img: Image.Image = out.images[0]
    buf = io.BytesIO()
    img.save(buf, format="PNG", optimize=True)
    ms = (time.perf_counter() - t0) * 1000.0
    return buf.getvalue(), ms

def generator_loop():
    global latest_id, last_error, last_gen_ms, generated_total

    while True:
        try:
            png, ms = render_png()

            latest_id += 1
            last_gen_ms = ms
            last_error = ""
            generated_total += 1

            if q.full():
                try:
                    q.get_nowait()
                except Empty:
                    pass

            q.put((latest_id, png), timeout=1)

            if DEVICE == "mps":
                try:
                    torch.mps.empty_cache()
                except Exception:
                    pass

        except Exception as e:
            last_error = repr(e)
            time.sleep(0.5)

@app.on_event("startup")
async def startup():
    load_pipeline()
    threading.Thread(target=generator_loop, daemon=True).start()

@app.get("/", response_class=HTMLResponse)
def root():
    try:
        with open(os.path.join(os.path.dirname(__file__), "index.html"), "r", encoding="utf-8") as f:
            return f.read()
    except Exception:
        return "<html><body style='background:black;color:white;font-family:system-ui'>index.html not found</body></html>"

@app.get("/health")
def health():
    return JSONResponse({
        "device": DEVICE,
        "dtype": str(DTYPE),
        "model": MODEL_ID,
        "w": W,
        "h": H,
        "steps": STEPS,
        "guidance": GUIDANCE,
        "queue": q.qsize(),
        "latest_id": latest_id,
        "last_gen_ms": last_gen_ms,
        "last_error": last_error,
        "generated_total": generated_total,
        "p_modern": P_MODERN,
        "p_classic": P_CLASSIC,
        "p_early": P_EARLY,
        "classic_people_weight": CLASSIC_PEOPLE_WEIGHT,
        "modern_people_weight": MODERN_PEOPLE_WEIGHT,
        "last_prompt": last_prompt[:400],
    })

@app.get("/next")
def next_frame():
    fid, png = q.get(timeout=600)
    return Response(
        content=png,
        media_type="image/png",
        headers={"X-Frame-Id": str(fid), "Cache-Control": "no-store"},
    )