import gradio as gr
from transformers import pipeline

# Load the text generation pipeline
generator = pipeline("text-generation", model="gpt2", device=-1)

# Local generation function (no API)
def generate_text(category, theme, tone, length):
    # Fallback defaults
    category = category or "story"
    tone = tone or "neutral"
    length = length or "short"
    theme = theme.strip() or "an interesting idea"

    # Determine length in tokens
    length_map = {
        "short": 100,
        "medium": 200,
        "long": 300
    }
    max_length = length_map.get(length.lower(), 50)

    # Build prompt
    prompt = (
        #f"Write a {length.lower()} {tone.lower()} {category.lower()} about {theme}."
        f"Write a {length.lower()} and {tone.lower()} {category.lower()} titled '{theme}'. It should have a beginning, middle, and end. Use vivid, imaginative language."
        f"'The {theme.title()}'. Write it in a creative and expressive style.\n\n"
    )

    # add specific instruction based on category
    if category.lower() == "poem":
        prompt += "Make sure it's in poetic form, with vivid imagery and emotion."
    elif category.lower() == "story":
        prompt += "Start with an engaging opening, include a conflict, and resolve it clearly."


    #line break for clarify
    prompt +="\n\n"
    
    try:
        result = generator(prompt, max_length=max_length, num_return_sequences=1)
        return result[0]["generated_text"]
    except Exception as e:
        return f"❌ Exception: {str(e)}"

# Create Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## ✨🧚✨ Story Spark Generator (Local ) ✨🧚✨")

    category = gr.Dropdown(["Story", "Poem"], label="Category")
    theme = gr.Textbox(label="Theme", placeholder="e.g. friendship, time travel, lost love")
    tone = gr.Dropdown(["Happy", "Sad", "Funny", "Dark", "Inspiring"], label="Tone")
    length = gr.Dropdown(["Short", "Medium", "Long"], label="Length")
    generate_button = gr.Button("Generate")
    output = gr.Textbox(label="Generated Content", lines=10)

    generate_button.click(fn=generate_text,
                          inputs=[category, theme, tone, length],
                          outputs=output)

demo.launch()