import gradio as gr
import pandas as pd
import os
import time
from setup import get_paper_agent

def generate_content_logic(paper_input, mode, progress_tracker=None):
    """
    Runs the agent with a prompt tailored to the selected mode (Flashcards vs Summary).
    Now also extracts the Paper Title.
    """
    flashcard_file = "flashcards_export.csv"
    summary_file = "summary_export.txt"
    title_file = "paper_title.txt" 

    # 1. Construct Prompt based on Mode
    # We explicitly ask the agent to save the title to 'paper_title.txt'
    base_instruction = (
        f"Step 1: Use 'get_paper_details' to find '{paper_input}'. "
        f"Step 2: From the tool output, extract the Title and save it to '{title_file}'. "
        f"Step 3: Extract the ID and use it to download the paper..."
      )

    if mode == "Summary Only":
        prompt = (base_instruction +
                  f"Then, analyze the paper and write a comprehensive summary. "
                  f"Save the summary text to '{summary_file}'.")
    else:
        prompt = (base_instruction +
                  f"Then, analyze the paper, create flashcards, and save them to '{flashcard_file}'.")

    # 2. Start Agent
    if progress_tracker:
        progress_tracker(0.2, desc=f"Agent Working: Finding '{paper_input}'...")

    try:
        agent = get_paper_agent()
        agent.run(prompt)
    except Exception as e:
        print(f"Agent run failed: {e}")
        return None, "error", None

    # 3. Process Output
    if progress_tracker:
        progress_tracker(0.8, desc="Processing Output...")

    # A. Retrieve Title (Fallback to user input if file missing)
    detected_title = paper_input
    if os.path.exists(title_file):
        try:
            with open(title_file, 'r') as f:
                detected_title = f.read().strip()
        except:
            pass

    # B. Retrieve Content
    if mode == "Summary Only":
        if os.path.exists(summary_file):
            with open(summary_file, 'r') as f:
                return f.read(), "summary", detected_title
    else:
        if os.path.exists(flashcard_file):
            try:
                df = pd.read_csv(flashcard_file)
                cards = df.to_dict('records')
                return cards, "flashcards", detected_title
            except:
                pass

    return None, "error", detected_title

def get_card_html(card_data, revealed):
    """Generates the HTML for the flashcard."""
    if not card_data:
        return "<div style='padding:50px; text-align:center;'>No flashcards available.</div>"

    content = card_data['Back'] if revealed else card_data['Front']
    side_label = "ANSWER (Back)" if revealed else "QUESTION (Front)"
    text_color = "#333" if revealed else "#000"
    bg_color = "#e6f7ff" if revealed else "#ffffff"
    border_color = "#4a90e2" if revealed else "#ccc"

    html = f"""
    <div style="
        padding: 40px;
        border: 2px solid {border_color};
        border-radius: 15px;
        min-height: 300px;
        display: flex;
        flex-direction: column;
        justify-content: center;
        align-items: center;
        text-align: center;
        background-color: {bg_color};
        box-shadow: 0 4px 12px 0 rgba(0,0,0,0.1);
        transition: all 0.3s ease;
    ">
        <div style="font-size: 0.9em; color: #888; text-transform: uppercase; letter-spacing: 2px; margin-bottom: 20px;">
            {side_label}
        </div>
        <div style="font-size: 1.5em; font-weight: 500; color: {text_color}; font-family: sans-serif; line-height: 1.4;">
            {content}
        </div>
    </div>
    """
    return html

# --- Event Handlers ---

def on_generate_click(paper_input, mode, progress=gr.Progress()):
    gr.Info(f"Starting analysis for {mode}...")

    # Now accepts 3 return values: data, type, title
    data, result_type, title = generate_content_logic(paper_input, mode, progress_tracker=progress)

    # Format the title for display
    title_md = f"## 📄 {title}" if title else ""

    # Handle Error
    if result_type == "error" or data is None:
        gr.Warning("Failed to generate content. Check logs.")
        return (
            [], 0, False, "", "0/0", None,      # Flashcard defaults
            gr.update(visible=False),           # Flashcard Area
            gr.update(visible=False, value=""), # Summary Area
            gr.update(visible=False, value="")  # Title Area
        )

    # Handle Summary
    if result_type == "summary":
        gr.Info("Summary Generated!")
        return (
            [], 0, False, "", "0/0", None,
            gr.update(visible=False),           # Hide Flashcard Area
            gr.update(visible=True, value=data),# Show Summary Area
            gr.update(visible=True, value=title_md) # Show Title
        )

    # Handle Flashcards
    else:
        cards = data
        index = 0
        html = get_card_html(cards[index], False)
        prog_text = f"Card 1 of {len(cards)}"
        filename = "flashcards_export.csv"

        gr.Info("Flashcards Ready!")
        return (
            cards, index, False, html, prog_text, filename,
            gr.update(visible=True),            # Show Flashcard Area
            gr.update(visible=False),           # Hide Summary Area
            gr.update(visible=True, value=title_md) # Show Title
        )

# Navigation handlers (Next/Prev/Reveal)
def on_next_click(cards, index):
    if not cards: return index, False, "", "0/0"
    new_index = (index + 1) % len(cards)
    html = get_card_html(cards[new_index], False)
    return new_index, False, html, f"Card {new_index + 1} of {len(cards)}"

def on_prev_click(cards, index):
    if not cards: return index, False, "", "0/0"
    new_index = (index - 1) % len(cards)
    html = get_card_html(cards[new_index], False)
    return new_index, False, html, f"Card {new_index + 1} of {len(cards)}"

def on_reveal_click(cards, index, revealed):
    if not cards: return revealed, "", "0/0"
    new_revealed = not revealed
    html = get_card_html(cards[index], new_revealed)
    return new_revealed, html, f"Card {index + 1} of {len(cards)}"

# --- Gradio Interface ---

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    # State Variables
    flashcards_state = gr.State(value=[])
    current_index_state = gr.State(value=0)
    is_revealed_state = gr.State(value=False)

    gr.Markdown("# 📑 Research Paper Assistant")
    gr.Markdown("Analyze a paper to generate Flashcards or a Summary.")

    # 1. Input Section
    with gr.Row():
        with gr.Column(scale=4):
            paper_input = gr.Textbox(label="Paper Topic / URL", placeholder="e.g. 'Attention is All You Need'")
            mode_radio = gr.Radio(["Flashcards", "Summary Only"], label="Output Mode", value="Flashcards")

        generate_btn = gr.Button("Generate", variant="primary", scale=1)

    # 2. Output Areas
    
    # --- NEW: Title Display Area ---
    paper_title_display = gr.Markdown(visible=False) 
    # -------------------------------

    # A. Summary Area (Hidden by default)
    summary_area = gr.Markdown(visible=False, label="Paper Summary")

    # B. Flashcard Area (Hidden by default)
    with gr.Column(visible=False) as flashcard_area:
        card_display = gr.HTML(value="", label="Flashcard")

        with gr.Row():
            prev_btn = gr.Button("← Previous", variant="secondary")
            reveal_btn = gr.Button("🔄 Reveal Answer", variant="primary")
            next_btn = gr.Button("Next →", variant="secondary")

        progress_label = gr.Label(value="Card 0/0", show_label=False, color="gray")

        gr.Markdown("### Export")
        download_file = gr.File(label="Download Flashcards (.csv)")

    # 3. Connect Events

    generate_btn.click(
        fn=on_generate_click,
        inputs=[paper_input, mode_radio],
        outputs=[
            flashcards_state, current_index_state, is_revealed_state,
            card_display, progress_label, download_file, # Flashcard UI
            flashcard_area, summary_area,                # Visibility
            paper_title_display                          # <--- Update Title
        ]
    )

    next_btn.click(fn=on_next_click, inputs=[flashcards_state, current_index_state], outputs=[current_index_state, is_revealed_state, card_display, progress_label])
    prev_btn.click(fn=on_prev_click, inputs=[flashcards_state, current_index_state], outputs=[current_index_state, is_revealed_state, card_display, progress_label])
    reveal_btn.click(fn=on_reveal_click, inputs=[flashcards_state, current_index_state, is_revealed_state], outputs=[is_revealed_state, card_display, progress_label])

demo.launch(demo.launch(share=True, ssr_mode=False))