import gradio as gr
from dotenv import load_dotenv
from research_manager import ResearchManager
import markdown
import re
from file_processor import process_file, get_file_icon, format_file_size

load_dotenv(override=True)


async def run_research(query: str, model_choice: str, conversation_history: list, attachments: list, progress=gr.Progress()):
    """Run research and yield updates for both report and references"""
    status_messages = []
    final_report_md = ""
    references_list = []
    total_searches = 10  # Default, will be updated
    writing_progress = 0.75  # Track dummy progress during writing

    progress(0, desc="🚀 Starting research...")

    # Show the query at the top with action buttons
    # Escape single quotes for JavaScript
    query_escaped = query.replace("'", "'")

    query_display = f'''
    <div class="user-query">
        <div class="query-header">
            <div class="query-label">Your Research Query:</div>
            <div class="query-actions">
                <button class="icon-btn" onclick="navigator.clipboard.writeText('{query_escaped}'); this.innerHTML='✓ Copied'; setTimeout(() => this.innerHTML='📋 Copy', 2000);" title="Copy query">
                    📋 Copy
                </button>
                <button class="icon-btn edit-btn" onclick="document.getElementById('edit-query-btn').click();" title="Edit query">
                    ✏️ Edit
                </button>
            </div>
        </div>
        <div class="query-text">{query}</div>
    </div>
    '''

    # Collect all chunks and parse structured messages
    async for chunk in ResearchManager(model_choice).run(query, conversation_history, attachments):
        # Parse structured messages (format: TYPE|data)
        if "|" in chunk:
            msg_type, msg_data = chunk.split("|", 1)

            if msg_type == "INIT":
                progress(0.10, desc="🤖 Agents initialized")
                status_messages.append(msg_data)

            elif msg_type == "PLANNING_COMPLETE":
                total_searches = int(msg_data)
                progress(0.20, desc=f"📋 Planning complete - {total_searches} searches queued")
                status_messages.append(f"Planning complete - {total_searches} searches to perform")

            elif msg_type == "SEARCH_PROGRESS":
                parts = msg_data.split("|")
                current = int(parts[0])
                total = int(parts[1])
                # Progress from 25% to 70% during searches (45% range for 10 searches)
                search_progress = 0.25 + (current / total) * 0.45
                progress(search_progress, desc=f"🔍 Searching {current}/{total}...")
                status_messages.append(f"Searching {current}/{total}...")

            elif msg_type == "SEARCH_COMPLETE":
                progress(0.70, desc="✅ All searches complete")
                status_messages.append(msg_data)

            elif msg_type == "WRITING_START":
                # Start writing at 75%
                writing_progress = 0.75
                progress(writing_progress, desc="✍️ Writing comprehensive report...")
                status_messages.append("Writing report...")

                # Simulate progress during writing (75% -> 85%)
                import asyncio
                async def simulate_writing_progress():
                    for i in range(3):
                        await asyncio.sleep(2)  # Update every 2 seconds
                        nonlocal writing_progress
                        writing_progress = min(0.85, writing_progress + 0.05)
                        progress(writing_progress, desc="✍️ Writing comprehensive report...")

                # Start dummy progress in background
                asyncio.create_task(simulate_writing_progress())

            elif msg_type == "REPORT_READY":
                # Report is ready! Process and show it immediately
                final_report_md = msg_data

                # Extract references from markdown
                references_list = extract_references(final_report_md)

                # Remove the References section from the report markdown
                report_without_refs = remove_references_section(final_report_md)

                # Convert markdown to HTML for better rendering
                final_html = markdown.markdown(
                    report_without_refs,
                    extensions=['extra', 'codehilite', 'tables', 'fenced_code']
                )

                # Make inline citations clickable
                final_html = make_citations_clickable(final_html)

                # Format references as clean HTML list
                references_html = format_references_html(references_list)

                # Mark progress as complete and BREAK the loop
                # This ensures Gradio closes the progress bar immediately
                progress(1.0, desc="✅ Research complete!")

                # Add query display and action buttons to the final report
                report_actions = '''
                <div class="report-actions">
                    <div class="left-actions">
                        <button class="icon-btn" onclick="
                            const reportText = document.querySelector('.output-box').innerText;
                            navigator.clipboard.writeText(reportText);
                            this.innerHTML='✓ Copied';
                            setTimeout(() => this.innerHTML='📋 Copy Report', 2000);
                        " title="Copy entire report">
                            📋 Copy Report
                        </button>
                        <button class="icon-btn" onclick="document.getElementById('rewrite-btn').click();" title="Regenerate report">
                            🔄 Rewrite
                        </button>
                    </div>
                    <div class="feedback-buttons">
                        <button class="icon-btn feedback-btn" onclick="this.classList.toggle('active'); document.querySelector('.feedback-btn.dislike').classList.remove('active');" title="Good response">
                            👍 Like
                        </button>
                        <button class="icon-btn feedback-btn dislike" onclick="this.classList.toggle('active'); document.querySelector('.feedback-btn:not(.dislike)').classList.remove('active');" title="Poor response">
                            👎 Dislike
                        </button>
                    </div>
                </div>
                '''
                final_html_with_query = query_display + final_html + report_actions

                # YIELD THE FINAL REPORT and stop here
                # Email will be sent but we won't wait for it
                yield final_html_with_query, references_html

                # Stop processing - don't wait for email messages
                return

            elif msg_type == "EMAIL_START":
                # Won't reach here because we return after REPORT_READY
                pass

            elif msg_type == "COMPLETE":
                # Won't reach here because we return after REPORT_READY
                pass

        # Show status messages in report area while processing
        if not final_report_md:
            status_html = query_display + '<div class="status-container">'
            for msg in status_messages:
                status_html += f'<div class="status-message">• {msg}</div>'
            status_html += '</div>'
            yield status_html, ""  # Empty references while processing


def remove_references_section(markdown_text):
    """Remove the References section from markdown"""
    # Remove everything from ## References onwards
    ref_pattern = r'##\s*References\s*\n.*'
    cleaned = re.sub(ref_pattern, '', markdown_text, flags=re.DOTALL | re.IGNORECASE)
    return cleaned.strip()


def extract_references(markdown_text):
    """Extract references section from markdown"""
    references = []

    # Find the References section
    ref_pattern = r'##\s*References\s*\n(.*?)(?=\n##|\Z)'
    match = re.search(ref_pattern, markdown_text, re.DOTALL | re.IGNORECASE)

    if match:
        ref_section = match.group(1)
        # Extract numbered list items with markdown links
        # Pattern: 1. [Title](URL) or just plain URLs
        list_items = re.findall(r'\d+\.\s*\[([^\]]+)\]\(([^\)]+)\)', ref_section)
        if list_items:
            references = [(title.strip(), url.strip()) for title, url in list_items]
        else:
            # Fallback: extract plain URLs
            urls = re.findall(r'https?://[^\s<>"{}|\\^`\[\]]+', ref_section)
            references = [(url, url) for url in urls]

    return references


def make_citations_clickable(html_text):
    """Convert inline citations [1], [2], etc. to clickable links that switch to References tab"""
    # Pattern to match [1], [2, 3], [1, 2, 3], etc.
    def replace_citation(match):
        full_match = match.group(0)
        numbers = match.group(1)

        # Split by comma and convert to clickable links
        nums = [n.strip() for n in numbers.split(',')]
        clickable_nums = []

        for num in nums:
            # Create a link that triggers tab switch via JavaScript
            # Using a more reliable Gradio tab selector
            clickable_nums.append(
                f'<a href="javascript:void(0)" class="citation-link" '
                f'data-ref-id="{num}" '
                f'onclick="'
                f'const tabs = document.querySelectorAll(\'button[role=tab]\'); '
                f'if (tabs.length > 1) {{ '
                f'tabs[1].click(); '
                f'setTimeout(() => {{ '
                f'const ref = document.getElementById(\'ref-{num}\'); '
                f'if (ref) {{ '
                f'ref.scrollIntoView({{ behavior: \'smooth\', block: \'center\' }}); '
                f'ref.style.transition = \'all 0.3s\'; '
                f'ref.style.boxShadow = \'0 0 20px rgba(16, 185, 129, 0.5)\'; '
                f'setTimeout(() => {{ ref.style.boxShadow = \'\'; }}, 2000); '
                f'}} '
                f'}}, 300); '
                f'}} '
                f'" '
                f'title="Click to view reference {num}">{num}</a>'
            )

        return '[' + ', '.join(clickable_nums) + ']'

    # Match patterns like [1], [2, 3], [1, 2, 3, 4], etc.
    citation_pattern = r'\[(\d+(?:\s*,\s*\d+)*)\]'
    return re.sub(citation_pattern, replace_citation, html_text)


def format_references_html(references):
    """Format references as clean HTML with IDs for linking"""
    if not references:
        return '<div class="no-references">No references found</div>'

    html = '<div class="references-list">'
    for idx, (title, url) in enumerate(references, 1):
        html += f'''
        <div class="reference-item" id="ref-{idx}">
            <span class="ref-number">[{idx}]</span>
            <a href="{url}" target="_blank" class="ref-link">
                {title}
            </a>
        </div>
        '''
    html += '</div>'
    return html


async def run_simple_search(query: str, model_choice: str, conversation_history: list, attachments: list, progress=gr.Progress()):
    """Run a quick follow-up search without full research workflow"""
    progress(0, desc="🔍 Quick search starting...")

    # Escape single quotes for JavaScript
    query_escaped = query.replace("'", "&#39;")

    query_display = f'''
    <div class="user-query simple-search">
        <div class="query-header">
            <div class="query-label">Your Follow-up Question:</div>
            <div class="query-actions">
                <button class="icon-btn" onclick="navigator.clipboard.writeText('{query_escaped}'); this.innerHTML='✓ Copied'; setTimeout(() => this.innerHTML='📋 Copy', 2000);" title="Copy query">
                    📋 Copy
                </button>
            </div>
        </div>
        <div class="query-text">{query}</div>
    </div>
    '''

    answer_text = ""

    async for chunk in ResearchManager(model_choice).run_simple_search(query, conversation_history, attachments):
        # Parse structured messages
        if "|" in chunk:
            msg_type, msg_data = chunk.split("|", 1)

            if msg_type == "SIMPLE_SEARCH_START":
                progress(0.3, desc="🔍 Searching...")
                yield query_display + '<div class="status-message">• Searching...</div>', ""

            elif msg_type == "SIMPLE_SEARCH_COMPLETE":
                answer_text = msg_data
                progress(1.0, desc="✅ Complete!")

                # Convert markdown to HTML
                answer_html = markdown.markdown(
                    answer_text,
                    extensions=['extra', 'codehilite', 'tables', 'fenced_code']
                )

                # Add answer actions
                answer_actions = '''
                <div class="answer-actions">
                    <button class="icon-btn" onclick="
                        const answerText = document.querySelector('.simple-answer').innerText;
                        navigator.clipboard.writeText(answerText);
                        this.innerHTML='✓ Copied';
                        setTimeout(() => this.innerHTML='📋 Copy Answer', 2000);
                    " title="Copy answer">
                        📋 Copy Answer
                    </button>
                </div>
                '''

                final_html = query_display + '<div class="simple-answer">' + answer_html + '</div>' + answer_actions
                yield final_html, ""
                return

            elif msg_type == "SIMPLE_SEARCH_ERROR":
                error_html = query_display + f'<div class="error-message">Error: {msg_data}</div>'
                yield error_html, ""
                return

    # Fallback if no complete message
    if answer_text:
        answer_html = markdown.markdown(answer_text, extensions=['extra', 'codehilite', 'tables', 'fenced_code'])
        yield query_display + '<div class="simple-answer">' + answer_html + '</div>', ""


# Create custom dark theme
luntre_theme = gr.themes.Base(
    primary_hue="green",
    secondary_hue="blue",
    neutral_hue="slate",
).set(
    # Dark background colors
    background_fill_primary="#1C1C1C",
    background_fill_secondary="#121212",

    # Input colors
    input_background_fill="#2D2D2D",
    input_background_fill_focus="#374151",
    input_border_color="#374151",
    input_border_color_focus="#10B981",

    # Button colors
    button_primary_background_fill="#10B981",
    button_primary_background_fill_hover="#059669",
    button_primary_text_color="#FFFFFF",
    button_secondary_background_fill="#3B82F6",
    button_secondary_background_fill_hover="#2563EB",

    # Text colors
    body_text_color="#E5E7EB",
    body_text_color_subdued="#9CA3AF",

    # Block colors
    block_background_fill="#2D2D2D",
    block_border_color="#374151",
)


# Custom CSS
custom_css = """
/* Import stylish fonts */
@import url('https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;500;600;700&display=swap');
@import url('https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@300;400;500;600;700&display=swap');

/* Global dark theme overrides */
.gradio-container {
    background: linear-gradient(135deg, #1C1C1C 0%, #121212 100%) !important;
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
}

/* Header styling */
.app-header {
    text-align: center;
    padding: 2rem 1rem 1rem 1rem;
    background: linear-gradient(135deg, #1e3a2e 0%, #1a2332 100%);
    border-radius: 16px;
    margin-bottom: 2rem;
    border: 1px solid #374151;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.3);
}

.brand-name {
    font-size: 2.5rem;
    font-weight: 500;
    background: linear-gradient(135deg, #10B981 0%, #3B82F6 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    margin-bottom: 0.5rem;
    letter-spacing: 0.12em;
    font-family: 'Space Grotesk', 'Outfit', sans-serif;
    text-transform: uppercase;
}

.page-title {
    font-size: 1.25rem;
    color: #9CA3AF;
    font-weight: 400;
}

/* Modern Minimal Input at Bottom - Unified Container */
.input-container-bottom {
    background: rgba(45, 45, 45, 0.6) !important;
    padding: 1.5rem;
    padding-bottom: 1.5rem !important;
    border-radius: 24px;
    border: 1px solid rgba(55, 65, 81, 0.4);
    margin: 1.5rem auto 2rem auto;
    max-width: 900px;
    position: relative;
    z-index: 100;
    backdrop-filter: blur(12px);
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.3);
    overflow: visible !important;
}

/* Ensure all children allow overflow for dropdown */
.input-container-bottom > *,
.controls-row {
    overflow: visible !important;
}

/* Force all children containers to be transparent */
.input-container-bottom > * {
    background: transparent !important;
}

.input-container-bottom .controls-row {
    background: transparent !important;
}

.chat-input textarea {
    background: transparent !important;
    color: #E5E7EB !important;
    border: none !important;
    border-radius: 0 !important;
    font-size: 0.95rem !important;
    padding: 0.5rem 0.75rem !important;
    line-height: 1.6 !important;
    resize: none !important;
    transition: all 0.2s ease !important;
    text-align: left !important;
}

.chat-input textarea:focus {
    background: transparent !important;
    box-shadow: none !important;
    outline: none !important;
}

.chat-input textarea::placeholder {
    color: #6B7280 !important;
    opacity: 0.7;
}

/* Generic textarea - only for textareas NOT in chat-input */
textarea:not(.chat-input textarea) {
    background: rgba(45, 45, 45, 0.6) !important;
    color: #E5E7EB !important;
    border: 1px solid rgba(55, 65, 81, 0.4) !important;
    border-radius: 8px !important;
    font-size: 1rem !important;
    transition: all 0.3s ease !important;
}

textarea:not(.chat-input textarea):focus {
    border-color: rgba(16, 185, 129, 0.5) !important;
    box-shadow: 0 0 0 2px rgba(16, 185, 129, 0.08) !important;
}

/* User query display at top of results */
.user-query {
    background: linear-gradient(135deg, #1e3a2e 0%, #1a2332 100%);
    border: 1px solid #374151;
    border-radius: 12px;
    padding: 1.5rem;
    margin-bottom: 2rem;
}

.query-header {
    display: flex;
    justify-content: space-between;
    align-items: center;
    margin-bottom: 0.75rem;
}

.query-label {
    color: #10B981;
    font-size: 0.875rem;
    font-weight: 600;
    text-transform: uppercase;
    letter-spacing: 0.05em;
}

.query-actions {
    display: flex;
    gap: 0.5rem;
}

.query-text {
    color: #E5E7EB;
    font-size: 1.125rem;
    line-height: 1.6;
    font-weight: 400;
}

/* Simple answer styling */
.simple-answer {
    background: rgba(45, 45, 45, 0.3);
    border: 1px solid rgba(55, 65, 81, 0.3);
    border-radius: 12px;
    padding: 1.5rem;
    margin-top: 1rem;
    color: #E5E7EB;
    line-height: 1.7;
}

.answer-actions {
    display: flex;
    justify-content: flex-start;
    align-items: center;
    margin-top: 1rem;
    padding-top: 1rem;
    border-top: 1px solid #374151;
}

.error-message {
    background: rgba(239, 68, 68, 0.1);
    border: 1px solid rgba(239, 68, 68, 0.3);
    border-radius: 8px;
    padding: 1rem;
    margin-top: 1rem;
    color: #FCA5A5;
}

/* Report action buttons */
.report-actions {
    display: flex;
    justify-content: space-between;
    align-items: center;
    margin-top: 2rem;
    padding-top: 1.5rem;
    border-top: 1px solid #374151;
}

.left-actions {
    display: flex;
    gap: 0.5rem;
}

.feedback-buttons {
    display: flex;
    gap: 0.5rem;
}

/* Modern Minimal Icon Buttons */
.icon-btn {
    background: rgba(45, 45, 45, 0.4) !important;
    border: 1px solid rgba(55, 65, 81, 0.3) !important;
    color: #9CA3AF !important;
    padding: 0.5rem 0.875rem !important;
    border-radius: 8px !important;
    font-size: 0.85rem !important;
    font-weight: 500 !important;
    cursor: pointer !important;
    transition: all 0.2s ease !important;
    display: inline-flex;
    align-items: center;
    gap: 0.375rem;
    letter-spacing: 0.01em;
}

.icon-btn:hover {
    background: rgba(16, 185, 129, 0.12) !important;
    border-color: rgba(16, 185, 129, 0.3) !important;
    color: #10B981 !important;
    transform: translateY(-1px) !important;
    box-shadow: 0 2px 12px rgba(16, 185, 129, 0.15) !important;
}

.icon-btn:active {
    transform: translateY(0) !important;
}

.feedback-btn.active {
    background: rgba(16, 185, 129, 0.15) !important;
    border-color: rgba(16, 185, 129, 0.4) !important;
    color: #10B981 !important;
}

/* Welcome message */
.welcome-message {
    color: #9CA3AF;
    text-align: center;
    padding: 4rem 2rem;
    font-size: 1.125rem;
}

/* Controls Row - Prevent wrapping, keep all buttons inline */
.controls-row {
    display: flex !important;
    flex-wrap: nowrap !important;
    gap: 0.5rem;
    margin-top: 0.625rem;
    align-items: stretch;
    justify-content: flex-start;
    background: transparent !important;
    padding: 0;
    border-radius: 0;
    border: none;
    overflow: visible;
}

/* Model Dropdown - Inline with other buttons */
.model-dropdown {
    flex: 0 0 auto !important;
    width: fit-content !important;
    position: relative !important;
    overflow: visible !important;
}

/* Remove ALL backgrounds from ALL elements inside model-dropdown */
.model-dropdown,
.model-dropdown * {
    background: transparent !important;
    border: none !important;
    padding: 0 !important;
    margin: 0 !important;
    overflow: visible !important;
}

/* Ensure wrap is positioned context for dropdown */
.model-dropdown .wrap {
    position: relative !important;
    overflow: visible !important;
}

/* NOW style ONLY the .wrap to create the button */
.model-dropdown .wrap {
    background: rgba(55, 55, 55, 0.4) !important;
    border: 1px solid rgba(55, 65, 81, 0.4) !important;
    border-radius: 8px !important;
    min-width: 160px !important;
    max-width: 180px !important;
}

/* Style the actual dropdown input/button - restore padding */
.model-dropdown input[type="text"],
.model-dropdown .wrap input {
    background: transparent !important;
    border: none !important;
    color: #E5E7EB !important;
    padding: 0.4rem 1rem !important;
    font-size: 0.875rem !important;
    font-weight: 500 !important;
    cursor: pointer !important;
    height: auto !important;
    min-height: auto !important;
    /* Prevent typing but allow clicks */
    user-select: none !important;
    caret-color: transparent !important;
}

/* Prevent text editing on focus */
.model-dropdown input[type="text"]:focus {
    outline: none !important;
}

/* Hover state for wrapper - only when dropdown is closed */
.model-dropdown .wrap:hover {
    background: rgba(16, 185, 129, 0.1) !important;
    border-color: rgba(16, 185, 129, 0.6) !important;
    transform: translateY(-1px);
    box-shadow: 0 4px 12px rgba(16, 185, 129, 0.2);
}

.model-dropdown .wrap:hover input {
    color: #10B981 !important;
}

/* Disable wrap hover animation when dropdown is open */
.model-dropdown:focus-within .wrap,
.model-dropdown:has(ul[role="listbox"]) .wrap {
    pointer-events: none !important;
}

/* But keep dropdown menu interactive */
.model-dropdown ul[role="listbox"] {
    pointer-events: auto !important;
}

/* Dropdown menu - position directly above the button using transform */
.model-dropdown ul[role="listbox"] {
    background: rgb(45, 45, 45) !important;
    border: 1px solid rgba(55, 65, 81, 0.6) !important;
    border-radius: 8px !important;
    box-shadow: 0 8px 24px rgba(0, 0, 0, 0.6) !important;
    backdrop-filter: blur(12px);
    position: absolute !important;
    bottom: 100% !important;
    top: auto !important;
    left: 0 !important;
    right: 0 !important;
    margin-bottom: 0.5rem !important;
    margin-top: 0 !important;
    transform: translateZ(0) !important;
    z-index: 99999 !important;
    max-height: 300px !important;
    overflow-y: auto !important;
    pointer-events: auto !important;
    isolation: isolate !important;
}

.model-dropdown ul[role="listbox"] li {
    color: #9CA3AF !important;
    padding: 0.5rem 1rem !important;
    font-size: 0.875rem !important;
    cursor: pointer !important;
}

.model-dropdown ul[role="listbox"] li:hover {
    background: rgba(16, 185, 129, 0.1) !important;
    color: #10B981 !important;
}

.model-dropdown ul[role="listbox"] li[data-selected="true"] {
    background: rgba(16, 185, 129, 0.15) !important;
    color: #10B981 !important;
    font-weight: 600 !important;
}

/* Prevent chatbox from interfering when dropdown is open */
.model-dropdown:focus-within ~ .chat-input textarea,
.input-container-bottom:has(.model-dropdown ul[role="listbox"]) .chat-input textarea {
    pointer-events: none !important;
    z-index: 1 !important;
}

/* Mode buttons - Research and Search */
.mode-button {
    background: rgba(55, 55, 55, 0.4) !important;
    border: 1px solid rgba(55, 65, 81, 0.4) !important;
    color: #9CA3AF !important;
    border-radius: 8px !important;
    font-weight: 500 !important;
    padding: 0.4rem 0.875rem !important;
    font-size: 0.8125rem !important;
    letter-spacing: 0.01em !important;
    transition: all 0.2s cubic-bezier(0.4, 0, 0.2, 1) !important;
    box-shadow: none !important;
    text-transform: none !important;
    flex: 1;
    min-height: 0 !important;
    cursor: pointer !important;
    white-space: nowrap !important;
}

/* Active mode button */
.mode-button.active-mode {
    background: rgba(16, 185, 129, 0.15) !important;
    border-color: rgba(16, 185, 129, 0.6) !important;
    color: #10B981 !important;
    box-shadow: 0 0 0 2px rgba(16, 185, 129, 0.1) !important;
}

/* Hover state for mode buttons */
.mode-button:hover {
    background: rgba(16, 185, 129, 0.08) !important;
    border-color: rgba(16, 185, 129, 0.4) !important;
    color: #10B981 !important;
    transform: translateY(-1px) !important;
}

/* Active mode button hover - keep it glowing */
.mode-button.active-mode:hover {
    background: rgba(16, 185, 129, 0.2) !important;
    border-color: rgba(16, 185, 129, 0.7) !important;
    box-shadow: 0 0 0 2px rgba(16, 185, 129, 0.15) !important;
}

/* Attach button - icon only, compact */
.attach-button {
    background: rgba(55, 55, 55, 0.4) !important;
    border: 1px solid rgba(55, 65, 81, 0.4) !important;
    color: #9CA3AF !important;
    border-radius: 8px !important;
    font-weight: 500 !important;
    padding: 0.4rem 0.75rem !important;
    font-size: 1rem !important;
    transition: all 0.2s cubic-bezier(0.4, 0, 0.2, 1) !important;
    box-shadow: none !important;
    text-transform: none !important;
    flex: 0 0 auto;
    min-height: 0 !important;
    min-width: 2.5rem !important;
}

.attach-button:hover {
    background: rgba(59, 130, 246, 0.1) !important;
    border-color: rgba(59, 130, 246, 0.4) !important;
    color: #3B82F6 !important;
    transform: translateY(-1px) !important;
}

/* Clear button - icon only, compact */
.clear-button {
    background: rgba(55, 55, 55, 0.4) !important;
    border: 1px solid rgba(55, 65, 81, 0.4) !important;
    color: #9CA3AF !important;
    border-radius: 8px !important;
    font-weight: 500 !important;
    padding: 0.4rem 0.75rem !important;
    font-size: 1rem !important;
    transition: all 0.2s cubic-bezier(0.4, 0, 0.2, 1) !important;
    box-shadow: none !important;
    text-transform: none !important;
    flex: 0 0 auto;
    min-height: 0 !important;
    min-width: 2.5rem !important;
}

.clear-button:hover {
    background: rgba(239, 68, 68, 0.1) !important;
    border-color: rgba(239, 68, 68, 0.4) !important;
    color: #EF4444 !important;
    transform: translateY(-1px) !important;
}

/* Attachments display area - no space when empty */
.attachments-display-area {
    margin-bottom: 0;
    min-height: 0;
}

.attachments-display-area:not(:empty) {
    margin-bottom: 0.75rem;
}

.attachments-container {
    background: rgba(45, 45, 45, 0.3);
    border: 1px solid rgba(55, 65, 81, 0.3);
    border-radius: 8px;
    padding: 0.75rem;
    color: #E5E7EB;
    font-size: 0.875rem;
}

.attachments-header {
    color: #10B981;
    font-weight: 600;
    margin-bottom: 0.5rem;
}

/* Attachment badges */
.attachment-badge {
    display: inline-block;
    background: rgba(16, 185, 129, 0.1);
    border: 1px solid rgba(16, 185, 129, 0.3);
    border-radius: 6px;
    padding: 0.375rem 0.75rem;
    margin: 0.25rem 0.25rem 0.25rem 0;
    color: #10B981;
    font-size: 0.8rem;
    transition: all 0.2s;
}

.attachment-badge:hover {
    background: rgba(16, 185, 129, 0.15);
    border-color: rgba(16, 185, 129, 0.4);
}

.remove-attachment {
    background: none;
    border: none;
    color: #EF4444;
    cursor: pointer;
    margin-left: 0.5rem;
    font-weight: bold;
    font-size: 0.9rem;
    padding: 0;
    transition: color 0.2s;
}

.remove-attachment:hover {
    color: #DC2626;
}

/* Modern Minimal Tabs */
.tab-nav {
    background: transparent;
    border: none;
    margin-bottom: 0;
    padding: 0 0 0.5rem 0;
}

.tab-nav button {
    background: transparent !important;
    color: #6B7280 !important;
    border: none !important;
    border-radius: 8px !important;
    padding: 0.625rem 1.25rem !important;
    font-size: 0.9rem !important;
    font-weight: 500 !important;
    transition: all 0.2s ease !important;
    margin-right: 0.5rem !important;
}

.tab-nav button:hover {
    background: rgba(16, 185, 129, 0.08) !important;
    color: #9CA3AF !important;
}

.tab-nav button.selected {
    background: linear-gradient(135deg, rgba(16, 185, 129, 0.12) 0%, rgba(59, 130, 246, 0.08) 100%) !important;
    color: #10B981 !important;
    font-weight: 600 !important;
}

/* Modern Output Container */
.output-box {
    background: rgba(45, 45, 45, 0.3);
    border: 1px solid rgba(55, 65, 81, 0.3);
    border-radius: 12px;
    min-height: 400px;
    max-height: calc(100vh - 450px);
    overflow-y: auto;
    padding: 2rem;
    backdrop-filter: blur(8px);
}

/* Status messages */
.status-container {
    padding: 1rem;
}

.status-message {
    color: #10B981;
    font-size: 0.95rem;
    padding: 0.5rem 0;
    animation: fadeIn 0.3s ease;
}

@keyframes fadeIn {
    from { opacity: 0; transform: translateY(-10px); }
    to { opacity: 1; transform: translateY(0); }
}

/* Report content */
.output-box h1 {
    color: #10B981;
    font-size: 2rem;
    margin-bottom: 1rem;
    border-bottom: 2px solid #374151;
    padding-bottom: 0.5rem;
}

.output-box h2 {
    color: #3B82F6;
    font-size: 1.5rem;
    margin-top: 2rem;
    margin-bottom: 1rem;
}

.output-box h3 {
    color: #9CA3AF;
    font-size: 1.25rem;
    margin-top: 1.5rem;
    margin-bottom: 0.75rem;
}

.output-box p {
    color: #E5E7EB;
    line-height: 1.7;
    margin-bottom: 1rem;
}

.output-box a {
    color: #3B82F6;
    text-decoration: none;
    border-bottom: 1px solid transparent;
    transition: all 0.2s;
}

.output-box a:hover {
    color: #10B981;
    border-bottom-color: #10B981;
}

.output-box ul, .output-box ol {
    color: #E5E7EB;
    margin-left: 1.5rem;
    margin-bottom: 1rem;
}

.output-box li {
    margin-bottom: 0.5rem;
    line-height: 1.6;
}

/* Inline citations - clickable */
.citation-link {
    color: #3B82F6 !important;
    text-decoration: none;
    font-weight: 600;
    cursor: pointer;
    transition: all 0.2s;
    padding: 0 2px;
}

.citation-link:hover {
    color: #10B981 !important;
    text-decoration: underline;
}

/* References list */
.references-list {
    display: flex;
    flex-direction: column;
    gap: 1rem;
}

.reference-item {
    display: flex;
    align-items: flex-start;
    gap: 1rem;
    padding: 1rem;
    background: #1C1C1C;
    border: 1px solid #374151;
    border-radius: 8px;
    transition: all 0.3s ease;
    scroll-margin-top: 20px;
}

.reference-item:hover {
    border-color: #10B981;
    background: #2D2D2D;
    transform: translateX(5px);
}

/* Highlight effect when scrolling to a reference */
.reference-item:target {
    border-color: #10B981;
    background: #2D2D2D;
    box-shadow: 0 0 20px rgba(16, 185, 129, 0.3);
}

/* Smooth scrolling for the entire page */
html {
    scroll-behavior: smooth;
}

/* Make sure tabs are accessible */
button[role="tab"] {
    cursor: pointer;
}

.ref-number {
    color: #10B981;
    font-weight: 700;
    font-size: 1rem;
    min-width: 2rem;
}

.ref-link {
    color: #3B82F6;
    text-decoration: none;
    flex: 1;
    word-break: break-word;
    transition: color 0.2s;
}

.ref-link:hover {
    color: #10B981;
}

.no-references {
    color: #9CA3AF;
    text-align: center;
    padding: 2rem;
    font-style: italic;
}

/* Scrollbar styling */
::-webkit-scrollbar {
    width: 10px;
}

::-webkit-scrollbar-track {
    background: #1C1C1C;
}

::-webkit-scrollbar-thumb {
    background: #374151;
    border-radius: 5px;
}

::-webkit-scrollbar-thumb:hover {
    background: #10B981;
}

/* Loading animation */
@keyframes pulse {
    0%, 100% { opacity: 1; }
    50% { opacity: 0.5; }
}

.loading {
    animation: pulse 2s ease-in-out infinite;
}

/* Add JavaScript to make dropdown readonly and position it correctly */
</style>
<script>
document.addEventListener('DOMContentLoaded', function() {
    // Make model dropdown input readonly
    setTimeout(function() {
        const modelDropdown = document.querySelector('.model-dropdown input[type="text"]');
        if (modelDropdown) {
            modelDropdown.setAttribute('readonly', 'readonly');
            modelDropdown.addEventListener('keydown', function(e) {
                e.preventDefault();
            });
        }

        // Position dropdown menu above button
        const observer = new MutationObserver(function(mutations) {
            const dropdownMenu = document.querySelector('.model-dropdown ul[role="listbox"]');
            const buttonWrap = document.querySelector('.model-dropdown .wrap');

            if (dropdownMenu && buttonWrap) {
                const rect = buttonWrap.getBoundingClientRect();
                const menuHeight = dropdownMenu.offsetHeight;

                // Position above the button
                dropdownMenu.style.top = (rect.top - menuHeight - 8) + 'px';
                dropdownMenu.style.left = rect.left + 'px';
                dropdownMenu.style.width = rect.width + 'px';
            }
        });

        observer.observe(document.body, {
            childList: true,
            subtree: true
        });
    }, 500);
});
</script>
<style>
"""


# Build the UI
with gr.Blocks(theme=luntre_theme, css=custom_css, title="Luntre AI - Deep Research") as demo:
    # State to store current query for rewrite functionality
    current_query_state = gr.State("")
    # State to store conversation history
    conversation_history_state = gr.State([])
    # State to track current mode: "research" or "search"
    current_mode_state = gr.State("research")
    # State to store attachments
    attachments_state = gr.State([])

    # Header
    gr.HTML("""
        <div class="app-header">
            <div class="brand-name">Luntre AI</div>
            <div class="page-title">Deep Research</div>
        </div>
    """)

    # Output Section with Tabs (appears ABOVE input)
    with gr.Tabs(elem_classes="tab-nav"):
        with gr.Tab("Report"):
            report_output = gr.HTML(
                value="<div class='welcome-message'>Welcome! Enter your research query below to get started.</div>",
                elem_classes="output-box"
            )

        with gr.Tab("References"):
            references_output = gr.HTML(
                value="<div class='no-references'>No references yet. Run a research query to see sources.</div>",
                elem_classes="output-box"
            )

    # Input Section at the BOTTOM (chat-style)
    with gr.Group(elem_classes="input-container-bottom"):
        # Hidden file upload component
        file_upload = gr.File(
            label="Upload Files",
            file_types=[".txt", ".md", ".pdf", ".docx", ".doc", ".xlsx", ".xls", ".csv",
                       ".json", ".py", ".js", ".ts", ".java", ".cpp", ".html", ".log"],
            file_count="multiple",
            visible=False,
            elem_id="file-upload-input"
        )

        # Attachments display area (only visible when files are attached)
        attachments_display = gr.HTML(value="", elem_classes="attachments-display-area")

        query_input = gr.Textbox(
            label="",
            placeholder="What would you like to research? (e.g., 'What are the latest developments in quantum computing?')",
            lines=3,
            max_lines=5,
            show_label=False,
            container=False,
            elem_classes="chat-input",

        )

        with gr.Row(elem_classes="controls-row"):
            research_mode_btn = gr.Button("Deep Research", variant="primary", elem_classes="mode-button research-mode active-mode", elem_id="research-mode-btn")
            search_mode_btn = gr.Button("Search", variant="secondary", elem_classes="mode-button search-mode", elem_id="search-mode-btn")
            attach_btn = gr.Button("📎", variant="secondary", elem_classes="attach-button", elem_id="attach-btn")
            clear_conv_btn = gr.Button("🗑️", variant="secondary", elem_classes="clear-button")
            model_selector = gr.Dropdown(
                choices=[
                    "Gemini 2.5 Flash",
                    "Gemini 2.0 Flash",
                    "Gemini 2.0 Pro",
                    "Llama 3.3",
                ],
                value="Gemini 2.5 Flash",
                label="",
                show_label=False,
                container=False,
                interactive=True,
                elem_classes="model-dropdown"
            )
            submit_button = gr.Button("Submit", visible=False, elem_id="submit-button") # NEW HIDDEN BUTTON

        # Hidden buttons for programmatic access
        edit_btn = gr.Button("Edit", variant="secondary", visible=False, elem_id="edit-query-btn")
        rewrite_btn = gr.Button("Rewrite", variant="secondary", visible=False, elem_id="rewrite-btn")

    # Event handlers
    def update_query_state(query):
        """Store the current query in state"""
        return query

    def load_query_for_edit(stored_query):
        """Load stored query back into input for editing"""
        return stored_query

    def clear_input():
        """Clear the input box"""
        return ""

    def add_query_to_history(query, conversation_history):
        """Add user query to conversation history"""
        import datetime
        conversation_history.append({
            "type": "query",
            "content": query,
            "timestamp": datetime.datetime.now().isoformat()
        })
        return conversation_history

    def add_report_to_history(report_html, conversation_history):
        """Add report to conversation history"""
        import datetime
        conversation_history.append({
            "type": "report",
            "content": report_html,
            "timestamp": datetime.datetime.now().isoformat()
        })
        return conversation_history

    def add_simple_search_to_history(answer_html, conversation_history):
        """Add simple search answer to conversation history"""
        import datetime
        conversation_history.append({
            "type": "simple_search",
            "content": answer_html,
            "timestamp": datetime.datetime.now().isoformat()
        })
        return conversation_history

    def clear_conversation():
        """Clear conversation history"""
        return []

    def switch_to_research_mode():
        """Switch to research mode"""
        return "research"

    def switch_to_search_mode():
        """Switch to search mode"""
        return "search"

    def handle_file_upload(files, current_attachments):
        """Process uploaded files and update attachments state"""
        if not files:
            return current_attachments, format_attachments_display(current_attachments)

        # Handle single file or list of files
        if not isinstance(files, list):
            files = [files]

        for file in files:
            if file is not None:
                # Process the file
                file_data = process_file(file.name)
                if file_data:
                    current_attachments.append(file_data)

        display_html = format_attachments_display(current_attachments)
        return current_attachments, display_html

    def format_attachments_display(attachments):
        """Generate HTML for attachment badges"""
        if not attachments:
            return ""

        html = '<div class="attachments-container">'
        html += f'<div class="attachments-header">📎 Attached Files ({len(attachments)})</div>'

        for idx, att in enumerate(attachments):
            size_str = format_file_size(att['size_bytes'])
            icon = get_file_icon(att['file_type'])
            html += f'''
            <span class="attachment-badge" id="attachment-{idx}">
                {icon} {att['filename']} ({size_str})
                <button class="remove-attachment" onclick="document.getElementById('remove-att-{idx}').click();">✕</button>
            </span>
            '''

        html += '</div>'
        return html

    def remove_attachment(attachments, idx):
        """Remove attachment at given index"""
        if 0 <= idx < len(attachments):
            attachments.pop(idx)
        return attachments, format_attachments_display(attachments)

    def clear_attachments():
        """Clear all attachments"""
        return [], ""

    # Mode switching - Research button
    research_mode_btn.click(
        fn=switch_to_research_mode,
        inputs=[],
        outputs=[current_mode_state],
        queue=False,
        js="""
        () => {
            // Add active class to research button, remove from search button
            const researchBtn = document.getElementById('research-mode-btn');
            const searchBtn = document.getElementById('search-mode-btn');
            if (researchBtn) researchBtn.classList.add('active-mode');
            if (searchBtn) searchBtn.classList.remove('active-mode');
        }
        """
    )

    # Mode switching - Search button
    search_mode_btn.click(
        fn=switch_to_search_mode,
        inputs=[],
        outputs=[current_mode_state],
        queue=False,
        js="""
        () => {
            // Add active class to search button, remove from research button
            const researchBtn = document.getElementById('research-mode-btn');
            const searchBtn = document.getElementById('search-mode-btn');
            if (researchBtn) researchBtn.classList.remove('active-mode');
            if (searchBtn) searchBtn.classList.add('active-mode');
        }
        """
    )

    # Attach button - trigger file upload dialog
    attach_btn.click(
        fn=None,
        inputs=[],
        outputs=[],
        js="""
        () => {
            // Trigger the hidden file input
            const fileInput = document.getElementById('file-upload-input');
            if (fileInput) {
                const actualInput = fileInput.querySelector('input[type="file"]');
                if (actualInput) {
                    actualInput.click();
                }
            }
        }
        """
    )

    # File upload handler
    file_upload.change(
        fn=handle_file_upload,
        inputs=[file_upload, attachments_state],
        outputs=[attachments_state, attachments_display],
        queue=False
    )

    # Edit and rerun
    edit_event = edit_btn.click(
        fn=load_query_for_edit,
        inputs=[current_query_state],
        outputs=[query_input],
        queue=False
    )

    # Rewrite (run again with same query)
    rewrite_event = rewrite_btn.click(
        fn=run_research,
        inputs=[current_query_state, model_selector, conversation_history_state, attachments_state],
        outputs=[report_output, references_output]
    )

    # Clear conversation
    def reset_conversation():
        """Reset conversation, display welcome message, and clear attachments"""
        return (
            [],  # conversation_history
            "<div class='welcome-message'>Welcome! Enter your research query below to get started.</div>",  # report_output
            "<div class='no-references'>No references yet. Run a research query to see sources.</div>",  # references_output
            "research",  # current_mode
            [],  # attachments
            ""  # attachments_display
        )

    clear_conv_event = clear_conv_btn.click(
        fn=reset_conversation,
        inputs=[],
        outputs=[conversation_history_state, report_output, references_output, current_mode_state, attachments_state, attachments_display],
        queue=False,
        js="""
        () => {
            // Reset to research mode visually
            const researchBtn = document.getElementById('research-mode-btn');
            const searchBtn = document.getElementById('search-mode-btn');
            if (researchBtn) researchBtn.classList.add('active-mode');
            if (searchBtn) searchBtn.classList.remove('active-mode');
        }
        """
    )

    # Enter key submission - runs based on current mode
    # We need separate event handlers for research and search modes
    # Store the event handler reference
    submit_event_state = gr.State(None)

    def run_based_on_mode(query, model, history, attachments, mode):
        """Wrapper to route to correct function based on mode"""
        if mode == "research":
            return run_research(query, model, history, attachments)
        else:
            return run_simple_search(query, model, history, attachments)

    submit_button.click(
        fn=update_query_state,
        inputs=[query_input],
        outputs=[current_query_state],
        queue=False
    ).then(
        fn=add_query_to_history,
        inputs=[current_query_state, conversation_history_state],
        outputs=[conversation_history_state],
        queue=False
    ).then(
        fn=clear_input,
        inputs=[],
        outputs=[query_input],
        queue=False
    ).then(
        fn=run_based_on_mode,
        inputs=[current_query_state, model_selector, conversation_history_state, attachments_state, current_mode_state],
        outputs=[report_output, references_output]
    ).then(
        fn=lambda mode: "search" if mode == "research" else mode,
        inputs=[current_mode_state],
        outputs=[current_mode_state],
        queue=False,
        js="""
        (mode) => {
            // If we just ran research, switch to search mode
            if (mode === "research") {
                const researchBtn = document.getElementById('research-mode-btn');
                const searchBtn = document.getElementById('search-mode-btn');
                if (researchBtn) researchBtn.classList.remove('active-mode');
                if (searchBtn) searchBtn.classList.add('active-mode');
            }
            return mode === "research" ? "search" : mode;
        }
        """
    )


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

    gr.HTML("""
    <script>
    document.addEventListener('DOMContentLoaded', function() {
        console.log('DOM Content Loaded - Attempting to set up Enter key listener.');
        const queryInput = document.querySelector('.chat-input textarea');
        const submitButton = document.getElementById('submit-button');

        if (queryInput) {
            console.log('queryInput found:', queryInput);
        } else {
            console.log('queryInput NOT found.');
        }

        if (submitButton) {
            console.log('submitButton found:', submitButton);
        } else {
            console.log('submitButton NOT found.');
        }

        if (queryInput && submitButton) {
            queryInput.addEventListener('keydown', function(e) {
                console.log('Keydown event:', e.key, 'Shift:', e.shiftKey);
                if (e.key === 'Enter' && !e.shiftKey) {
                    console.log('Enter (without Shift) pressed. Preventing default and clicking submit button.');
                    e.preventDefault(); // Prevent new line
                    submitButton.click(); // Trigger the hidden submit button
                } else if (e.key === 'Enter' && e.shiftKey) {
                    console.log('Shift+Enter pressed. Allowing new line.');
                    // Allow default behavior for new line
                }
            });
            console.log('Keydown listener attached to queryInput.');
        } else {
            console.log('Could not attach keydown listener: queryInput or submitButton not found.');
        }
    });
    </script>
    """)