import os
import gc
import random
import gradio as gr
import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont
from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes
import glob

import spaces
from diffusers import AutoPipelineForImage2Image
from saint_oracle import generate_saint_text

# ────────────────────────────────────────────────
# THEME
# ────────────────────────────────────────────────
colors.fire_red = colors.Color(
    name="fire_red",
    c50="#FFF5F0", c100="#FFE8DB", c200="#FFD0B5", c300="#FFB088",
    c400="#FF8C5A", c500="#FF6B35", c600="#E8531F", c700="#CC4317",
    c800="#A63812", c900="#80300F", c950="#5C220A",
)

class HolySmokesTheme(Soft):
    def __init__(self, *, primary_hue=colors.fire_red, secondary_hue=colors.fire_red, neutral_hue=colors.slate, **kwargs):
        super().__init__(primary_hue=primary_hue, secondary_hue=secondary_hue, neutral_hue=neutral_hue, **kwargs)
        super().set(
            body_background_fill="#f7f3ea",
            button_primary_background_fill="linear-gradient(135deg,#b48c2c,#ffd700)",
            button_primary_background_fill_hover="linear-gradient(135deg,#c89b34,#ffeb3b)",
            button_primary_shadow="0 6px 20px rgba(180,140,44,0.5)",
        )

theme = HolySmokesTheme()

# ────────────────────────────────────────────────
# MODEL – back to sd-turbo (your original fast & reliable model)
# ────────────────────────────────────────────────
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.float16 if torch.cuda.is_available() else torch.float32
pipe = AutoPipelineForImage2Image.from_pretrained(
    "stabilityai/sd-turbo",
    torch_dtype=dtype,
).to(device)

MAX_SEED = np.iinfo(np.int32).max

# ────────────────────────────────────────────────
# DYNAMIC ASSETS
# ────────────────────────────────────────────────
FRAME_FILES = sorted(
    glob.glob("frame_*.png"),
    key=lambda x: int(''.join(filter(str.isdigit, os.path.basename(x))))
)
FRAME_LABELS = [f"Frame {i+1}" for i in range(len(FRAME_FILES))]
BASE_FRAMES = [Image.open(f).convert("RGBA") for f in FRAME_FILES]

SAINT_BASE_FILES = sorted(
    glob.glob("saint_base_*.png"),
    key=lambda x: int(''.join(filter(str.isdigit, os.path.basename(x))))
)
SAINT_BASES = [Image.open(f).convert("RGBA") for f in SAINT_BASE_FILES]

if FRAME_FILES:
    FRAME_W, FRAME_H = BASE_FRAMES[0].size
else:
    raise RuntimeError("No frame_*.png files found!")

# ────────────────────────────────────────────────
# OPTIONS
# ────────────────────────────────────────────────
VIBE_OPTIONS = ["Saint of Drama", "Saint of Gay Chaos", "Saint of Delulu"]
COLOR_OPTIONS = ["Divine Gold", "Infernal Red", "Celestial Blue", "Pastel Angel", "Neon Drag", "Baroque Sepia"]
PROP_OPTIONS = ["Candles & Incense", "Money & Bills", "Phones & Screens", "Crowns & Halos", "Flowers & Thorns"]
SINS = ["Weaponized Gossip", "Luxury Without Income", "Texting The Ex", "Chronic Delusion", "Main Character Syndrome", "Gay Chaos at Brunch", "Financial Delusion"]
RANDOM_NAMES = ["Santa Delulu", "Gloria del Caos", "Our Lady of Brunch", "Saint Overthinker", "Madre de la Drama", "San Ex-Toxic"]

# ────────────────────────────────────────────────
# CARTOON SAINT PROMPT – optimized for clear illustration
# ────────────────────────────────────────────────
def generate_prompt(name, sin, personality, vibe, color, prop):
    return f"""
Make this person a ridiculously funny Saint illustration with large glowing golden halo.
Exact face preserved, cartoon style, vibrant colors, exaggerated campy baroque details.
Robes in {color}, holding {prop}, embodying {vibe} energy and the sin of {sin}.
Funny theatrical pose, queer-coded drama, painterly illustration style.
Clearly cartoon/illustration, not realistic photo. Preserve real identity perfectly.
"""

# ────────────────────────────────────────────────
# IMAGE HELPERS – ribbon fixed to your frame
# ────────────────────────────────────────────────
def resize_and_center(img: Image.Image) -> Image.Image:
    if img is None: return None
    img = img.convert("RGB")
    target_h = int(FRAME_H * 0.95)
    target_w = int(target_h * img.width / img.height)
    resized = img.resize((target_w, target_h), Image.LANCZOS)
    canvas = Image.new("RGB", (FRAME_W, FRAME_H), (20, 20, 30))
    canvas.paste(resized, ((FRAME_W - target_w) // 2, (FRAME_H - target_h) // 2))
    return canvas

def apply_frame(img: Image.Image, frame_label: str) -> Image.Image:
    idx = FRAME_LABELS.index(frame_label)
    frame = BASE_FRAMES[idx]
    out = img.convert("RGBA").resize((FRAME_W, FRAME_H), Image.LANCZOS)
    out.paste(frame, (0, 0), frame)
    return out

def draw_name_on_ribbon(img: Image.Image, name: str) -> Image.Image:
    if not name: return img
    img = img.convert("RGBA")
    draw = ImageDraw.Draw(img)
    w, h = img.size

    # ALIGNED TO YOUR FRAME'S LOWER RIBBON BANNER
    x1 = int(w * 0.02)
    x2 = int(w * 0.98)
    y1 = int(h * 0.68)
    y2 = int(h * 0.95)
    zone_width = x2 - x1
    zone_height = y2 - y1

    text = f"SAINT {name.upper()}"
    font_path = "font.ttf"
    base_font_size = 92 if os.path.isfile(font_path) else 40
    font_size = base_font_size
    font = None
    while font_size >= 26:
        try:
            font = ImageFont.truetype(font_path, font_size) if os.path.isfile(font_path) else ImageFont.load_default()
        except:
            font = ImageFont.load_default()
        bbox = draw.textbbox((0, 0), text, font=font)
        tw = bbox[2] - bbox[0]
        th = bbox[3] - bbox[1]
        if tw <= zone_width * 0.88 and th <= zone_height * 0.78:
            break
        font_size -= 3

    x = x1 + (zone_width - tw) // 2
    y = y1 + (zone_height - th) // 2

    for dx, dy in [(-3,-3), (-3,3), (3,-3), (3,3)]:
        draw.text((x + dx, y + dy), text, fill="black", font=font)
    draw.text((x, y), text, fill="black", font=font)
    return img

def apply_frame_and_ribbon(base_img: Image.Image, frame_label: str, name: str) -> Image.Image:
    if base_img is None: return None
    framed = apply_frame(base_img, frame_label)
    return draw_name_on_ribbon(framed, name)

# ────────────────────────────────────────────────
# ORACLE CARD – dynamic + auto-size + dark text
# ────────────────────────────────────────────────
def generate_oracle(name, sin=None, personality=None, vibe=None):
    name_clean = name.strip() if name else "Unnamed Saint"
    title = f"Saint {name_clean}"
    bio = f"Born from glitter, gossip, and deeply questionable decisions, this chaotic nun wanders the mortal realm spreading drama, delusion, and weaponized camp. Their sin is {sin or 'chaos'}."
    roast = "Their holiness is as real as their lashes: heavy, synthetic, and absolutely non-refundable."
    prayer = f"May your sins be glamorous, your chaos be blessed, and your enemies stay forever pressed. {vibe or ''}"
    return f"""
<div class="oracle-card">
    <div class="oracle-header"><b>The Holy Smokes Oracle Speaks</b></div>
    <div class="oracle-name" style="color:#000 !important;">{title}</div>
    <div class="oracle-body" style="color:#000 !important; white-space: pre-wrap;">
        <b>TITLE:</b> {title}<br><br>
        <b>BIO:</b> {bio}<br><br>
        <b>ROAST:</b> {roast}<br><br>
        <b>PRAYER:</b> {prayer}
    </div>
</div>
"""

def build_download_image(saint_img: Image.Image, oracle_text: str, saint_name: str) -> str:
    saint_img = saint_img.convert("RGB")
    w, h = saint_img.size
    lines = [line.strip() for line in oracle_text.split("\n") if line.strip()]
    canvas = Image.new("RGB", (w, h + 120 + len(lines)*28), "white")
    canvas.paste(saint_img, (0, 0))
    draw = ImageDraw.Draw(canvas)
    try:
        font = ImageFont.truetype("font.ttf", 24)
        title_font = ImageFont.truetype("font.ttf", 28)
    except:
        font = title_font = ImageFont.load_default()
    y = h + 20
    draw.text((40, y), f"The Holy Smokes Oracle Speaks — Saint {saint_name}", fill="black", font=title_font)
    y += 40
    for line in lines:
        draw.text((40, y), line, fill="black", font=font)
        y += 28
    out_path = "/tmp/holy_smokes_saint_and_oracle.jpg"
    canvas.save(out_path, format="JPEG", quality=95)
    return out_path

def format_info(seed_val, images):
    if images:
        try:
            first = images[0]
            path = first[0] if isinstance(first, (tuple, list)) else first
            im = Image.open(path if isinstance(path, str) else path.name)
            ow, oh = im.size
            return f"**Seed:** `{int(seed_val)}`\n\n**Original:** {ow}×{oh} → **Output:** {FRAME_W}×{FRAME_H}"
        except:
            pass
    return f"**Seed:** `{int(seed_val)}`"

def generate_local_oracle(name, sin, personality, vibe, color, prop):
    try:
        return generate_saint_text(
            name or "This Sinner",
            sin or "Unspecified Sin",
            personality or "Unspecified Personality",
            vibe or "Saint of Drama",
            color or "Divine Gold",
            prop or "Candles & Incense",
        )
    except Exception:
        return (
            "SAINT BIO\nThe oracle choked on your drama but you are still canonized.\n\n"
            "PRAYER\nMay your Wi-Fi be stable and your sins aesthetic."
        )

def randomize_fields():
    return (
        random.choice(SINS),
        random.choice(VIBE_OPTIONS),
        random.choice(COLOR_OPTIONS),
        random.choice(PROP_OPTIONS),
        random.choice(FRAME_LABELS),
    )

def random_saint():
    base_raw = random.choice(SAINT_BASES)
    base = resize_and_center(base_raw)
    name = random.choice(RANDOM_NAMES)
    sin = random.choice(SINS)
    personality = "Chaotic but blessed"
    vibe = random.choice(VIBE_OPTIONS)
    color = random.choice(COLOR_OPTIONS)
    prop = random.choice(PROP_OPTIONS)
    frame_label = random.choice(FRAME_LABELS)
    framed = apply_frame_and_ribbon(base, frame_label, name)
    saint_text = generate_local_oracle(name, sin, personality, vibe, color, prop)
    oracle_html = generate_oracle(name)
    download_path = build_download_image(framed, saint_text, name)
    return framed, oracle_html, name, sin, personality, vibe, color, prop, frame_label, download_path, base

@spaces.GPU(duration=120)
def run_canonizer(
    images,
    name,
    sin,
    personality,
    vibe,
    color,
    prop,
    frame_label,
    seed,
    randomize_seed,
    guidance_scale,
    steps,
    current_base,
    progress=gr.Progress(track_tqdm=True),
):
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    source = None
    if images and images[0]:
        first = images[0]
        path = first[0] if isinstance(first, (tuple, list)) else first
        try:
            source = Image.open(path if isinstance(path, str) else path.name).convert("RGB")
        except:
            pass

    if source is None and current_base is not None:
        source = current_base.convert("RGB")

    if source is None:
        raise gr.Error("No image to canonize.\nUpload a photo or generate a random saint first.")

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(int(seed))

    base_small = source.resize((512, 512), Image.LANCZOS)
    prompt = generate_prompt(name, sin, personality, vibe, color, prop)

    try:
        result_small = pipe(
            image=base_small,
            prompt=prompt,
            guidance_scale=float(guidance_scale),
            num_inference_steps=int(steps),
            strength=0.45,
            generator=generator,
        ).images[0]
    except Exception as e:
        print("Generation error:", str(e))
        raise gr.Error("Generation failed. Try fewer steps or lower guidance.")

    unframed = resize_and_center(result_small)
    framed = apply_frame_and_ribbon(unframed, frame_label, name or "Unnamed Saint")

    saint_text = generate_local_oracle(name, sin, personality, vibe, color, prop)
    oracle_html = generate_oracle(name or "Unnamed Saint")
    download_path = build_download_image(framed, saint_text, name or "saint")

    return framed, oracle_html, seed, download_path, unframed

def update_preview_frame(frame_label, current_unframed, name):
    if current_unframed is None:
        return None
    return apply_frame_and_ribbon(current_unframed, frame_label, name)

def auto_preview_upload(images, frame_label, name):
    if not images or not images[0]:
        return None
    first = images[0]
    path = first[0] if isinstance(first, (tuple, list)) else first
    try:
        raw = Image.open(path if isinstance(path, str) else path.name).convert("RGB")
        centered = resize_and_center(raw)
        return apply_frame_and_ribbon(centered, frame_label, name or "Unnamed Saint")
    except:
        return None

# ────────────────────────────────────────────────
# CSS – full frame fill + dark oracle + auto-size
# ────────────────────────────────────────────────
css = """
#col-container { max-width: 1100px; margin: 0 auto; }
.hs-header { text-align:center; padding:40px 20px; background:linear-gradient(135deg,#000,#2b0050,#5a1a8a); border-radius:20px; margin-bottom:24px; box-shadow:0 12px 44px rgba(0,0,0,.4); color:white; }
.hs-header img { width:180px; margin-bottom:12px; filter:drop-shadow(0 0 20px #ffd700); display:block; margin-left:auto; margin-right:auto; }

.oracle-card { background:white !important; border:4px solid #b48c2c; border-radius:22px; padding:36px 40px; box-shadow:0 12px 36px rgba(180,140,44,0.45); margin:28px auto; max-width:580px; height:auto; }
.oracle-header { font-size:1.45rem; font-weight:900; text-transform:uppercase; letter-spacing:0.12em; color:#b48c2c; }
.oracle-name { font-size:1.4rem; font-weight:900; color:#000 !important; margin:16px 0; }
.oracle-body { font-size:1.08rem; line-height:1.8; color:#000 !important; font-weight:600; text-align:left; white-space:pre-wrap; }

#gen-btn { background:linear-gradient(135deg,#b48c2c,#ffd700) !important; color:white !important; font-weight:800 !important; padding:16px 32px !important; border-radius:18px !important; box-shadow:0 8px 24px rgba(180,140,44,0.6) !important; }
#output-img { border-radius:0 !important; object-fit:cover !important; width:100% !important; height:100% !important; }
"""

# ────────────────────────────────────────────────
# UI
# ────────────────────────────────────────────────
with gr.Blocks(title="Holy Smokes Saintifier 3.0", theme=theme, css=css) as demo:
    current_unframed = gr.State(None)

    with gr.Column(elem_id="col-container"):
        gr.HTML("""
<div class="hs-header">
    <img src="/file=HSLogo.png" onerror="this.src='/file=HSLogo1.png'; this.onerror=null;" alt="Holy Smokes Logo">
    <h1>Holy Smokes Saintifier</h1>
    <p>Upload your chaos or summon a random saint. Either way, you're getting canonized.</p>
</div>
        """)

        with gr.Row(equal_height=False):
            with gr.Column(scale=1):
                gr.HTML('<p class="stl">📤 Upload Photo + Chaos References</p>')
                images = gr.Gallery(
                    label="Upload main photo + optional extras",
                    type="filepath",
                    columns=2,
                    rows=1,
                    height=260,
                    allow_preview=True,
                    object_fit="contain",
                )
                gr.HTML('<p class="stl" style="margin-top:12px">📝 Who is this saint?</p>')
                name = gr.Textbox(label="Name", placeholder="Your name or alter ego")
                sin = gr.Textbox(label="Sin", placeholder="Weaponized gossip, texting the ex, etc.")
                personality = gr.Textbox(label="Personality", placeholder="Delulu, dramatic, soft life, etc.")
                vibe = gr.Dropdown(VIBE_OPTIONS, value="Saint of Drama", label="Vibe")
                color = gr.Dropdown(COLOR_OPTIONS, value="Divine Gold", label="Palette")
                prop = gr.Dropdown(PROP_OPTIONS, value="Candles & Incense", label="Prop")
                frame_choice = gr.Dropdown(FRAME_LABELS, value="Frame 1", label="Frame")

                with gr.Row():
                    random_saint_btn = gr.Button("🔥 Random Saint", variant="secondary")
                    random_btn = gr.Button("🎲 Random Chaos (fields only)", variant="secondary")

                with gr.Row():
                    canonize_btn = gr.Button("✨ Canonize Me", variant="primary", elem_id="gen-btn")
                    clear_button = gr.Button("🗑️ Clear", variant="secondary")

                reveal_oracle_btn = gr.Button("🔮 Reveal Oracle", variant="secondary")

                with gr.Accordion("⚙️ Advanced (FireRed)", open=False):
                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                    randomize_seed = gr.Checkbox(label="🎲 Randomize seed", value=True)
                    guidance_scale = gr.Slider(
                        label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0,
                    )
                    steps = gr.Slider(
                        label="Inference Steps", minimum=1, maximum=30, step=1, value=4,
                    )

            with gr.Column(scale=1):
                gr.HTML('<p class="stl">🖼️ Your Saint</p>')
                output_image = gr.Image(
                    show_label=False,
                    interactive=False,
                    format="png",
                    width="100%",
                    height=FRAME_H,
                    elem_id="output-img",
                )
                oracle_md = gr.HTML()
                info_box = gr.Markdown(
                    value="*Generate or summon a saint to see details here.*",
                    elem_id="info-box",
                )
                download_btn = gr.DownloadButton(
                    "📥 Download Saint + Oracle JPG",
                    visible=True,
                )

    # ────────────────────────────────────────────────
    # EVENTS
    # ────────────────────────────────────────────────

    frame_choice.change(update_preview_frame, inputs=[frame_choice, current_unframed, name], outputs=output_image)
    images.change(auto_preview_upload, inputs=[images, frame_choice, name], outputs=output_image)

    canonize_btn.click(
        run_canonizer,
        inputs=[images, name, sin, personality, vibe, color, prop, frame_choice, seed, randomize_seed, guidance_scale, steps, current_unframed],
        outputs=[output_image, oracle_md, seed, download_btn, current_unframed],
    ).then(format_info, inputs=[seed, images], outputs=info_box)

    reveal_oracle_btn.click(
        generate_oracle,
        inputs=[name],
        outputs=[oracle_md],
    )

    random_saint_btn.click(random_saint, outputs=[output_image, oracle_md, name, sin, personality, vibe, color, prop, frame_choice, download_btn, current_unframed])
    random_btn.click(randomize_fields, outputs=[sin, vibe, color, prop, frame_choice])

    clear_button.click(
        lambda: (None, "", "", "", "Saint of Drama", "Divine Gold", "Candles & Incense", "Frame 1", None, "*Generate or summon a saint to see details here.*", "", None, None),
        outputs=[images, name, sin, personality, vibe, color, prop, frame_choice, output_image, info_box, oracle_md, download_btn, current_unframed],
    )

    demo.load(random_saint, outputs=[output_image, oracle_md, name, sin, personality, vibe, color, prop, frame_choice, download_btn, current_unframed])

if __name__ == "__main__":
    demo.launch(allowed_paths=["."])