"""
Dispatch AI — Meme Generator
Gradio app for AI/tech memes. Top+bottom caption input, FLUX image gen, PIL text overlay.
"""

import os
import io
import random
import gradio as gr
from huggingface_hub import InferenceClient
from PIL import Image, ImageDraw, ImageFont

# --- Configuration -----------------------------------------------------------
HF_TOKEN = os.environ.get("HF_TOKEN", None)
MODEL_ID = "black-forest-labs/FLUX.1-schnell"

client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)

BG_COLOR = "#0A0F1A"
ACCENT = "#1FE0E6"

# Preset templates: name -> (base_image_prompt, default_top, default_bottom)
TEMPLATES = {
    "This is Fine AI": (
        "A cartoon character sitting in a burning AI server room surrounded by flames and melting GPUs, "
        "smiling calmly and saying 'this is fine', meme style, comic illustration, white background panels",
        "MY GPU IS ON FIRE",
        "THIS IS FINE",
    ),
    "Distracted Developer": (
        "A meme photo of a developer walking with his laptop but looking back at a passing shiny new robot "
        "labeled 'NEW AI MODEL', laptop labeled 'MY CURRENT MODEL', meme style photo",
        "MY CURRENT MODEL",
        "NEW HYPED AI MODEL",
    ),
    "Drake Format": (
        "Drake hotline bling meme format two-panel image, top panel Drake disapproving with hand up, "
        "bottom panel Drake approving pointing finger, meme template",
        "WRITING CODE MYSELF",
        "ASKING AI TO DO IT",
    ),
    "Expanding Brain": (
        "Expanding brain meme format, four panels showing progressively glowing brighter brains, meme template",
        "USING AI",
        "FINE-TUNING AI ON MY DATA",
    ),
    "Galaxy Brain": (
        "Galaxy brain meme format, three panels with progressively more cosmic brain, meme template",
        "PROMPTING CHATGPT",
        "TRAINING MY OWN MODEL",
    ),
    "Stonks": (
        "Stonks meme format, crudely drawn man in suit standing in front of upward arrow chart, meme meme",
        "CLOUD API COSTS",
        "ON-DEVICE INFERENCE",
    ),
    "Two Buttons": (
        "Two buttons meme format, sweating man deciding between two red buttons, meme template",
        "PAY $0.06 PER 1K TOKENS",
        "RUN ON MY PHONE",
    ),
    "Is This a Pigeon": (
        "Is this a pigeon meme format, anime character looking at butterfly confused, meme template",
        "IS THIS AN AGENT?",
        "IT'S A WHILE LOOP",
    ),
}

# Random AI topics for the Random AI Topic button
AI_TOPICS = [
    "AGI is coming next year (again)",
    "My model hallucinated a whole paper",
    "Prompt engineering is real engineering",
    "I fine-tuned on my DMs",
    "The model refused my prompt",
    "I ran out of GPU credits",
    "My agent looped forever",
    "Context window too small for my TODO list",
    "Open source beat closed source",
    "Inference costs more than my rent",
    "My dataset was just vibes",
    "RLHF made my bot too polite",
    "I forgot to clip the gradients",
    "The model learned to lie",
    "Quantization ruined my quality",
    "My LoRA saved my marriage",
    "Distillation taught the student to cheat",
    "The tokenizer hates Arabic",
    "I spent $0 on training and won the benchmark",
    "Agents are just while loops with vibes",
    "My gradient exploded and so did my career",
    "I asked the model to debug itself",
    "The embedding space is just vibes in high dimensions",
]

# Try to find a usable font
FONT_CANDIDATES = [
    "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
    "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf",
    "C:\\Windows\\Fonts\\arialbd.ttf",
    "arial.ttf",
]


def load_font(size):
    for path in FONT_CANDIDATES:
        if os.path.exists(path):
            return ImageFont.truetype(path, size)
    return ImageFont.load_default()


def random_topic():
    return random.choice(AI_TOPICS)


# --- Core function -----------------------------------------------------------
def generate_meme(top_caption, bottom_caption, template_name):
    """Generate a meme image: generate base via FLUX, then overlay text."""
    if template_name and template_name in TEMPLATES:
        base_prompt, def_top, def_bot = TEMPLATES[template_name]
        top = top_caption.strip() if top_caption.strip() else def_top
        bot = bottom_caption.strip() if bottom_caption.strip() else def_bot
    else:
        base_prompt = "A humorous meme style illustration related to artificial intelligence, comic style"
        top = top_caption.strip() if top_caption.strip() else "TOP TEXT"
        bot = bottom_caption.strip() if bottom_caption.strip() else "BOTTOM TEXT"

    w, h = 1024, 1024

    try:
        image = client.text_to_image(base_prompt, width=w, height=h)
        if not isinstance(image, Image.Image):
            image = Image.open(io.BytesIO(image)) if hasattr(image, "read") else Image.open(image)
    except Exception as e:
        image = Image.new("RGB", (w, h), BG_COLOR)
        d = ImageDraw.Draw(image)
        d.text((20, h // 2), f"Error: {e}", fill=ACCENT)

    image = overlay_meme_text(image, top.upper(), bot.upper())
    return image, f"✅ Meme generated! Top: '{top}' | Bottom: '{bot}'"


def overlay_meme_text(img, top_text, bottom_text):
    """Draw classic Impact-font style meme text with black outline."""
    img = img.convert("RGB")
    draw = ImageDraw.Draw(img)
    W, H = img.size

    font_size = max(24, int(W / 12))
    font = load_font(font_size)

    def wrap_text(text, max_chars):
        words = text.split()
        lines = []
        current = ""
        for word in words:
            if len(current + " " + word) <= max_chars:
                current = current + " " + word if current else word
            else:
                if current:
                    lines.append(current)
                current = word
        if current:
            lines.append(current)
        return lines

    def draw_text_with_outline(text, y, anchor="ma"):
        for dx in (-3, -2, -1, 1, 2, 3):
            for dy in (-3, -2, -1, 1, 2, 3):
                draw.text((W // 2 + dx, y + dy), text, fill="black", font=font, anchor=anchor)
        draw.text((W // 2, y), text, fill="white", font=font, anchor=anchor)

    # Top text
    max_chars = int(W / (font_size * 0.55))
    top_lines = wrap_text(top_text, max_chars)
    for i, line in enumerate(top_lines):
        draw_text_with_outline(line, int(H * 0.06) + i * (font_size + 4), anchor="ma")

    # Bottom text
    bot_lines = wrap_text(bottom_text, max_chars)
    total_h = len(bot_lines) * (font_size + 4)
    for i, line in enumerate(bot_lines):
        y = int(H * 0.94) - total_h + i * (font_size + 4) + font_size
        draw_text_with_outline(line, y, anchor="ma")

    return img


def load_template(template_name):
    if template_name and template_name in TEMPLATES:
        _, def_top, def_bot = TEMPLATES[template_name]
        return def_top, def_bot
    return "", ""


# --- UI -----------------------------------------------------------------------
CSS = """
#dispatch-header h1 {
    color: #FFFFFF; font-size: 2.2rem; margin: 0;
    background: linear-gradient(90deg, #1FE0E6 0%, #FFFFFF 60%);
    -webkit-background-clip: text; -webkit-text-fill-color: transparent;
}
#dispatch-header p { color: #1FE0E6; font-size: 1.05rem; margin: 6px 0 0 0; }
.dispatch-footer { text-align: center; color: #8A8F9C; font-size: 0.9rem; padding-top: 8px; }
"""

with gr.Blocks(
    title="Dispatch AI — Meme Generator",
    theme=gr.themes.Base(
        primary_hue="cyan",
        secondary_hue="cyan",
        neutral_hue="slate",
        font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui"],
    ).set(
        body_background_fill="#0A0F1A",
        body_background_fill_dark="#0A0F1A",
        body_text_color="#FFFFFF",
        body_text_color_dark="#FFFFFF",
        block_background_fill="#0E1424",
        block_background_fill_dark="#0E1424",
        block_border_color="#1FE0E6",
        block_border_width="1px",
        block_label_text_color="#1FE0E6",
        block_title_text_color="#1FE0E6",
        button_primary_background_fill="#1FE0E6",
        button_primary_background_fill_dark="#1FE0E6",
        button_primary_text_color="#0A0F1A",
        button_primary_border_color="#1FE0E6",
        input_background_fill="#0E1424",
        input_background_fill_dark="#0E1424",
        input_border_color="#1FE0E6",
        input_border_width="1px",
    ),
    css=CSS,
) as demo:
    with gr.Column(elem_id="dispatch-header"):
        gr.Markdown(
            """
            # Dispatch AI — Meme Generator
            Generate AI/tech memes with FLUX.1-schnell + PIL text overlay · Dispatch AI (FZE) · UAE
            """
        )

    with gr.Row():
        with gr.Column(scale=1):
            template_select = gr.Dropdown(
                list(TEMPLATES.keys()),
                label="Meme Template",
                value="This is Fine AI",
                info="Pick a template to auto-fill captions and base image",
            )
            load_template_btn = gr.Button("Load Template Captions", variant="secondary")
            top_caption = gr.Textbox(label="Top Caption", placeholder="TOP TEXT", lines=1)
            bottom_caption = gr.Textbox(label="Bottom Caption", placeholder="BOTTOM TEXT", lines=1)
            random_topic_btn = gr.Button("🎲 Random AI Topic", variant="secondary")
            topic_box = gr.Textbox(label="Random AI Topic", interactive=True, lines=1)
            generate_btn = gr.Button("🤣 Generate Meme", variant="primary")
        with gr.Column(scale=2):
            output_image = gr.Image(label="Generated Meme", type="pil", show_download_button=True)
            status_box = gr.Textbox(label="Status", interactive=False)

    # Events
    load_template_btn.click(load_template, inputs=template_select, outputs=[top_caption, bottom_caption])
    random_topic_btn.click(random_topic, outputs=topic_box)
    generate_btn.click(
        generate_meme,
        inputs=[top_caption, bottom_caption, template_select],
        outputs=[output_image, status_box],
    )

    gr.Markdown(
        """
        <div class="dispatch-footer">
        © 2026 Dispatch AI (FZE) · UAE · License 10818 · Model: FLUX.1-schnell · Text overlay via Pillow
        </div>
        """
    )

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