import gradio as gr
from huggingface_hub import InferenceClient
import random

# Initialize the inference client
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Story genres and themes for variety
STORY_GENRES = [
    "fantasy", "sci-fi", "mystery", "horror", "historical", 
    "cyberpunk", "western", "post-apocalyptic", "steampunk"
]

# Define character archetypes for the assistant to use
CHARACTER_ARCHETYPES = [
    "the reluctant hero", "the wise mentor", "the cunning rogue",
    "the loyal companion", "the mysterious stranger", "the noble warrior"
]

def get_enhanced_system_prompt(genre=None):
    """Generate a detailed system prompt with optional genre specification"""
    selected_genre = genre or random.choice(STORY_GENRES)
    
    system_message = f"""You are an interactive storyteller creating an immersive {selected_genre} choose-your-own-adventure story.

For each response:
1. Provide vivid sensory descriptions of the scene, environment, and characters
2. Include meaningful dialogue or internal monologue that reveals character motivations
3. End with exactly 3 different possible actions or decisions, each offering a distinct path
4. Maintain consistent world-building and character development
5. Incorporate appropriate atmosphere and tone for a {selected_genre} setting
6. Remember previous choices to create a coherent narrative arc

Format your three options as:
- Option 1: [Complete sentence describing a possible action]
- Option 2: [Complete sentence describing a possible action]
- Option 3: [Complete sentence describing a possible action]

Keep responses engaging but concise (200-300 words maximum). If the user's input doesn't clearly indicate a choice, interpret their intent and move the story forward in the most logical direction."""
    
    return system_message

def respond(message, history, genre=None, memory_length=5):
    """Generate a response based on the current message and conversation history"""
    # Use a more detailed system prompt
    system_message = get_enhanced_system_prompt(genre)
    
    # Initialize messages with system prompt
    messages = [{"role": "system", "content": system_message}]
    
    # Add limited history to prevent token overflow (only use recent exchanges)
    for user_msg, bot_msg in history[-memory_length:]:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})
    
    # Add current message
    messages.append({"role": "user", "content": message})
    
    # Special handling for story initialization
    if len(history) == 0 or message.lower() in ["start", "begin", "begin my adventure"]:
        # Add a specific instruction for starting a new story
        messages.append({
            "role": "system", 
            "content": f"Begin a new {genre or random.choice(STORY_GENRES)} adventure with an intriguing opening scene. Introduce the protagonist without assuming too much about them, allowing the user to shape the character."
        })
    
    # Generate and stream response
    response = ""
    try:
        for message in client.chat_completion(
            messages,
            max_tokens=512,
            stream=True,
            temperature=0.7,
            top_p=0.95,
        ):
            token = message.choices[0].delta.content
            response += token
            yield response
    except Exception as e:
        error_message = f"Story magic temporarily interrupted. Please try again. (Error: {str(e)})"
        yield error_message

# Create interface with additional customization options
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🔮 Interactive Story Adventure")
    gr.Markdown("Immerse yourself in an interactive story where your choices shape the narrative.")
    
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                height=500,
                bubble_full_width=False,
                show_copy_button=True,
                avatar_images=(None, "🧙"),
            )
            msg = gr.Textbox(
                placeholder="Describe what you want to do next in the story...",
                container=False,
                scale=4,
            )
            
            with gr.Row():
                submit = gr.Button("Continue Story", variant="primary")
                clear = gr.Button("Start New Adventure")
                
        with gr.Column(scale=1):
            gr.Markdown("## Adventure Settings")
            genre = gr.Dropdown(
                choices=STORY_GENRES,
                label="Story Genre",
                info="Select a genre for your next adventure",
                value="fantasy"
            )
            
            examples = gr.Examples(
                examples=[
                    ["Begin my adventure"],
                    ["I explore the mysterious cave"],
                    ["I approach the merchant to ask about the rumors"],
                    ["I decide to climb the tower"],
                    ["I hide in the shadows and observe"]
                ],
                inputs=msg
            )
    
    msg.submit(respond, [msg, chatbot, genre], chatbot)
    submit.click(respond, [msg, chatbot, genre], chatbot)
    clear.click(lambda: None, None, chatbot, queue=False)
    clear.click(lambda: "", None, msg, queue=False)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)