import re
import gradio as gr
import torch
from transformers import pipeline

# SmolLM2-360M fits easily in 16GB RAM and runs on CPU
pipe = pipeline(
    "text-generation",
    model="HuggingFaceTB/SmolLM2-360M-Instruct",
    torch_dtype=torch.float32,
    device="cpu",
)

SYSTEM = "You are an RPG narrator. Convert the user's life situation into a hero's journey."

def generate_hero(name, goal, challenge):
    name = name.strip() or "Stranger"
    messages = [
        {"role": "system", "content": SYSTEM},
        {"role": "user", "content": (
            f"Name: {name}\nGoal: {goal}\nBiggest Challenge: {challenge}\n\n"
            "Respond in exactly this format:\n"
            "Hero Class: <class>\n"
            "Main Quest: <quest>\n"
            "Story Introduction: <2-3 sentence intro>"
        )},
    ]

    out = pipe(messages, max_new_tokens=200, do_sample=False)
    text = out[0]["generated_text"][-1]["content"].strip()

    def extract(label):
        m = re.search(rf"{label}:\s*(.+?)(?=\n[A-Z]|$)", text, re.DOTALL)
        return m.group(1).strip() if m else ""

    return extract("Hero Class"), extract("Main Quest"), extract("Story Introduction")

with gr.Blocks(title="Hero Generator") as app:
    gr.Markdown("# ⚔️ Hero Generator")

    with gr.Row():
        name      = gr.Textbox(label="Name",              placeholder="Your name")
        goal      = gr.Textbox(label="Goal",              placeholder="e.g. learn to code")
    challenge     = gr.Textbox(label="Biggest Challenge", placeholder="e.g. lack of focus")

    btn = gr.Button("Generate Hero", variant="primary")

    hero_class = gr.Textbox(label="Hero Class")
    quest      = gr.Textbox(label="Main Quest")
    story      = gr.Textbox(label="Story Introduction", lines=5)

    btn.click(generate_hero, inputs=[name, goal, challenge],
              outputs=[hero_class, quest, story])

app.launch()