import streamlit as st
from transformers import pipeline

# Set page configuration
st.set_page_config(page_title="AI Puzzle Solver", layout="centered")

# Title and description
st.title("🧠 AI Puzzle Solver Game")
st.write("Enter a riddle or puzzle, and the AI will try to solve it!")

# Load a small, public model that works everywhere
@st.cache_resource
def load_model():
    return pipeline("text-generation", model="sshleifer/tiny-gpt2", max_new_tokens=50)

generator = load_model()

# User input
puzzle_input = st.text_area("Type your riddle or puzzle:")

# Solve button
if st.button("Solve Puzzle"):
    if puzzle_input.strip() == "":
        st.warning("Please enter a puzzle.")
    else:
        with st.spinner("AI is thinking..."):
            prompt = f"Riddle: {puzzle_input}\nAnswer:"
            result = generator(prompt)[0]['generated_text']
            answer = result.split("Answer:")[-1].strip()
            st.success("✅ AI's Best Guess:")
            st.markdown(f"**{answer}**")