import gradio as gr
import random
import os
from PIL import Image, ImageDraw, ImageFont
import textwrap
import torch
import spaces

device = "cuda" if torch.cuda.is_available() else "cpu"

@spaces.GPU(enable_queue=True)
def load_pipeline():
    from diffusers import AutoPipelineForImage2Image
    pipe = AutoPipelineForImage2Image.from_pretrained(
        "stabilityai/stable-diffusion-2-1",
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        variant="fp16" if device == "cuda" else None,
        use_safetensors=True,
    )
    if device == "cuda":
        pipe.enable_model_cpu_offload()
    return pipe

pipe = load_pipeline()

# ====================== HELPERS ======================
def list_files(prefix):
    return [f for f in os.listdir(".") if f.startswith(prefix) and f.lower().endswith((".png", ".jpg", ".jpeg"))]

def get_random_base():
    bases = list_files("saint_base_")
    return Image.open(random.choice(bases)).convert("RGB") if bases else Image.new("RGB", (512, 512), "#2a0d4a")

def create_saint_card(generated_img: Image.Image, saint_name: str, chaotic_sin: str):
    frames = list_files("frame_")
    frame_path = frames[0] if frames else None
    if frame_path:
        frame = Image.open(frame_path).convert("RGBA")
    else:
        frame = Image.new("RGBA", (800, 1000), (42, 13, 74, 255))
    
    # Adjust these numbers so the photo sits perfectly inside your frame
    inner_x, inner_y, inner_w, inner_h = 90, 140, 620, 620
    saint_resized = generated_img.resize((inner_w, inner_h), Image.Resampling.LANCZOS)
    frame.paste(saint_resized, (inner_x, inner_y))
    
    draw = ImageDraw.Draw(frame)
    try:
        font_big = ImageFont.truetype("FONT.ttf", 62)
        font_med = ImageFont.truetype("FONT.ttf", 44)
        font_small = ImageFont.truetype("FONT.ttf", 34)
    except:
        font_big = font_med = font_small = ImageFont.load_default()
    
    draw.text((130, 55), saint_name.upper(), font=font_big, fill="#f5c242", stroke_width=6, stroke_fill="#2a0d4a")
    draw.text((150, 740), chaotic_sin.upper(), font=font_med, fill="#ffd700", stroke_width=4, stroke_fill="#2a0d4a")
    draw.text((200, 860), "STILL HAS TO CONFESS", font=font_small, fill="#ffffff", stroke_width=3, stroke_fill="#2a0d4a")
    draw.text((280, 905), "HOLY SMOKES", font=font_small, fill="#f5c242", stroke_width=3, stroke_fill="#2a0d4a")
    
    return frame

def create_oracle_card(saint_name: str, chaotic_sin: str, vibe: str, accessory: str):
    width, height = 620, 860
    parchment = Image.new("RGB", (width, height), "#f5e8c7")
    draw = ImageDraw.Draw(parchment)
    draw.rectangle((25, 25, width-25, height-25), outline="#d4a017", width=24)
    draw.rectangle((45, 45, width-45, height-45), outline="#2a0d4a", width=10)
    
    try:
        font_title = ImageFont.truetype("FONT.ttf", 44)
        font_body = ImageFont.truetype("FONT.ttf", 29)
        font_banner = ImageFont.truetype("FONT.ttf", 34)
    except:
        font_title = font_body = font_banner = ImageFont.load_default()
    
    prayer = f"""Prayer to {saint_name} the Unbothered

Oh {saint_name}, 3 a.m. comeback daily.
Bless me with {accessory} as my dignity's grease
and that unbothered face slays.
Edges snatched, receipts gone because of {chaotic_sin}.

Werk.

Amen. ✨"""

    draw.text((width//2, 90), "Prayer to", font=font_title, fill="#2a0d4a", anchor="mm")
    draw.text((width//2, 145), f"{saint_name} the Unbothered", font=font_title, fill="#2a0d4a", anchor="mm")
    wrapped = textwrap.fill(prayer.split("\n\n")[1], width=34)
    draw.text((72, 230), wrapped, font=font_body, fill="#2a0d4a", spacing=8)
    draw.text((width//2, height - 95), "STILL HAS TO CONFESS", font=font_banner, fill="#d4a017", anchor="mm")
    draw.text((width//2, height - 52), "HOLY SMOKES", font=font_banner, fill="#d4a017", anchor="mm")
    
    return parchment

@spaces.GPU
def canonize(photo, saint_name, chaotic_sin, divine_vibe, sacred_accessory, drama, debauch, delulu):
    if not saint_name: saint_name = "Santa Vibra Raro"
    if not chaotic_sin: chaotic_sin = "Still has to confess"
    
    init_image = photo.convert("RGB") if photo is not None else get_random_base()
    
    prompt = (f"tarot card style portrait of Saint {saint_name}, {divine_vibe.lower()}, "
              f"holding {sacred_accessory.lower()}, ethereal holy glow, dramatic lighting, masterpiece")
    
    negative = "blurry, ugly, deformed, extra limbs, bad anatomy, low quality"
    
    strength = 0.35 + (delulu * 0.065)
    guidance = 6.0 + (drama * 1.5)
    steps = 18 + (debauch * 5)
    
    result = pipe(
        prompt=prompt,
        image=init_image,
        strength=strength,
        guidance_scale=guidance,
        num_inference_steps=steps,
        negative_prompt=negative,
    ).images[0]
    
    saint_card = create_saint_card(result, saint_name, chaotic_sin)
    oracle_card = create_oracle_card(saint_name, chaotic_sin, divine_vibe, sacred_accessory)
    
    return saint_card, oracle_card

# ====================== GRADIO UI WITH GALAXY LOOK ======================
with gr.Blocks(
    title="HOLY SMOKES SAINTIFIER",
    theme=gr.themes.Dark(),
    css="""
    body { background: radial-gradient(circle at center, #2a0d4a 0%, #000022 70%, #000011 100%) !important; }
    .gradio-container { background: transparent !important; }
    .phone-frame { 
        background: #111; 
        border-radius: 55px; 
        padding: 18px; 
        box-shadow: 0 0 60px rgba(245,194,66,0.7), inset 0 0 30px rgba(255,255,255,0.1);
        max-width: 360px; 
        margin: 20px auto;
    }
    """
) as demo:
    gr.HTML("""
    <div style="text-align:center; padding:40px 0 10px;">
        <h1 style="font-size:4.2rem; color:#f5c242; font-family:'UnifrakturCook', cursive; text-shadow:0 0 30px #f5c242;">HOLY SMOKES</h1>
        <h2 style="font-size:3rem; color:#f5c242; font-family:'UnifrakturCook', cursive; margin-top:-20px;">SAINTIFIER</h2>
        <p style="font-size:1.5rem; color:#e0c0ff;">Upload chaos → Get canonized 🔥</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            photo = gr.Image(label="📸 Upload your photo (optional)", type="pil", height=320)
            saint_name = gr.Textbox(label="👑 Saint Name", value="Santa Vibra Raro")
            chaotic_sin = gr.Textbox(label="💥 Your Most Chaotic Sin", value="Still has to confess")
            
            with gr.Row():
                gr.Button("🎲 Random Saint Base", variant="secondary").click(
                    fn=get_random_base, outputs=photo)
                canonize_btn = gr.Button("✨ CANONIZE ME", variant="primary", size="large")
        
        with gr.Column(scale=1):
            with gr.Row():
                with gr.Column():
                    gr.Markdown("**Saint Card** (with your frame)")
                    saint_output = gr.Image(label=None, height=680, show_label=False, container=False)
                
                with gr.Column():
                    gr.Markdown("**Oracle / Prayer Card**")
                    oracle_output = gr.Image(label=None, height=680, show_label=False, container=False)
    
    with gr.Accordion("🌟 Advanced Chaos Options (mapped to img2img)", open=False):
        divine_vibe = gr.Dropdown(choices=["Saint of Divine Titles", "Saint of Glitter Chaos", "Saint of Unbothered Vibes", "Saint of Delulu Energy"], value="Saint of Divine Titles", label="Your Divine Vibe")
        sacred_accessory = gr.Dropdown(choices=["Money & Bills", "Sequins & Drama", "Crystal Ball of Chaos", "Receipts & Regrets"], value="Money & Bills", label="Your Sacred Accessory")
        with gr.Row():
            drama = gr.Slider(1, 5, value=3, step=1, label="🎭 Drama Intensity")
            debauch = gr.Slider(1, 5, value=2, step=1, label="🍷 Level of Debauch")
            delulu = gr.Slider(1, 10, value=6, step=1, label="🌀 Level of Delulu")
    
    canonize_btn.click(
        fn=canonize,
        inputs=[photo, saint_name, chaotic_sin, divine_vibe, sacred_accessory, drama, debauch, delulu],
        outputs=[saint_output, oracle_output]
    )

    # Show your actual files in the interface (great for testing)
    gr.Markdown("### Your files currently loaded:")
    gr.Markdown(f"**Frames:** {list_files('frame_')}")
    gr.Markdown(f"**Saint Bases:** {list_files('saint_base_')}")

demo.launch()