import gradio as gr
from transformers import pipeline

# Load the model using pipeline
generator = pipeline("text-generation", model="distilgpt2")

def generate_story(prompt, temperature, top_p, max_length):
    # Clamp temperature to at least 0.01 to avoid division by zero
    temperature = max(0.01, float(temperature))
    
    # Generate the text
    result = generator(
        prompt,
        max_length=int(max_length),
        temperature=temperature,
        top_p=float(top_p),
        do_sample=True,
        num_return_sequences=1,
        repetition_penalty=1.2,        # <--- FIX 1: Taxes words it has already used
        no_repeat_ngram_size=2,        # <--- FIX 2: Prevents any 2-word phrase from repeating
        truncation=True
    )
    
    return result[0]["generated_text"]

# Build the Gradio interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Game Narrative Generator (Patched Version)")
    gr.Markdown("Designed for writing game stories — character dialogue, quest descriptions, item lore, and scene-setting.")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(label="Story Prompt", lines=3)
            temperature = gr.Slider(0.1, 2.0, value=0.7, label="Temperature (Creativity)")
            top_p = gr.Slider(0.1, 1.0, value=0.9, label="Top-p (Word Diversity)")
            max_length = gr.Slider(20, 200, value=100, step=1, label="Max Length")
            
            generate_btn = gr.Button("Generate Narrative", variant="primary")
            
        with gr.Column():
            output_text = gr.Textbox(label="Generated Output", lines=10)
            
    gr.Examples(
        examples=[["The warrior entered the dungeon and"],
            ["You found a legendary sword. Its description reads:"],["The village elder whispered the secret of the forest:"],
            ["QUEST LOG: Your mission is to"],["The final boss appeared, and it was"]
        ],
        inputs=prompt_input
    )
    
    generate_btn.click(
        fn=generate_story,
        inputs=[prompt_input, temperature, top_p, max_length],
        outputs=output_text
    )

demo.launch()