import gradio as gr
from transformers import pipeline

# Initialize the o3 mini model (update the model identifier as needed)
o3_mini = pipeline("text2text-generation", model="o3mini-template-generator")

def generate_template(ref1, ref2, ref3, ref4, ref5):
    # Collect non-empty references
    references = [ref.strip() for ref in [ref1, ref2, ref3, ref4, ref5] if ref.strip()]
    if not references:
        return "Please provide at least one piece of reference copy."
    
    # Combine all references into one prompt text
    combined_references = "\n\n".join(references)
    
    # Create a prompt that instructs the model to analyze the provided references
    # and generate a structured template following the e-commerce copywriting guidelines.
    prompt = (
        "Analyze the following reference copy examples and generate a structured e-commerce "
        "copywriting template that includes a header, template blocks, and control structures. "
        "Follow this guide document:\n\n"
        "E-commerce Copywriting Template Documentation (refer to internal guide for details)\n\n"
        "Reference Copy:\n"
        f"{combined_references}\n\n"
        "Generate the final structured template output."
    )
    
    # Call the o3 mini model with the prompt
    result = o3_mini(prompt, max_length=512, truncation=True)
    
    # The model returns a list of dictionaries; extract the generated text.
    template_output = result[0]['generated_text']
    return template_output

iface = gr.Interface(
    fn=generate_template,
    inputs=[
        gr.Textbox(label="Reference Copy 1", lines=3, placeholder="Paste your reference copy here"),
        gr.Textbox(label="Reference Copy 2 (optional)", lines=3, placeholder="Paste your reference copy here"),
        gr.Textbox(label="Reference Copy 3 (optional)", lines=3, placeholder="Paste your reference copy here"),
        gr.Textbox(label="Reference Copy 4 (optional)", lines=3, placeholder="Paste your reference copy here"),
        gr.Textbox(label="Reference Copy 5 (optional)", lines=3, placeholder="Paste your reference copy here")
    ],
    outputs=gr.Textbox(label="Structured Template"),
    title="E-commerce Template Generator using O3 Mini",
    description=(
        "Provide up to five pieces of reference copy. This app calls the o3 mini model "
        "to analyze your inputs and generate a structured template based on our e-commerce copywriting guidelines."
    )
)

iface.launch()