import gradio as gr
from transformers import pipeline, set_seed
from PIL import Image, ImageDraw, ImageFont
import textwrap
import io

# Lazy-load models
_captioner = None
_generator = None

def get_captioner():
    global _captioner
    if _captioner is None:
        _captioner = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
    return _captioner

def get_generator():
    global _generator
    if _generator is None:
        _generator = pipeline("text-generation", model="distilgpt2")
    return _generator

def wrap_text(text, width=22):
    return "\n".join(textwrap.wrap(text, width=width))

def draw_meme(img: Image.Image, caption: str, position: str = "bottom", uppercase: bool = True) -> Image.Image:
    if uppercase:
        caption = caption.upper()

    # Prepare drawing
    draw = ImageDraw.Draw(img)
    W, H = img.size

    # Choose a font (fallback to default if Impact is not available)
    try:
        font = ImageFont.truetype("Impact.ttf", size=max(24, W // 15))
    except Exception:
        font = ImageFont.load_default()

    # Wrap caption to fit width
    wrapped = wrap_text(caption, width=max(18, W // 25))

    # Determine text size
    bbox = draw.multiline_textbbox((0, 0), wrapped, font=font, align="center")
    text_w = bbox[2] - bbox[0]
    text_h = bbox[3] - bbox[1]

    # Choose y coordinate
    margin = int(0.03 * H)
    if position == "top":
        y = margin
    elif position == "center":
        y = (H - text_h) // 2
    else:
        y = H - text_h - margin

    # Draw black outline then white text (classic meme style)
    x = (W - text_w) // 2
    outline = max(2, W // 200)
    for dx in (-outline, outline):
        for dy in (-outline, outline):
            draw.multiline_text((x + dx, y + dy), wrapped, font=font, fill="black", align="center")
    draw.multiline_text((x, y), wrapped, font=font, fill="white", align="center")

    return img

def explain_caption(caption: str) -> str:
    cap = caption.lower()
    cues = []
    if any(w in cap for w in ["when ", "me:", "you ", "pov", "that moment"]):
        cues.append("relatable setup")
    if any(w in cap for w in ["but", "yet", "although", "instead"]):
        cues.append("incongruity/twist")
    if any(w in cap for w in ["too", "very", "super", "always", "never"]):
        cues.append("mild exaggeration")
    if any(w in cap for w in ["monday", "wifi", "coffee", "deadline", "sleep"]):
        cues.append("everyday theme")

    pieces = [
        "Short, high-contrast phrasing (classic meme style).",
        "Focuses on a common situation to trigger recognition.",
        "Adds a small twist that creates incongruity (the core of the joke)."
    ]
    if cues:
        pieces.append("Detected humor cues: " + ", ".join(cues) + ".")
    return " ".join(pieces)

def generate_meme(image, style, uppercase, seed):
    if image is None:
        return None, "Please upload an image.", ""

    if seed is not None and str(seed).strip() != "":
        try:
            set_seed(int(seed))
        except Exception:
            pass

    # Step 1: caption the image
    captioner = get_captioner()
    raw_caption = captioner(image)[0]["generated_text"]

    # Step 2: create a witty meme caption from the caption
    prompt = (
        "Write a short, witty meme caption (max 12 words) based on this image description: "
        f"'{raw_caption}'. Use relatable internet humor. Avoid names, no offensive content."
    )
    generator = get_generator()
    gen = generator(prompt, max_new_tokens=24, temperature=0.9, top_p=0.95, do_sample=True)[0]["generated_text"]
    # Get only the newly generated tail after the prompt
    meme_caption = gen[len(prompt):].strip().split("\n")[0]
    # Clean fallbacks
    if len(meme_caption) < 3 or len(meme_caption.split()) < 2:
        meme_caption = "When Monday hits before the coffee"

    # Step 3: draw caption onto the image
    img = image.convert("RGB").copy()
    img_out = draw_meme(img, meme_caption, position=style, uppercase=uppercase)

    # Step 4: generate a simple explanation
    explanation = explain_caption(meme_caption)

    return img_out, meme_caption, explanation

with gr.Blocks(title="AI Meme Generator + Explainer") as demo:
    gr.Markdown(
        """
        # AI Meme Generator + Explainer
        Upload an image → the app captions it, writes a short meme line, and explains why it works.
        Runs on CPU using lightweight models (BLIP-base + DistilGPT2).
        """
    )

    with gr.Row():
        image = gr.Image(type="pil", label="Upload image")
        with gr.Column():
            style = gr.Radio(choices=["top", "center", "bottom"], value="bottom", label="Caption position")
            uppercase = gr.Checkbox(value=True, label="Uppercase caption (meme style)")
            seed = gr.Textbox(value="", label="Seed (optional)", placeholder="e.g., 42")
            btn = gr.Button("Generate Meme")

    out_img = gr.Image(label="Meme")
    out_caption = gr.Textbox(label="Generated caption", lines=2)
    out_expl = gr.Textbox(label="Why it’s funny (explanation)", lines=4)

    btn.click(generate_meme, inputs=[image, style, uppercase, seed], outputs=[out_img, out_caption, out_expl])

if __name__ == "__main__":
    demo.launch()