import os
import sys
import gradio as gr

sys.path.append(os.path.dirname(os.path.dirname(__file__)))
from common import (
    PoetryModel, load_prompt_template, fill_prompt_template, 
    filter_output, create_fill_in_blank, select_words_to_blank,
    load_themes, load_interests, check_input_safety
)


def create_app(model: PoetryModel):
    """Create Poem Generator Gradio app with intelligent word selection and position-aware checking."""
    
    prompt_path = os.path.join(os.path.dirname(__file__), 'prompts.md')
    prompt_template = load_prompt_template(prompt_path)
    
    themes_path = os.path.join(os.path.dirname(__file__), 'themes.md')
    interests_path = os.path.join(os.path.dirname(__file__), 'interests.md')
    
    THEMES = load_themes(themes_path)
    INTERESTS = load_interests(interests_path)
    
    MAX_INTERESTS = 3
    
    def generate_poem(age, length, difficulty, selected_theme, selected_interests):
        # Convert age to int if it's a string
        try:
            age = int(age)
        except (ValueError, TypeError):
            age = 9 # Default
            
        # Build theme and interests strings
        theme_text = selected_theme if selected_theme else "general"
        interests_text = ", ".join(selected_interests) if selected_interests else "various things"
        
        prompt = fill_prompt_template(
            prompt_template,
            age=age,
            length=length,
            theme=theme_text,
            interests=interests_text
        )
        
        # Generate poem with debugging output
        print(f"\n{'='*70}\n🎨 GENERATING POEM\n{'='*70}")
        print(f"Age: {age} | Length: {length} | Difficulty: {difficulty}")
        print(f"Themes: {theme_text}")
        print(f"Interests: {interests_text}")
        print(f"\nPrompt sent to model:\n{'-'*70}\n{prompt}\n{'-'*70}")
        
        poem = model.generate(prompt, max_tokens=350, temperature=0.7)
        poem = filter_output(poem)
        
        print(f"\n✓ Generated Poem:\n{'-'*70}\n{poem}\n{'-'*70}")
        
        # Use AI to intelligently select words to blank
        selected_words = select_words_to_blank(poem, difficulty, age, model)
        
        # Create blanks with exact positions
        blanked, correct_answers, positions, all_words = create_fill_in_blank(poem, difficulty, selected_words)
        
        print(f"\n📊 Blank Information:")
        print(f"Number of blanks: {len(positions)}")
        print(f"Positions: {positions}")
        print(f"Correct answers: {correct_answers}")
        print(f"\nBlanked poem:\n{'-'*70}\n{blanked}\n{'='*70}\n")
        
        # Format poem with line breaks for better display
        poem_lines = blanked.split('. ')
        # Use double line break for clear separation
        formatted_blanked = "<br><br>".join([line.strip() + ('.' if i < len(poem_lines)-1 else '') for i, line in enumerate(poem_lines) if line.strip()])
        
        # Create shuffled word bank
        import random
        word_bank = correct_answers.copy()
        random.shuffle(word_bank)
        word_bank_html = "**Word Bank:** " + " | ".join([f"<span style='background: #e0e7ff; padding: 4px 12px; border-radius: 12px; color: #667eea; font-weight: 600;'>{word}</span>" for word in word_bank])
        
        # Create input fields for each blank
        blank_components = []
        for i in range(len(correct_answers)):
            blank_components.append(gr.update(
                value="",
                visible=True,
                placeholder=f"Blank {i+1}",
                interactive=True,
                type = 'text'
            ))
        # Hide unused textboxes
        for i in range(len(correct_answers), 9):
            blank_components.append(gr.update(visible=False))
        
        # Store state with positions and correct answers
        state = {
            "correct_answers": correct_answers,
            "positions": positions,
            "all_words": all_words,
            "original_poem": poem,
            "word_bank": word_bank
        }
        
        return {
            output_markdown: gr.update(value=f"**Your Poem:**\n\n{formatted_blanked}\n\n{word_bank_html}", visible=True),
            blank_inputs: gr.update(visible=True),
            **{blank_textboxes[i]: blank_components[i] for i in range(9)},
            check_button: gr.update(visible=True),
            solution_button: gr.update(visible=False),
            result_markdown: gr.update(visible=False),
            solution_display: gr.update(visible=False),
            poem_state: state
        }
    
    def check_answers(*args):
        """Check user answers with position-aware validation and color coding."""
        state = args[-1]
        user_inputs = list(args[:-1])
        
        correct_answers = state.get("correct_answers", [])
        positions = state.get("positions", [])
        all_words = state.get("all_words", [])
        word_bank = state.get("word_bank", [])
        
        # Recreate word bank HTML
        word_bank_html = "**Word Bank:** " + " | ".join([f"<span style='background: #e0e7ff; padding: 4px 12px; border-radius: 12px; color: #667eea; font-weight: 600;'>{word}</span>" for word in word_bank])
        
        # Check each answer against correct answer at same position
        results = []
        for i, (user_word, correct_word) in enumerate(zip(user_inputs[:len(correct_answers)], correct_answers)):
            is_correct = (user_word or "").lower().strip() == correct_word.lower().strip()
            results.append(is_correct)
        
        # Build poem with user's answers color-coded (in place)
        display_words = all_words.copy()
        for i, pos in enumerate(positions):
            if i < len(user_inputs) and user_inputs[i]:
                # Color code based on correctness
                if results[i]:
                    display_words[pos] = f"<span style='color: #38ef7d; font-weight: bold;'>{user_inputs[i]}</span>"
                else:
                    display_words[pos] = f"<span style='color: #f5576c; font-weight: bold;'>{user_inputs[i]}</span>"
            else:
                # Empty input - keep underscores in poem and mark as incorrect (red)
                display_words[pos] = f"<span style='color: #f5576c; font-weight: bold;'>_______________</span>"
        
        # Format poem with line breaks (split at periods or newlines)
        poem_text = " ".join(display_words)
        # Try to detect natural line breaks in original poem
        poem_lines = poem_text.replace('\n', '. ').split('. ')
        formatted_poem = "<br><br>".join([line.strip() + ('.' if i < len(poem_lines)-1 and not line.endswith('.') else '') for i, line in enumerate(poem_lines) if line.strip()])
        
        # Generate result message
        if all(results):
            message = "## 🎉 Perfect! All answers are correct!"
            show_solution = False
        else:
            incorrect_count = sum(1 for r in results if not r)
            empty_count = sum(1 for i, inp in enumerate(user_inputs[:len(correct_answers)]) if not inp)
            if empty_count > 0:
                message = f"## ⚠️ {empty_count} blank(s) still empty, {incorrect_count} incorrect. Keep trying!"
            else:
                message = f"## ❌ {incorrect_count} incorrect. Try again or check the solution!"
            show_solution = True
        
        # Update the main output_markdown instead of result_markdown to avoid duplication
        return {
            output_markdown: gr.update(value=f"{message}\n\n**Your Poem:**\n\n{formatted_poem}\n\n{word_bank_html}", visible=True),
            result_markdown: gr.update(visible=False), # Hide result markdown as we use main output
            solution_button: gr.update(visible=show_solution),
            **{blank_textboxes[i]: gr.update() for i in range(9)}
        }
    
    def show_solution(state):
        """Show the complete poem with all correct answers in green."""
        correct_answers = state.get("correct_answers", [])
        positions = state.get("positions", [])
        all_words = state.get("all_words", [])
        
        # Build poem with all correct answers in green
        solution_words = all_words.copy()
        for i, pos in enumerate(positions):
            solution_words[pos] = f"<span style='color: #38ef7d; font-weight: bold;'>{correct_answers[i]}</span>"
        
        # Format poem with line breaks (split at periods or newlines)
        poem_text = " ".join(solution_words)
        poem_lines = poem_text.replace('\n', '. ').split('. ')
        formatted_poem = "<br><br>".join([line.strip() + ('.' if i < len(poem_lines)-1 and not line.endswith('.') else '') for i, line in enumerate(poem_lines) if line.strip()])
        
        # Update main output_markdown to show solution in place
        return {
            output_markdown: gr.update(value=f"## ✅ Complete Poem (Solution)\n\n{formatted_poem}", visible=True),
            solution_display: gr.update(visible=False), # Hide separate solution display
            solution_button: gr.update(visible=False),
            result_markdown: gr.update(visible=False)
        }
    
    with gr.Blocks(title="Poem Generator") as demo:
        gr.Markdown("# 📝 Poem Generator with Fill-in-the-Blank Quiz")
        gr.Markdown("Create a personalized poem and test your vocabulary!")
        
        with gr.Row():
            age_input = gr.Textbox(value="9", label="Age", placeholder="Enter age (6-99)")
            length_dropdown = gr.Dropdown(["short", "medium", "long"], value="short", label="Length")
            difficulty_dropdown = gr.Dropdown(["Easy", "Medium", "Hard"], value="Easy", label="Difficulty")
        
        with gr.Row():
            theme_dropdown = gr.Dropdown(
                choices=THEMES,
                label="Theme (select one)",
                value=THEMES[0] if THEMES else None
            )
        
        with gr.Row():
            interests_checkboxes = gr.CheckboxGroup(
                choices=INTERESTS,
                label=f"Interests (select up to {MAX_INTERESTS})",
                value=[]
            )
        
        generate_button = gr.Button("✨ Generate Poem", variant="primary")
        
        output_markdown = gr.Markdown("", visible=False)
        
        # Container for blank input fields
        with gr.Column(visible=False) as blank_inputs:
            gr.Markdown("### Fill in the blanks")
            with gr.Row():
                blank_textboxes = [
                    gr.Textbox(label="", visible=False, placeholder=f"Blank {i+1}", value="_______________", scale=1, show_label=False, type="text")
                    for i in range(9)
                ]
        
        with gr.Row():
            check_button = gr.Button("Check Answers", visible=False, variant="primary")
            solution_button = gr.Button("Show Solution", visible=False, variant="secondary")
        
        result_markdown = gr.Markdown("", visible=False)
        solution_display = gr.Markdown("", visible=False)
        poem_state = gr.State({})
        
        # Limit interests selection
        def limit_interests(selected):
            if len(selected) > MAX_INTERESTS:
                return selected[:MAX_INTERESTS]
            return selected
        
        interests_checkboxes.change(
            limit_interests,
            inputs=[interests_checkboxes],
            outputs=[interests_checkboxes]
        )
        
        generate_button.click(
            generate_poem,
            inputs=[age_input, length_dropdown, difficulty_dropdown, theme_dropdown, interests_checkboxes],
            outputs=[output_markdown, blank_inputs] + blank_textboxes + [check_button, solution_button, result_markdown, solution_display, poem_state],
        )
        
        check_button.click(
            check_answers,
            inputs=blank_textboxes + [poem_state],
            outputs=[output_markdown, result_markdown, solution_button] + blank_textboxes
        )
        
        solution_button.click(
            show_solution,
            inputs=[poem_state],
            outputs=[output_markdown, solution_display, solution_button, result_markdown]
        )

        # Enable the Check Answers button only when all required blanks are filled
        def validate_blanks(*args):
            # args: first 9 are textbox values, last is state
            *vals, state = args
            n_required = len(state.get("correct_answers", [])) if isinstance(state, dict) else 0
            # consider only first n_required textboxes
            all_filled = n_required > 0 and all((vals[i] or "").strip() for i in range(n_required))
            return {check_button: gr.update(interactive=all_filled)}

        # Attach change handlers to all blank textboxes
        for tb in blank_textboxes:
            tb.change(
                validate_blanks,
                inputs=blank_textboxes + [poem_state],
                outputs=[check_button]
            )
    
    return demo


if __name__ == "__main__":
    from dotenv import load_dotenv
    load_dotenv()
    
    model = PoetryModel()
    app = create_app(model)
    app.launch()