import gradio as gr
import os
from huggingface_hub import InferenceClient, __version__ as hf_version
import random
from typing import Generator, Dict, List, Tuple, Optional
import logging  # Added logging for better debugging

# Configure logging with DEBUG level and add version info
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logging.debug(f"Using huggingface_hub version: {hf_version}")

# Get token from environment variable
hf_token = os.environ.get("HF_TOKEN")
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta", token=hf_token)

# Story genres with genre-specific example prompts
GENRE_EXAMPLES = {
    "fairy tale": [
        "I follow the shimmer of fairy dust into a hidden forest",
        "A tiny dragon appears at my window, asking for help to find its mother",
        "A friendly witch invites me into her cozy cottage, offering a warm cup of tea"
    ],
    "fantasy": [
        "I enter the ancient forest seeking the wizard's tower",
        "I approach the dragon cautiously with my shield raised",
        "I try to bargain with the elven council for safe passage"
    ],
    "sci-fi": [
        "I investigate the strange signal coming from the abandoned planet",
        "I negotiate with the alien ambassador about the peace treaty",
        "I try to repair my damaged spacecraft before oxygen runs out"
    ],
    "mystery": [
        "I examine the crime scene for overlooked evidence",
        "I question the nervous butler about the night of the murder",
        "I follow the suspicious figure through the foggy streets"
    ],
    "horror": [
        "I read the forbidden text while the candles flicker",
        "I hide under the bed as footsteps approach",
        "I investigate the strange noises coming from the attic"
    ],
    "western": [
        "I challenge the outlaw to a duel at high noon",
        "I track the bandits through the desert canyon",
        "I enter the saloon looking for information"
    ],
    "cyberpunk": [
        "I jack into the corporate mainframe to steal data",
        "I hide in the neon-lit alleyway from corporate security",
        "I meet my mysterious client in the underground bar"
    ],
    "historical": [
        "I join the resistance against the occupying forces",
        "I navigate the dangerous politics of the royal court",
        "I set sail on a voyage to discover new lands"
    ],
    "post-apocalyptic": [
        "I scavenge the abandoned shopping mall for supplies",
        "I navigate through the radioactive zone using my old map",
        "I hide from the approaching group of raiders"
    ],
    "steampunk": [
        "I pilot my airship through the lightning storm",
        "I present my new invention to the Royal Academy",
        "I sneak aboard the emperor's armored train"
    ]
}

# 2. Add constants at the top for magic numbers
MAX_HISTORY_LENGTH = 20
MEMORY_WINDOW = 5  # Reduced from 10 to limit context
MAX_TOKENS = 1024  # Reduced from 2048 for faster responses
TEMPERATURE = 0.7  # Slightly reduced for faster convergence
TOP_P = 0.95
MIN_RESPONSE_LENGTH = 100  # Reduced from 200 for quicker display

def get_examples_for_genre(genre):
    """Get example prompts specific to the selected genre"""
    return GENRE_EXAMPLES.get(genre, GENRE_EXAMPLES["fantasy"])

def get_enhanced_system_prompt(genre=None):
    """Generate a detailed system prompt with optional genre specification"""
    selected_genre = genre or "fantasy"
    system_message = f"""You are an interactive storyteller creating an immersive {selected_genre} choose-your-own-adventure story.
For each response you MUST:
1. Write 100-200 words describing the scene, using vivid sensory details
2. Always use second-person perspective ("you", "your") to maintain reader immersion
3. Include dialogue or your character's thoughts that reveal personality and motivations
4. Create a strong sense of atmosphere appropriate for {selected_genre}
5. End with EXACTLY THREE numbered choices and NOTHING ELSE AFTER THEM:
1. [Complete sentence in second-person starting with a verb]
2. [Complete sentence in second-person starting with a verb]
3. [Complete sentence in second-person starting with a verb]
CRITICAL RULES:
- Provide only ONE set of three choices at the very end of your response
- Never continue the story after giving choices
- Never provide additional choices
- Keep all narrative before the choices
- End every response with exactly three numbered options
- Each choice must start with "You" followed by a verb
Remember: The story continues ONLY when the player makes a choice."""
    return system_message

def create_story_summary(chat_history):
    """Create a concise summary of the story so far if the history gets too long"""
    if len(chat_history) <= 2:
        return None
    
    story_text = ""
    for i in range(0, len(chat_history), 2):
        if i+1 < len(chat_history):
            story_text += f"User: {chat_history[i]}\nStory: {chat_history[i+1]}\n\n"
    
    summary_instruction = {
        "role": "system",
        "content": "The conversation history is getting long. Please create a brief summary of the key plot points and character development so far to help maintain context without exceeding token limits."
    }
    return summary_instruction

# Modified function for proper Gradio format (lists)
def respond(message: str, chat_history: List[Tuple[str, str]], genre: Optional[str] = None, use_full_memory: bool = True) -> Tuple[str, List[Tuple[str, str]]]:
    """Generate a response based on the current message and conversation history."""
    if not message.strip():
        return "", chat_history
    
    try:
        # Start with system prompt
        api_messages = [{"role": "system", "content": get_enhanced_system_prompt(genre)}]
        logging.debug(f"System Message: {api_messages[0]}")
        
        # Add chat history - convert from tuples to API format
        if chat_history and use_full_memory:
            for user_msg, bot_msg in chat_history[-MEMORY_WINDOW:]:
                api_messages.extend([
                    {"role": "user", "content": str(user_msg)},
                    {"role": "assistant", "content": str(bot_msg)}
                ])
            logging.debug(f"Chat History Messages: {api_messages[1:]}")
        
        # Add current message
        api_messages.append({"role": "user", "content": str(message)})
        logging.debug(f"Final Message List: {api_messages}")

        # Make API call without timeout
        logging.debug("Making API call...")
        response = client.chat_completion(
            messages=api_messages,
            max_tokens=MAX_TOKENS,
            temperature=TEMPERATURE,
            top_p=TOP_P
        )
        logging.debug("API call completed")
        
        # Extract response
        bot_message = response.choices[0].message.content
        logging.debug(f"Bot Response: {bot_message[:100]}...")
        
        # Update history using tuple format [(user_msg, bot_msg), ...]
        updated_history = list(chat_history)  # Create a copy
        updated_history.append((message, bot_message))  # Add as tuple
        return "", updated_history
    
    except Exception as e:
        logging.error("Error in respond function", exc_info=True)
        error_msg = f"Story magic temporarily interrupted. Please try again. (Error: {str(e)})"
        return "", list(chat_history) + [(message, error_msg)]

def save_story(chat_history):
    """Convert chat history to markdown and return as downloadable file"""
    if not chat_history:
        return gr.File.update(value=None)
    
    try:
        story_text = "# My Adventure\n\n"
        for user_msg, bot_msg in chat_history:
            story_text += f"**Player:** {user_msg}\n\n"
            story_text += f"**Story:** {bot_msg}\n\n---\n\n"
        
        # Create temporary file
        temp_file = "my_story.md"
        with open(temp_file, "w", encoding="utf-8") as f:
            f.write(story_text)
        
        return temp_file
        
    except Exception as e:
        logging.error(f"Error saving story: {e}")
        return gr.File.update(value=None)

# Add this function to get a custom avatar image URL
def get_storyteller_avatar_url():
    """Get a URL for the storyteller avatar from a free image service"""
    # Using an external wizard avatar image
    return "https://api.dicebear.com/7.x/bottts/svg?seed=wizard&backgroundColor=b6e3f4&eyes=bulging"

# Add this before your gr.Blocks definition
custom_css = """
.compact-file-output > div {
    min-height: 0 !important;
    padding: 0 !important;
}
.compact-file-output .file-preview {
    margin: 0 !important;
    display: flex;
    align-items: center;
}
.compact-btn {
    padding: 0.5rem !important;
    min-height: 0 !important;
    height: auto !important;
    line-height: 1.2 !important;
}
"""

with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
    # Header section with improved instructions
    gr.Markdown("""
    # 🔮 AI Story Studio
    **Collaborate with AI to craft your own adventure, one scene at a time.**  
    Pick a genre, start with a prompt or write your own, and guide the story with your choices.  
    > **Tip:** The more detail you provide, the deeper the story becomes.
    """)
    
    wizard_avatar = get_storyteller_avatar_url()
    
    with gr.Row():
        with gr.Column(scale=3):
            # Chat window + user input - USING LIST FORMAT
            chatbot = gr.Chatbot(
                height=400,
                bubble_full_width=True,
                show_copy_button=True,
                avatar_images=(None, wizard_avatar),
                container=True,
                scale=1,
                min_width=800,
                value=[],  # Empty list for messages
                render=True
            )
            msg = gr.Textbox(
                placeholder="Describe your next move...",
                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=list(GENRE_EXAMPLES.keys()),
                label="Story Genre",
                info="Choose the theme of your next adventure",
                value="fantasy"
            )
            full_memory = gr.Checkbox(
                label="Full Story Memory",
                value=True,
                info="When enabled, the AI tries to remember the entire story. If disabled, only the last few exchanges are used."
            )
            
            gr.Markdown("## Story Starters")
            
            # Create three placeholder buttons for story starters
            starter_btn1 = gr.Button("Starter 1", scale=1, min_width=250, elem_classes="compact-btn")
            starter_btn2 = gr.Button("Starter 2", scale=1, min_width=250, elem_classes="compact-btn")
            starter_btn3 = gr.Button("Starter 3", scale=1, min_width=250, elem_classes="compact-btn")
            starter_buttons = [starter_btn1, starter_btn2, starter_btn3]

            # Simplified update function
            def update_starter_buttons(selected_genre):
                examples = get_examples_for_genre(selected_genre)
                results = []
                for i in range(3):
                    if i < len(examples):
                        results.append(examples[i])
                    else:
                        results.append("")
                return tuple(results)

            # New direct handler for starter clicks
            def use_starter(starter_text: str, history: List[Tuple[str, str]], selected_genre: str, memory_flag: bool) -> Tuple[str, List[Tuple[str, str]]]:
                """Handle starter button clicks with proper message formatting"""
                if not starter_text:
                    return "", history
                
                try:
                    # Use the respond function for consistent handling
                    _, updated_history = respond(
                        message=starter_text,
                        chat_history=history,
                        genre=selected_genre,
                        use_full_memory=memory_flag
                    )
                    return "", updated_history
                                
                except Exception as e:
                    error_msg = f"Story magic temporarily interrupted. Please try again. (Error: {str(e)})"
                    return "", list(history) + [(starter_text, error_msg)]

            # Simplified button connections
            for starter_button in starter_buttons:
                starter_button.click(
                    fn=use_starter,
                    inputs=[starter_button, chatbot, genre, full_memory],
                    outputs=[msg, chatbot],
                    queue=True
                )

            # Update buttons when genre changes
            genre.change(
                fn=update_starter_buttons,
                inputs=[genre],
                outputs=starter_buttons
            )

    # Handler for user input
    msg.submit(
        fn=respond,
        inputs=[msg, chatbot, genre, full_memory],
        outputs=[msg, chatbot]
    )
    submit.click(
        fn=respond,
        inputs=[msg, chatbot, genre, full_memory],
        outputs=[msg, chatbot]
    )
    
    # Clear the chatbot for a new adventure
    clear.click(lambda: [], None, chatbot, queue=False)
    clear.click(lambda: "", None, msg, queue=False)
    
    # "Download My Story" row with improved layout
    with gr.Row(equal_height=True):  # Force equal height for all children
        # Use Column for the button to control width
        with gr.Column(scale=4):
            save_btn = gr.Button("Download My Story", variant="secondary", size="lg")
        
        # Use Column for the file output with matching height
        with gr.Column(scale=1):
            story_output = gr.File(
                label=None,  # Remove the label that adds extra height
                file_count="single",
                file_types=[".md"],
                interactive=False,
                visible=True,
                elem_classes="compact-file-output"  # Optional: for custom CSS styling
            )
    
    # Connect the save button to the save_story function
    save_btn.click(
        fn=save_story,
        inputs=[chatbot],
        outputs=story_output,
        queue=False  # Process immediately
    )

    # Initialize buttons with default fantasy genre examples
    initial_examples = get_examples_for_genre("fantasy")
    initial_button_data = tuple(
        initial_examples[i] if i < len(initial_examples) else ""
        for i in range(3)
    )
    
    # Update button text on page load
    demo.load(
        fn=lambda: initial_button_data,
        outputs=starter_buttons,
        queue=False
    )

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