import gradio as gr
from huggingface_hub import InferenceClient
import os
from datetime import datetime
import json

# Initialize clients
client = InferenceClient("EvanZhouDev/open-genmoji", token=os.getenv("HUGGINGFACE_API_TOKEN"))
llm = InferenceClient("Qwen/Qwen2.5-72B-Instruct")

# Ensure output directories exist
os.makedirs("outputs/images", exist_ok=True)

# Define the process function
def process(prompt, steps, seed, quantize, guidance, width, height):
    print(f"Prompt: {prompt}")
    messages = [
        {
            "role": "system",
            "content": (
                "You are helping create a prompt for a Emoji generation image model. An emoji must be easily "
                "interpreted when small so details must be exaggerated to be clear. Your goal is to use descriptions "
                "to achieve this.\n\nYou will receive a user description, and you must rephrase it to consist of "
                "short phrases separated by periods, adding detail to everything the user provides.\n\nAdd describe "
                "the color of all parts or components of the emoji. Unless otherwise specified by the user, do not "
                "describe people. Do not describe the background of the image. Your output should be in the format:\n\n"
                "```emoji of {description}. {addon phrases}. 3D lighting. no cast shadows.```\n\nThe description "
                "should be a 1 sentence of your interpretation of the emoji. Then, you may choose to add addon phrases."
                " You must use the following in the given scenarios:\n\n- \"cute.\": If generating anything that's not "
                "an object, and also not a human\n- \"enlarged head in cartoon style.\": ONLY animals\n- \"head is "
                "turned towards viewer.\": ONLY humans or animals\n- \"detailed texture.\": ONLY objects\n\nFurther "
                "addon phrases may be added to ensure the clarity of the emoji."
            ),
        },
        {"role": "user", "content": prompt},
    ]

    completion = llm.chat_completion(messages, max_tokens=100)
    response = completion.get("choices")[0].get("message").get("content").replace("```", "").replace("\n", "")
    print(f"Refined Prompt: {response}")

    time = datetime.now().strftime("%Y%m%d%H%M%S")
    image = client.text_to_image(response, steps=steps, seed=seed, quantize=quantize, guidance=guidance, width=width, height=height)
    image.save(f"outputs/images/{time}.png")

    with open(f"outputs/{time}.json", "w") as f:
        json.dump({"prompt": prompt, "refined_prompt": response, "image": f"outputs/images/{time}.png"}, f)

    return image

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Emoji Generator with Customizable Parameters")
    
    # Input fields
    with gr.Row():
        prompt_input = gr.Textbox(label="Enter a prompt")
        steps_input = gr.Slider(label="Steps", minimum=1, maximum=50, value=20, step=1)
    with gr.Row():
        seed_input = gr.Number(label="Seed", value=2, precision=0)
        quantize_input = gr.Slider(label="Quantize", minimum=1, maximum=16, value=8, step=1)
    with gr.Row():
        guidance_input = gr.Slider(label="Guidance", minimum=1.0, maximum=10.0, value=5.0, step=0.1)
        width_input = gr.Slider(label="Width", minimum=256, maximum=2048, value=1280, step=64)
        height_input = gr.Slider(label="Height", minimum=256, maximum=2048, value=640, step=64)

    # Output
    image_output = gr.Image(label="Generated Image")
    
    # Button to generate the image
    generate_button = gr.Button("Generate Image")
    
    # Define button click behavior
    generate_button.click(
        fn=process,
        inputs=[prompt_input, steps_input, seed_input, quantize_input, guidance_input, width_input, height_input],
        outputs=image_output,
    )

# Launch the app
if __name__ == "__main__":
    demo.launch(show_error=True)