import gradio as gr
import random
import pandas as pd
import os
import threading
from threading import Event
from utils.data_loader import get_random_example
from utils.models import generate_summaries, model_names
from utils.ui_helpers import toggle_context_display, update_feedback, get_context_html
from utils.leaderboard import load_leaderboard_data, save_leaderboard_data

# Global interrupt mechanism for model generation
generation_interrupt = Event()

# Feedback options for different voting outcomes
feedback_options = {
    "left": ["Model A: More complete", "Model A: More accurate", "Model A: More relevant", "Model A: Better written", "Model A: Better refusal (if applicable)"],
    "right": ["Model B: More complete", "Model B: More accurate", "Model B: More relevant", "Model B: Better written", "Model B: Better refusal (if applicable)"],
    "tie": ["Both complete", "Both accurate", "Both well written", "Both handle refusal well (if applicable)"],
    "neither": ["Both incomplete", "Both hallucinate", "Both irrelevant", "Both incorrectly refuse (if applicable)", "A is bad", "B is bad"]
}

def load_context():
    """Load a new question and context (fast operation)"""
    generation_interrupt.clear()
    example = get_random_example()
    
    context_desc = example.get('processed_context_desc', '')
    if context_desc:
        context_desc = f"<div class='context-topic'><span class='topic-label'>The question and context are about:</span> {context_desc}</div>"
    
    show_full = False
    context_html = get_context_html(example, show_full=show_full)
    
    return [
        example,
        gr.update(value=example['question']),
        gr.update(value=context_desc, visible=bool(context_desc)),
        gr.update(value=context_html),
        gr.update(value="Show Full Context", elem_classes=["context-toggle-button"]),
        show_full
    ]

def generate_model_summaries_with_timeout(example, timeout=60):
    """Run model inference in a separate thread with timeout for interruptibility"""
    import threading
    import time
    
    result = {
        "model_a": "",
        "model_b": "",
        "summary_a": "",
        "summary_b": "",
        "completed": False
    }
    
    if generation_interrupt.is_set():
        return result

    def run_generation():
        try:
            m_a_name, m_b_name = random.sample(model_names, 2)
            s_a, s_b = generate_summaries(example, m_a_name, m_b_name)
            
            if not generation_interrupt.is_set():
                result["model_a"] = m_a_name
                result["model_b"] = m_b_name
                result["summary_a"] = s_a
                result["summary_b"] = s_b
                result["completed"] = True
        except Exception as e:
            print(f"Error in generation thread: {e}")
    
    generation_thread = threading.Thread(target=run_generation)
    generation_thread.daemon = True
    generation_thread.start()
    
    # Uncomment this critical waiting code
    start_time = time.time()
    while time.time() - start_time < timeout:
        if generation_interrupt.is_set() or not generation_thread.is_alive() or result["completed"]:
            break
        time.sleep(0.1)
        
    return result

async def generate_model_summaries_with_timeout_async(example, timeout=30):
    """Async version that properly waits for the thread"""
    import asyncio
    import threading
    import time
    
    result = {
        "model_a": "",
        "model_b": "",
        "summary_a": "",
        "summary_b": "",
        "completed": False
    }
    
    if generation_interrupt.is_set():
        return result

    def run_generation():
        try:
            m_a_name, m_b_name = random.sample(model_names, 2)
            s_a, s_b = generate_summaries(example, m_a_name, m_b_name)
            
            if not generation_interrupt.is_set():
                result["model_a"] = m_a_name
                result["model_b"] = m_b_name
                result["summary_a"] = s_a
                result["summary_b"] = s_b
                result["completed"] = True
        except Exception as e:
            print(f"Error in generation thread: {e}")
    
    generation_thread = threading.Thread(target=run_generation)
    generation_thread.daemon = True
    generation_thread.start()
    
    # Use asyncio.sleep instead of time.sleep for async waiting
    start_time = time.time()
    while time.time() - start_time < timeout:
        if generation_interrupt.is_set() or not generation_thread.is_alive() or result["completed"]:
            break
        await asyncio.sleep(0.1)  # Non-blocking sleep
        
    return result

def process_generation_result(result):
    """Process the results from the threaded generation function"""
    if not result["completed"]:
        # Generation was interrupted or failed
        return [
            "", "", "", "", None, [], False, load_leaderboard_data(),
            gr.update(value="Generation was interrupted or timed out. Please try again."),
            gr.update(value="Generation was interrupted or timed out. Please try again."),
            gr.update(interactive=True, elem_classes=["vote-button"]),
            gr.update(interactive=True, elem_classes=["vote-button"]),
            gr.update(interactive=True, elem_classes=["vote-button"]),
            gr.update(interactive=True, elem_classes=["vote-button", "vote-button-neither"]),
            gr.update(choices=[], value=[], interactive=False, visible=False),
            gr.update(visible=False),
            gr.update(interactive=False, visible=True),
            gr.update(visible=False),
            gr.update(interactive=True),
            gr.update(elem_classes=[])
        ]
    
    # Generation completed successfully
    agg_results = load_leaderboard_data()
    return [
        result["model_a"], result["model_b"], 
        result["summary_a"], result["summary_b"],
        None, [], False, agg_results,
        gr.update(value=result["summary_a"]),
        gr.update(value=result["summary_b"]),
        gr.update(interactive=True, elem_classes=["vote-button"]),
        gr.update(interactive=True, elem_classes=["vote-button"]),
        gr.update(interactive=True, elem_classes=["vote-button"]),
        gr.update(interactive=True, elem_classes=["vote-button", "vote-button-neither"]),
        gr.update(choices=[], value=[], interactive=False, visible=False),
        gr.update(visible=False),
        gr.update(interactive=False, visible=True),
        gr.update(visible=False),
        gr.update(interactive=True),
        gr.update(elem_classes=[])
    ]
async def process_example_async(example):
    result = await generate_model_summaries_with_timeout_async(example)
    return process_generation_result(result)

def process_example_sync(example):
    result = generate_model_summaries_with_timeout(example)
    return process_generation_result(result)

def select_vote_improved(winner_choice):
    """Updates UI based on vote selection"""
    feedback_choices = feedback_options.get(winner_choice, [])

    btn_a_classes = ["vote-button"]
    btn_b_classes = ["vote-button"]
    btn_tie_classes = ["vote-button"]
    btn_neither_classes = ["vote-button", "vote-button-neither"]
    
    if winner_choice == 'left':
        btn_a_classes.append("selected")
    elif winner_choice == 'right':
        btn_b_classes.append("selected")
    elif winner_choice == 'tie':
        btn_tie_classes.append("selected")
    elif winner_choice == 'neither':
        btn_neither_classes.append("selected")

    return [
        winner_choice,
        gr.update(choices=feedback_choices, value=[], interactive=True, visible=True),
        gr.update(visible=True),
        gr.update(interactive=True),
        gr.update(elem_classes=btn_a_classes),
        gr.update(elem_classes=btn_b_classes),
        gr.update(elem_classes=btn_tie_classes),
        gr.update(elem_classes=btn_neither_classes)
    ]

def submit_vote_fixed(m_a, m_b, winner, feedback, current_results):
    """Processes vote and updates leaderboard"""
    if winner is None:
        print("Warning: Submit called without a winner selected.")
        return {}

    updated_results = current_results.copy()
    models_involved = [m_a, m_b]
    for model in models_involved:
         if model not in updated_results["wins"]:
            updated_results["wins"][model] = 0
            updated_results["losses"][model] = 0
            updated_results["ties"][model] = 0

    if winner == 'left':
        updated_results["wins"][m_a] = updated_results["wins"].get(m_a, 0) + 1
        updated_results["losses"][m_b] = updated_results["losses"].get(m_b, 0) + 1
    elif winner == 'right':
        updated_results["wins"][m_b] = updated_results["wins"].get(m_b, 0) + 1
        updated_results["losses"][m_a] = updated_results["losses"].get(m_a, 0) + 1
    elif winner == 'tie':
        updated_results["ties"][m_a] = updated_results["ties"].get(m_a, 0) + 1
        updated_results["ties"][m_b] = updated_results["ties"].get(m_b, 0) + 1

    updated_results["votes"] = updated_results.get("votes", 0) + 1
    save_leaderboard_data(updated_results)

    # Prepare Results Table
    results_list = []
    all_models = list(set(list(updated_results["wins"].keys()) + 
                          list(updated_results["losses"].keys()) + 
                          list(updated_results["ties"].keys())))

    for model in sorted(all_models):
        wins = updated_results["wins"].get(model, 0)
        losses = updated_results["losses"].get(model, 0)
        ties = updated_results["ties"].get(model, 0)
        total_comparisons = wins + losses + ties
        win_rate = (wins + 0.5 * ties) / total_comparisons if total_comparisons > 0 else 0.0
        results_list.append({
            "Model": model,
            "Win Rate (%)": f"{win_rate:.1%}",
            "Wins": wins,
            "Losses": losses,
            "Ties": ties,
            "Comparisons": total_comparisons
        })
        
    results_df = pd.DataFrame(results_list)
    if not results_df.empty:
        results_df['Win Rate Value'] = results_df['Win Rate (%)'].str.rstrip('%').astype('float') / 100.0
        results_df = results_df.sort_values(by='Win Rate Value', ascending=False).drop(columns=['Win Rate Value'])

    return [
        True, updated_results,
        gr.update(interactive=False), gr.update(interactive=False),
        gr.update(interactive=False), gr.update(interactive=False),
        gr.update(interactive=False), gr.update(visible=True),
        gr.update(visible=False), gr.update(visible=True),
        gr.update(interactive=False), gr.update(value=results_df, visible=True),
        gr.update(elem_classes=["results-revealed"]),
        gr.update(interactive=True), gr.update(value=m_a), gr.update(value=m_b)
    ]

# Create Gradio interface
with gr.Blocks(theme=gr.themes.Default(
    primary_hue=gr.themes.colors.orange,
    secondary_hue=gr.themes.colors.slate
)) as demo:
    # Load CSS
    css_path = os.path.join(os.getcwd(), 'static', 'styles.css')
    with open(css_path, 'r') as f:
        css_content = f.read()
    gr.HTML(f"<style>{css_content}</style>")

    # State Variables
    current_example = gr.State({})
    model_a_name = gr.State("")
    model_b_name = gr.State("")
    summary_a_text = gr.State("")
    summary_b_text = gr.State("")
    selected_winner = gr.State(None)
    feedback_list = gr.State([])
    show_results_state = gr.State(False)
    results_agg = gr.State({"wins": {}, "losses": {}, "ties": {}, "votes": 0})
    show_full_context = gr.State(False)

    # Create Tabs
    with gr.Tabs() as tabs:
        # Main Arena Tab
        with gr.TabItem("Arena", id="arena-tab"):
            gr.Markdown("# RAG Summarizer Arena")
            gr.Markdown("Compare summaries generated by different models based on the provided context and query. Select the better summary, or choose 'Tie' or 'Neither'. Your feedback helps evaluate model performance.")

            # Main container
            with gr.Column(elem_id="main-interface-area") as main_interface_area:
                # Query section
                with gr.Row(elem_id="query-title-row"):
                    gr.Markdown("### Query", elem_classes="section-heading")

                with gr.Row(elem_id="query-container"):
                    with gr.Row(elem_classes="query-box-row"):
                        query_display = gr.Markdown(value="Loading question...", elem_classes="query-text")
                    random_question_btn = gr.Button("🔄 Get Random Question", elem_classes="query-button")
                
                # Context description and display
                context_description = gr.Markdown("", elem_classes="context-description")
                
                with gr.Row(elem_id="context-header-row"):
                    gr.Markdown("### Context Provided", elem_classes="context-title")
                    context_toggle_btn = gr.Button("Show Full Context", elem_classes=["context-toggle-button"])
                    
                context_display = gr.HTML(value="Loading context...", label="Context Chunks")

                gr.Markdown("---")
                gr.Markdown("### Compare Summaries", elem_classes="section-heading")

                # Model summaries
                with gr.Row():
                    with gr.Column(scale=1):
                        with gr.Group(elem_classes=["summary-card", "summary-card-a"]):
                            summary_a_display = gr.Textbox(label="Model A", lines=10, interactive=False, show_copy_button=True)
                    with gr.Column(scale=1):
                        with gr.Group(elem_classes=["summary-card", "summary-card-b"]):
                            summary_b_display = gr.Textbox(label="Model B", lines=10, interactive=False, show_copy_button=True)

                # Voting section
                gr.Markdown("### Cast Your Vote", elem_classes="section-heading")
                with gr.Row():
                    vote_button_a = gr.Button("⬅️ Summary A is Better", elem_classes=["vote-button"])
                    vote_button_tie = gr.Button("🤝 Tie / Equally Good", elem_classes=["vote-button"])
                    vote_button_b = gr.Button("➡️ Summary B is Better", elem_classes=["vote-button"])
                    vote_button_neither = gr.Button("❌ Neither is Adequate", elem_classes=["vote-button", "vote-button-neither"])

                # Feedback and Submit sections
                with gr.Group(elem_classes=["feedback-section"], visible=False) as feedback_section:
                    feedback_checkboxes = gr.CheckboxGroup(label="Feedback (optional)", choices=[], interactive=False)
                submit_button = gr.Button("Submit Vote", variant="primary", interactive=False, elem_id="submit-button")

                # Results area
                with gr.Column(visible=False) as results_reveal_area:
                    gr.Markdown("---")
                    gr.Markdown("### ✅ Vote Submitted!", elem_classes="section-heading")
                     
                    # Model reveal section
                    with gr.Row():
                        with gr.Column(scale=1):
                            gr.Markdown("### Model A was actually:", elem_classes="section-heading")
                            model_a_reveal = gr.Markdown("", elem_classes="model-reveal model-a-reveal")
                        with gr.Column(scale=1):
                            gr.Markdown("### Model B was actually:", elem_classes="section-heading")
                            model_b_reveal = gr.Markdown("", elem_classes="model-reveal model-b-reveal")
                     
                    gr.HTML("<div style='height: 10px;'></div>")
                    
                    # Try another button
                    with gr.Row(elem_classes=["control-buttons"]):
                        try_another_btn = gr.Button("🔄 Try Another Question", elem_id="try-another-btn")

        # Leaderboard Tab
        with gr.TabItem("Leaderboard", id="leaderboard-tab"):
            gr.Markdown("# Model Performance Leaderboard")
            gr.Markdown("View aggregate performance statistics for all models. The table below shows win rates, wins, losses, and ties for each model based on all evaluations.")
            results_table_display = gr.DataFrame(label="Model Performance", interactive=False, wrap=True)

    # Generic function to handle starting a new example
    def handle_new_example_click():
        generation_interrupt.set()  # Interrupt any ongoing generation
        return load_context()[0]
        
    def update_ui_for_new_context(example):
        return [
            gr.update(value=example['question']),
            gr.update(value=example.get('processed_context_desc', ''), visible=bool(example.get('processed_context_desc', ''))),
            gr.update(value=get_context_html(example, False)),
            gr.update(value="Show Full Context", elem_classes=["context-toggle-button"]),
            False
        ]

    # Event handling
    # Toggle context display
    context_toggle_btn.click(
        fn=toggle_context_display,
        inputs=[current_example, show_full_context],
        outputs=[show_full_context, context_display, context_toggle_btn]
    )
    
    # Initial loading - context first, then summaries
    demo.load(
        fn=load_context,
        inputs=[],
        outputs=[current_example, query_display, context_description, context_display, 
                context_toggle_btn, show_full_context]
    ).then(
        fn=process_example_async,
        inputs=[current_example],
        outputs=[model_a_name, model_b_name, summary_a_text, summary_b_text,
                selected_winner, feedback_list, show_results_state, results_agg,
                summary_a_display, summary_b_display, vote_button_a, vote_button_b, 
                vote_button_tie, vote_button_neither, feedback_checkboxes, feedback_section, 
                submit_button, results_reveal_area, random_question_btn, main_interface_area]
    )

    # Random Question and Try Another buttons with interruption
    for btn in [random_question_btn, try_another_btn]:
        btn.click(
            fn=handle_new_example_click,
            inputs=[],
            outputs=[current_example]
        ).then(
            fn=update_ui_for_new_context,
            inputs=[current_example],
            outputs=[query_display, context_description, context_display, 
                    context_toggle_btn, show_full_context]
        ).then(
            fn=process_example_sync,
            inputs=[current_example],
            outputs=[model_a_name, model_b_name, summary_a_text, summary_b_text,
                    selected_winner, feedback_list, show_results_state, results_agg,
                    summary_a_display, summary_b_display, vote_button_a, vote_button_b, 
                    vote_button_tie, vote_button_neither, feedback_checkboxes, feedback_section, 
                    submit_button, results_reveal_area, random_question_btn, main_interface_area]
        )

    # Vote button handlers
    for btn, choice in zip(
        [vote_button_a, vote_button_b, vote_button_tie, vote_button_neither],
        ['left', 'right', 'tie', 'neither']
    ):
        btn.click(
            fn=lambda choice=choice: select_vote_improved(choice),
            inputs=None,
            outputs=[selected_winner, feedback_checkboxes, feedback_section, submit_button, 
                    vote_button_a, vote_button_b, vote_button_tie, vote_button_neither]
        )

    # Update feedback when checkboxes change
    feedback_checkboxes.change(
        fn=update_feedback,
        inputs=[feedback_checkboxes],
        outputs=[feedback_list]
    )

    # Process vote submission and reveal results
    submit_button.click(
        fn=submit_vote_fixed,
        inputs=[model_a_name, model_b_name, selected_winner, feedback_list, results_agg],
        outputs=[show_results_state, results_agg, vote_button_a, vote_button_b, 
                vote_button_tie, vote_button_neither, feedback_checkboxes,
                feedback_section, submit_button, results_reveal_area,
                random_question_btn, results_table_display, main_interface_area,
                context_toggle_btn, model_a_reveal, model_b_reveal]
    )

if __name__ == "__main__":
    demo.launch(debug=True)