import os
import time
import json
import gradio as gr
import requests
from groq import Groq
from dotenv import load_dotenv
from io import BytesIO
from PIL import Image

load_dotenv()

# Setup clients and credentials
IMGFLIP_USERNAME = os.getenv("IMGFLIP_USERNAME")
IMGFLIP_PASSWORD = os.getenv("IMGFLIP_PASSWORD")
client = Groq(api_key=os.getenv("GROQ_API_KEY"))

# Load memes from Imgflip API
def list_memes():
    # r = requests.get("https://api.imgflip.com/get_memes")
    # data = r.json()
    # if not data["success"]:
    #     return {}
    data = json.loads(open("./metadata.json","r").read())
    return {meme["name"]: {"id": meme["id"], "url": meme["url"]} for meme in data["data"]["memes"]}

MEME_DATABASE = list_memes()

# Llama-4-Scout selects the meme template + captions
def pick_meme_and_params(prompt: str) -> dict:
    meme_names = list(MEME_DATABASE.keys())
    meme_list_string = "\n".join(f"- {name}" for name in meme_names)

    system_message = "You are a JSON-output assistant that selects the best meme template and generates top and bottom captions. Make them funny and creative"
    user_message = (
        f"Here is the list of available meme templates:\n{meme_list_string}\n\n"
        f"Based on this list, choose the most suitable meme template for the prompt: '{prompt}'.\n\n"
        f"Return your answer in this JSON format:\n"
        "{\n"
        "\"template_name\": \"template name from the list\",\n"
        "\"top_text\": \"top caption\",\n"
        "\"bottom_text\": \"bottom caption\",\n"
        "\"font\": \"optional font (default to impact)\",\n"
        "\"max_font_size\": integer between 10 and 100 (optional)\n"
        "}"
        "Only return JSON, not extra information, no delimiter, only JSON"
    )
    print("preresponse")
    response = client.chat.completions.create(
        # model="meta-llama/llama-4-scout-17b-16e-instruct",
        model="llama3-70b-8192",
        messages=[{"role": "system", "content": system_message}, {"role": "user", "content": user_message}],
        # response_format="json"
    )
    # print(response.choices[0].message.content)
    json_response = json.loads(response.choices[0].message.content)
    return json_response

# Call Imgflip API to generate the meme
# def make_meme(config: dict) -> str:
#     template = MEME_DATABASE.get(config["template_name"])
#     if not template:
#         return f"Error: Meme template '{config['template_name']}' not found."

#     payload = {
#         "template_id": template["id"],
#         "username": IMGFLIP_USERNAME,
#         "password": IMGFLIP_PASSWORD,
#         "text0": config["top_text"],
#         "text1": config["bottom_text"]
#     }

#     if config.get("font"):
#         payload["font"] = config["font"]
#     if config.get("max_font_size"):
#         payload["max_font_size"] = config["max_font_size"]

#     r = requests.post("https://api.imgflip.com/caption_image", data=payload)
#     d = r.json()
#     return d["data"]["url"] if d.get("success") else f"Error: {d.get('error_message')}"

def make_meme(config: dict):
    template = MEME_DATABASE.get(config["template_name"])
    if not template:
        raise ValueError(f"Meme template '{config['template_name']}' not found.")

    payload = {
        "template_id": template["id"],
        "username": IMGFLIP_USERNAME,
        "password": IMGFLIP_PASSWORD,
        "text0": config["top_text"],
        "text1": config["bottom_text"]
    }

    if config.get("font"):
        payload["font"] = config["font"]
    if config.get("max_font_size"):
        payload["max_font_size"] = config["max_font_size"]

    r = requests.post("https://api.imgflip.com/caption_image", data=payload)
    d = r.json()
    print(d)
    # with open("response.json","w") as f:
    #     json.dump(d,f)
    if not d.get("success"):
        raise ValueError(f"Error from Imgflip API: {d.get('error_message')}")

    meme_url = d["data"]["url"]

    # Fetch the generated image from the returned URL
    img_response = requests.get(meme_url)
    if img_response.status_code != 200:
        raise ValueError("Failed to fetch the generated meme image.")

    # Convert image bytes to PIL image for Gradio
    image = Image.open(BytesIO(img_response.content))
    image.save("./img.png")
    return image

# Full meme generation pipeline
def generate_meme(prompt):
    try:
        config = pick_meme_and_params(prompt)
        print(type(config))
        print(config)
        return make_meme(config)
    except Exception as e:
        return f"Error generating meme: {e}"

# Gradio + MCP interface

with gr.Blocks() as demo:
    gr.Markdown("# 🖼️ Meme Context Protocol Tool")
    gr.Markdown("""
    **Your personal meme genie!**  
    Drop a prompt, and this meme conjurer uses magic to pick the perfect template and cook up hilarious captions faster than you can say "dank."  
    It’s the ultimate meme-making mischief machine – pun-powered, LLM-fueled, and chaos-approved.  
    Memes so fresh, they might just slap themselves.
    """)
    prompt_in = gr.Textbox(label="Enter your meme idea")
    out_img = gr.Image(label="Generated Meme")
    btn = gr.Button("Generate Meme")

    btn.click(fn=generate_meme, inputs=prompt_in, outputs=out_img)

if __name__ == "__main__":
    demo.launch(mcp_server=True, ssr_mode=False)