"""Streamlit Chat Interface for MCP Servers

A visual interface for testing deployed MCP servers on GCP Cloud Run.
Provides a Claude Desktop-like experience for bioinformatics workflows.
"""

import streamlit as st
import os
import uuid
import tempfile
from typing import List, Dict
from datetime import datetime
from pathlib import Path
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

from utils import (
    MCP_SERVERS,
    get_server_config,
    get_tools_config,
    get_server_categories,
    EXAMPLE_PROMPTS,
    ChatHandler,
    build_orchestration_trace,
    render_trace,
    render_trace_export,
    check_and_render_downloads
)

# Import provider system
from providers import (
    get_provider,
    get_available_providers,
    ChatMessage,
    GEMINI_AVAILABLE
)

# Import authentication and audit logging
from utils.auth import require_authentication, display_user_info, display_logout_button
from utils.audit_logger import get_audit_logger

# Import file validation and GCS handling
from utils.file_validator import validate_uploaded_file, sanitize_filename
from utils.gcs_handler import (
    is_gcs_path,
    validate_gcs_uri,
    get_gcs_file_metadata,
    get_gcs_file_content,
    get_gcs_files_metadata
)

# Page configuration
st.set_page_config(
    page_title="MCP Chat - Precision Medicine",
    page_icon="🧬",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for better styling
st.markdown("""
<style>
    .stChatMessage {
        padding: 1rem;
        border-radius: 0.5rem;
    }
    .server-card {
        padding: 1rem;
        border-radius: 0.5rem;
        border: 1px solid #ddd;
        margin: 0.5rem 0;
    }
    .production {
        border-left: 4px solid #28a745;
    }
    .mock {
        border-left: 4px solid #ffc107;
    }
</style>
""", unsafe_allow_html=True)


def is_cloud_run() -> bool:
    """Detect if running on GCP Cloud Run.

    Returns:
        bool: True if deployed on Cloud Run, False for local development
    """
    return os.getenv("K_SERVICE") is not None


def initialize_session_state():
    """Initialize Streamlit session state."""
    if "messages" not in st.session_state:
        st.session_state.messages = []
    if "selected_servers" not in st.session_state:
        st.session_state.selected_servers = ["fgbio"]

    # Initialize provider selection (Cloud Run only)
    if "llm_provider" not in st.session_state:
        st.session_state.llm_provider = "claude"  # Default to Claude

    if "llm_model" not in st.session_state:
        st.session_state.llm_model = "claude-sonnet-4-6"  # Default model

    # Initialize provider instance
    if "provider_instance" not in st.session_state:
        try:
            st.session_state.provider_instance = get_provider(
                provider_name=st.session_state.llm_provider
            )
        except (ValueError, ImportError) as e:
            st.session_state.provider_instance = None
            if is_cloud_run():
                st.error(f"Failed to initialize {st.session_state.llm_provider} provider: {str(e)}")

    # Keep chat_handler for backward compatibility (local development)
    if "chat_handler" not in st.session_state:
        api_key = os.getenv("ANTHROPIC_API_KEY")
        if api_key:
            try:
                st.session_state.chat_handler = ChatHandler(api_key)
            except ValueError:
                st.session_state.chat_handler = None
        else:
            st.session_state.chat_handler = None

    # Initialize session tracking
    if "session_id" not in st.session_state:
        st.session_state.session_id = f"sess_{uuid.uuid4().hex[:12]}"
    if "session_start_time" not in st.session_state:
        st.session_state.session_start_time = datetime.utcnow()
    if "total_queries" not in st.session_state:
        st.session_state.total_queries = 0
    if "total_tokens" not in st.session_state:
        st.session_state.total_tokens = 0
    if "total_cost" not in st.session_state:
        st.session_state.total_cost = 0.0

    # Initialize audit logger
    if "audit_logger" not in st.session_state:
        st.session_state.audit_logger = get_audit_logger()

    # Initialize trace storage
    if "traces" not in st.session_state:
        st.session_state.traces = {}  # message_index -> OrchestrationTrace

    # Initialize uploaded files storage
    if "uploaded_files" not in st.session_state:
        st.session_state.uploaded_files = {}  # filename -> {'path': temp_path, 'metadata': dict}


def prepare_file_for_mcp(uploaded_file, metadata: Dict) -> str:
    """Prepare uploaded file for MCP server access.

    Writes validated file to a temporary directory and returns the path
    for MCP servers to access.

    Args:
        uploaded_file: Streamlit UploadedFile object
        metadata: File metadata dict from validation

    Returns:
        str: Absolute path to the temporary file

    Example:
        >>> temp_path = prepare_file_for_mcp(uploaded_file, metadata)
        >>> # Pass temp_path to MCP tool call
    """
    # Create a temp directory if it doesn't exist
    temp_dir = Path(tempfile.gettempdir()) / "mcp_uploads"
    temp_dir.mkdir(exist_ok=True)

    # Use sanitized filename
    sanitized_name = metadata['sanitized_filename']
    temp_path = temp_dir / sanitized_name

    # Write file content
    content = uploaded_file.read()
    uploaded_file.seek(0)  # Reset file pointer

    with open(temp_path, 'wb') as f:
        f.write(content)

    return str(temp_path)


def render_sidebar():
    """Render sidebar with server selection and settings."""
    with st.sidebar:
        st.title("🧬 MCP Chat")
        st.markdown("---")

        # User info (SSO authentication)
        if "user" in st.session_state:
            display_user_info(st.session_state.user, location="sidebar")
            display_logout_button()

        st.markdown("---")

        # API Key status
        if st.session_state.chat_handler:
            st.success("✅ API Key Configured")
        else:
            st.error("❌ API Key Missing")
            st.info("Set ANTHROPIC_API_KEY environment variable")

        st.markdown("---")

        # Provider selection (available in all modes)
        st.subheader("LLM Provider")

        available_providers = get_available_providers()

        # Build provider options - show ALL providers, not just available ones
        provider_options = []
        provider_mapping = {}
        for provider_key, provider_info in available_providers.items():
            display_name = provider_info["name"]
            # Add indicator if API key is missing
            if not provider_info["available"]:
                display_name += " (API key required)"
            provider_options.append(display_name)
            provider_mapping[display_name] = provider_key

        if not provider_options:
            st.error("No LLM providers configured.")
            st.stop()

        # Provider dropdown
        current_provider_display = available_providers[st.session_state.llm_provider]["name"]
        selected_provider_display = st.selectbox(
            "Select Provider",
            options=provider_options,
            index=provider_options.index(current_provider_display) if current_provider_display in provider_options else 0,
            help="Choose LLM provider for analysis"
        )

        # Update provider if changed
        selected_provider_key = provider_mapping[selected_provider_display]
        if selected_provider_key != st.session_state.llm_provider:
            # Check if provider is available before switching
            provider_info = available_providers[selected_provider_key]
            if not provider_info["available"]:
                st.error(f"⚠️ {provider_info['name']} API key not configured")
                if "note" in provider_info:
                    st.info(provider_info["note"])
                st.info("Configure API key in deployment environment variables and redeploy")
                st.stop()

            st.session_state.llm_provider = selected_provider_key
            try:
                st.session_state.provider_instance = get_provider(provider_name=selected_provider_key)
                st.success(f"✅ Switched to {selected_provider_display}")
                st.rerun()
            except (ValueError, ImportError) as e:
                st.error(f"❌ Failed to initialize {selected_provider_display}: {str(e)}")
                st.stop()

        st.markdown("---")

        # Model selection
        st.subheader("Model Settings")

        # Dynamic model selection based on provider
        if st.session_state.provider_instance:
            available_providers = get_available_providers()
            provider_info = available_providers.get(st.session_state.llm_provider, {})
            available_models = provider_info.get("models", [])

            if available_models:
                model_options = [m["id"] for m in available_models]
                model_labels = [m["name"] for m in available_models]

                # Create mapping for display
                model_display_map = {m["name"]: m["id"] for m in available_models}

                current_model = st.session_state.llm_model
                try:
                    current_index = model_options.index(current_model)
                except ValueError:
                    current_index = 0
                    st.session_state.llm_model = model_options[0]

                selected_model_label = st.selectbox(
                    f"{provider_info['name']} Model",
                    options=model_labels,
                    index=current_index,
                    help=f"Select {provider_info['name']} model"
                )
                model = model_display_map[selected_model_label]
                st.session_state.llm_model = model
            else:
                model = provider_info.get("default_model", "claude-sonnet-4-6")
                st.session_state.llm_model = model
        else:
            # Local development - Claude only
            model = st.selectbox(
                "Claude Model",
                options=["claude-sonnet-4-6", "claude-opus-4-6", "claude-haiku-4-5"],
                index=0,
                help="Select Claude model. Sonnet recommended for most tasks."
            )

        max_tokens = st.slider(
            "Max Tokens",
            min_value=1024,
            max_value=8192,
            value=4096,
            step=512,
            help="Maximum response length"
        )

        st.markdown("---")

        # MCP Server Selection
        st.subheader("MCP Servers")

        categories = get_server_categories()

        # Production servers
        st.markdown("**Production Servers** (Real Analysis)")
        for server_name in categories["Production Servers (Real Analysis)"]:
            server = MCP_SERVERS[server_name]
            selected = st.checkbox(
                f"{server_name} ({server['tools_count']} tools)",
                value=server_name in st.session_state.selected_servers,
                key=f"checkbox_{server_name}",
                help=server['description']
            )
            if selected and server_name not in st.session_state.selected_servers:
                st.session_state.selected_servers.append(server_name)
            elif not selected and server_name in st.session_state.selected_servers:
                st.session_state.selected_servers.remove(server_name)

        st.markdown("**Mock Servers** (Demo Only)")
        with st.expander("Show Mock Servers"):
            for server_name in categories["Mock Servers (Workflow Demo)"]:
                server = MCP_SERVERS[server_name]
                selected = st.checkbox(
                    f"{server_name} ({server['tools_count']} tools)",
                    value=server_name in st.session_state.selected_servers,
                    key=f"checkbox_{server_name}",
                    help=server['description']
                )
                if selected and server_name not in st.session_state.selected_servers:
                    st.session_state.selected_servers.append(server_name)
                elif not selected and server_name in st.session_state.selected_servers:
                    st.session_state.selected_servers.remove(server_name)

        # Show selected count
        st.info(f"**{len(st.session_state.selected_servers)}** servers selected")

        st.markdown("---")

        # Orchestration Trace Settings
        st.subheader("🔍 Orchestration Trace")
        show_trace = st.toggle(
            "Show trace for responses",
            value=False,
            help="Display which MCP servers were called and in what order"
        )

        trace_style = "log"  # Default
        if show_trace:
            trace_style = st.selectbox(
                "Trace style",
                options=["log", "cards", "timeline", "mermaid"],
                format_func=lambda x: {
                    "log": "📝 Log View",
                    "cards": "🎴 Card View",
                    "timeline": "📈 Timeline View",
                    "mermaid": "📊 Sequence Diagram"
                }.get(x, x),
                help="Choose how to display the orchestration trace"
            )

        st.markdown("---")

        # File Upload Section
        st.subheader("📁 File Upload")
        st.caption("Upload bioinformatics data files")

        uploaded_files = st.file_uploader(
            "Choose files",
            type=['fasta', 'fa', 'fna', 'fastq', 'fq', 'vcf', 'gff', 'gtf', 'bed',
                  'csv', 'tsv', 'tab', 'txt', 'json', 'h5ad', 'h5',
                  'png', 'jpg', 'jpeg', 'tiff', 'tif'],
            accept_multiple_files=True,
            help="Upload sequence files, annotations, tabular data, or images"
        )

        if uploaded_files:
            st.caption(f"{len(uploaded_files)} file(s) selected")

            for uploaded_file in uploaded_files:
                # Validate each file
                is_valid, errors, metadata = validate_uploaded_file(uploaded_file)

                if is_valid:
                    # Store validated file
                    temp_path = prepare_file_for_mcp(uploaded_file, metadata)
                    st.session_state.uploaded_files[metadata['sanitized_filename']] = {
                        'path': temp_path,
                        'metadata': metadata,
                        'original_name': uploaded_file.name
                    }

                    # Show success with file info
                    with st.expander(f"✅ {uploaded_file.name}", expanded=False):
                        st.success("Valid file")
                        col1, col2 = st.columns(2)
                        col1.metric("Size", f"{metadata['size_mb']:.2f} MB")
                        col2.metric("Type", metadata['extension'])
                        if not metadata.get('is_binary', False):
                            st.caption(f"Lines: {metadata.get('line_count', 'N/A')}")
                else:
                    # Show validation errors
                    with st.expander(f"❌ {uploaded_file.name}", expanded=True):
                        st.error("Invalid file")
                        for error in errors:
                            st.write(f"- {error}")

        # GCS Path Input Section
        st.caption("Or provide GCS bucket path")

        gcs_uri = st.text_input(
            "GCS URI",
            placeholder="gs://bucket-name/path/to/file.fastq or gs://bucket-name/path/to/folder/",
            help="Enter a Google Cloud Storage URI to a file or folder. For folders, all files will be loaded.",
            key="gcs_uri_input"
        )

        if gcs_uri and gcs_uri.strip():
            # Get metadata for file(s) - handles both single files and folders
            success, metadata_list, meta_error = get_gcs_files_metadata(gcs_uri, use_mock=False)

            if success:
                # Store all files found
                for metadata in metadata_list:
                    file_gcs_uri = metadata['gcs_uri']

                    # Store GCS file reference
                    st.session_state.uploaded_files[metadata['sanitized_filename']] = {
                        'path': file_gcs_uri,  # Store GCS path directly
                        'metadata': metadata,
                        'original_name': metadata['filename'],
                        'source': 'gcs'
                    }

                # Show success - different display for single vs multiple files
                if len(metadata_list) == 1:
                    # Single file
                    metadata = metadata_list[0]
                    with st.expander(f"✅ {metadata['filename']} (GCS)", expanded=False):
                        st.success("Valid GCS URI")
                        col1, col2 = st.columns(2)
                        col1.metric("Bucket", metadata['bucket'])
                        col2.metric("Type", metadata['extension'])
                        st.caption(f"Path: {metadata['blob_path']}")

                        # Try to get content for small text files
                        if not metadata.get('is_binary', False) and metadata['size_mb'] < 0.05:
                            success_content, content, content_error = get_gcs_file_content(metadata['gcs_uri'], max_size_bytes=50000)
                            if success_content:
                                st.caption(f"✅ Content loaded ({len(content)} chars)")
                                # Store content in metadata for inline inclusion
                                metadata['_gcs_content'] = content
                            else:
                                st.caption(f"⚠️ Content not loaded: {content_error}")
                else:
                    # Multiple files from folder
                    with st.expander(f"✅ {len(metadata_list)} files loaded from folder", expanded=True):
                        st.success(f"Found {len(metadata_list)} files in GCS folder")

                        # Show file list
                        for metadata in metadata_list:
                            col1, col2, col3 = st.columns([3, 1, 1])
                            col1.write(f"📄 {metadata['filename']}")
                            col2.write(f"{metadata['size_mb']:.2f} MB")
                            col3.write(metadata['extension'])
            else:
                st.error(f"Error: {meta_error}")

        # Show currently uploaded files
        if st.session_state.uploaded_files:
            st.caption(f"📎 {len(st.session_state.uploaded_files)} file(s) available for MCP")

            # Button to clear all uploaded files
            if st.button("Clear Files", key="clear_files"):
                # Clean up temp files
                for file_info in st.session_state.uploaded_files.values():
                    try:
                        if os.path.exists(file_info['path']):
                            os.remove(file_info['path'])
                    except Exception:
                        pass
                st.session_state.uploaded_files = {}
                st.rerun()

        st.markdown("---")

        # Example prompts
        st.subheader("Example Prompts")
        selected_example = st.selectbox(
            "Load an example:",
            options=[""] + list(EXAMPLE_PROMPTS.keys()),
            index=0
        )

        if selected_example:
            # Show preview of the prompt
            preview = EXAMPLE_PROMPTS[selected_example][:120]
            if len(EXAMPLE_PROMPTS[selected_example]) > 120:
                preview += "..."
            st.caption(preview)

            if st.button("Send Prompt", use_container_width=True):
                st.session_state.pending_example = EXAMPLE_PROMPTS[selected_example]
                st.rerun()

        st.markdown("---")

        # Session statistics
        st.subheader("Session Stats")
        col1, col2 = st.columns(2)
        col1.metric("Queries", st.session_state.total_queries)
        col2.metric("Tokens", f"{st.session_state.total_tokens:,}")
        st.metric("Cost", f"${st.session_state.total_cost:.4f}")

        # Session duration
        if st.session_state.session_start_time:
            duration = datetime.utcnow() - st.session_state.session_start_time
            duration_mins = int(duration.total_seconds() / 60)
            st.caption(f"Session: {duration_mins} min")

        st.markdown("---")

        # Clear chat button
        if st.button("🗑️ Clear Chat", use_container_width=True):
            st.session_state.messages = []
            st.session_state.traces = {}  # Clear traces too
            st.rerun()

        st.markdown("---")

        # Token Benchmark section
        st.subheader("Token Benchmark")
        with st.expander("Run benchmark", expanded=False):
            st.caption("Compare token usage across prompts and providers")

            # Prompt selection
            all_prompt_names = list(EXAMPLE_PROMPTS.keys())
            selected_prompts = st.multiselect(
                "Prompts to benchmark",
                options=all_prompt_names,
                default=all_prompt_names,
                key="benchmark_prompts"
            )

            # Provider selection
            benchmark_providers = []
            avail = get_available_providers()
            if avail.get("claude", {}).get("available"):
                run_claude = st.checkbox("Claude", value=True, key="bench_claude")
                if run_claude:
                    claude_models = avail["claude"].get("models", [])
                    claude_model_id = st.selectbox(
                        "Claude model",
                        [m["id"] for m in claude_models],
                        key="bench_claude_model"
                    )
                    benchmark_providers.append(("claude", claude_model_id))
            if avail.get("gemini", {}).get("available"):
                run_gemini = st.checkbox("Gemini", value=True, key="bench_gemini")
                if run_gemini:
                    gemini_models = avail["gemini"].get("models", [])
                    gemini_model_id = st.selectbox(
                        "Gemini model",
                        [m["id"] for m in gemini_models],
                        key="bench_gemini_model"
                    )
                    benchmark_providers.append(("gemini", gemini_model_id))

            # Track completed phases
            phases_done = st.session_state.get("benchmark_phases_done", set())

            # Show progress if a phase is running
            if st.session_state.get("benchmark_running"):
                phase = st.session_state.get("benchmark_phase", "")
                queue = st.session_state.get("benchmark_queue", [])
                done = st.session_state.get("benchmark_done", [])
                total = len(queue) + len(done)
                if total > 0:
                    st.progress(len(done) / total,
                                text=f"{phase.title()}: {len(done)+1}/{total}...")
                if st.button("Cancel", key="cancel_benchmark",
                             use_container_width=True):
                    st.session_state["benchmark_running"] = False
                    st.session_state["benchmark_queue"] = []
                    st.rerun()
            else:
                # Phase status indicators
                for phase_name in ["cold", "warm", "repeat"]:
                    if phase_name in phases_done:
                        st.caption(f"  {phase_name.title()} — done")

                # Phase buttons
                can_run = selected_prompts and benchmark_providers
                col1, col2, col3 = st.columns(3)
                with col1:
                    if st.button("Cold", key="run_cold",
                                 use_container_width=True,
                                 disabled=not can_run):
                        _start_benchmark_phase(
                            selected_prompts, benchmark_providers, "cold")
                with col2:
                    if st.button("Warm", key="run_warm",
                                 use_container_width=True,
                                 disabled=not can_run or "cold" not in phases_done):
                        _start_benchmark_phase(
                            selected_prompts, benchmark_providers, "warm")
                with col3:
                    if st.button("Repeat", key="run_repeat",
                                 use_container_width=True,
                                 disabled=not can_run or "warm" not in phases_done):
                        _start_benchmark_phase(
                            selected_prompts, benchmark_providers, "repeat")

                # Reset button (only if any phase completed)
                if phases_done:
                    if st.button("Reset All Results", key="reset_benchmark",
                                 use_container_width=True):
                        st.session_state.pop("benchmark_phases_done", None)
                        st.session_state.pop("benchmark_all_results", None)
                        st.session_state.pop("benchmark_results", None)
                        st.session_state.pop("benchmark_csv", None)
                        st.rerun()

    return model, max_tokens, show_trace, trace_style


def render_server_status():
    """Render server status cards."""
    st.subheader("Active MCP Servers")

    if not st.session_state.selected_servers:
        st.warning("⚠️ No servers selected. Select servers from the sidebar.")
        return

    cols = st.columns(3)
    for idx, server_name in enumerate(st.session_state.selected_servers):
        server = MCP_SERVERS[server_name]
        with cols[idx % 3]:
            status_class = "production" if server["status"] == "production" else "mock"
            st.markdown(f"""
            <div class="server-card {status_class}">
                <h4>{server_name}</h4>
                <p>{server['description']}</p>
                <small>{server['tools_count']} tools • {server['status']}</small>
            </div>
            """, unsafe_allow_html=True)


def render_chat_history(show_trace: bool = False, trace_style: str = "log"):
    """Render chat message history.

    Args:
        show_trace: Whether to show orchestration traces
        trace_style: Style for rendering traces
    """
    for i, message in enumerate(st.session_state.messages):
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

            # Show usage info if available
            if "usage" in message and message["usage"]:
                with st.expander("Token Usage"):
                    usage = message["usage"]
                    cache_read = usage.get("cache_read_tokens", 0)
                    cache_creation = usage.get("cache_creation_tokens", 0)
                    iterations = usage.get("iterations", 1)

                    if cache_read > 0:
                        col1, col2, col3, col4 = st.columns(4)
                        col1.metric("Input", f"{usage.get('input_tokens', 0):,}")
                        col2.metric("Output", f"{usage.get('output_tokens', 0):,}")
                        col3.metric("Total", f"{usage.get('total_tokens', 0):,}")
                        col4.metric("Cached", f"{cache_read:,}")
                    else:
                        col1, col2, col3 = st.columns(3)
                        col1.metric("Input", f"{usage.get('input_tokens', 0):,}")
                        col2.metric("Output", f"{usage.get('output_tokens', 0):,}")
                        col3.metric("Total", f"{usage.get('total_tokens', 0):,}")

                    # Show extra details if multi-iteration or cache write
                    details = []
                    if iterations > 1:
                        details.append(f"Iterations: {iterations}")
                    if cache_creation > 0:
                        details.append(f"Cache Write: {cache_creation:,} tokens")
                    if cache_read > 0:
                        hit_rate = cache_read / (cache_read + usage.get("input_tokens", 1))
                        details.append(f"Cache Hit Rate: {hit_rate:.0%}")
                    if details:
                        st.caption(" | ".join(details))

            # Show trace for assistant messages if enabled
            if message["role"] == "assistant" and show_trace:
                if i in st.session_state.traces:
                    trace = st.session_state.traces[i]
                    render_trace(trace, style=trace_style)

                    # Add export buttons if trace has data
                    if trace.tool_calls:
                        render_trace_export(trace, trace_id=i)

            # Check for downloadable files (PDFs from patient reports)
            if message["role"] == "assistant":
                trace = st.session_state.traces.get(i) if hasattr(st.session_state, 'traces') else None
                check_and_render_downloads(
                    trace=trace,
                    content=message["content"],
                    key_prefix=f"msg_{i}"
                )


def handle_user_input(prompt: str, model: str, max_tokens: int):
    """Handle user input and get response from Claude.

    Args:
        prompt: User's message
        model: Claude model to use
        max_tokens: Maximum tokens for response
    """
    # Get user info for audit logging
    user = st.session_state.get("user")
    if not user:
        # Development mode - create mock user
        user = {
            "email": "developer@localhost",
            "user_id": "dev_user",
            "display_name": "Developer"
        }

    # Add user message to history
    st.session_state.messages.append({"role": "user", "content": prompt})

    # Display user message immediately (before spinner starts)
    with st.chat_message("user"):
        st.markdown(prompt)

    # Get MCP server config
    mcp_servers = get_server_config(st.session_state.selected_servers)

    if not mcp_servers:
        with st.chat_message("assistant"):
            error_msg = "⚠️ No MCP servers selected. Please select at least one server from the sidebar."
            st.error(error_msg)
            st.session_state.messages.append({"role": "assistant", "content": error_msg})
        return

    # Log the query (audit trail)
    st.session_state.audit_logger.log_mcp_query(
        user_email=user["email"],
        user_id=user["user_id"],
        servers=st.session_state.selected_servers,
        prompt=prompt,
        model=model,
        session_id=st.session_state.session_id
    )

    # Track query start time
    import time
    query_start_time = datetime.utcnow()
    start_time = time.time()

    # Show thinking indicator (outside chat message so it shows before rerun)
    with st.spinner("🤔 Thinking..."):
        try:
            # Use provider abstraction for all modes (supports mock MCP)
            if st.session_state.provider_instance:
                # Convert messages to ChatMessage objects
                chat_messages = [
                    ChatMessage(role=msg["role"], content=msg["content"])
                    for msg in st.session_state.messages
                ]

                # Send via provider abstraction
                chat_response = st.session_state.provider_instance.send_message(
                    messages=chat_messages,
                    mcp_servers=mcp_servers,
                    model=model,
                    max_tokens=max_tokens,
                    uploaded_files=st.session_state.uploaded_files
                )

                # Calculate response time
                duration_ms = (time.time() - start_time) * 1000

                # Extract response content and usage
                response_text = chat_response.content
                usage = None
                if chat_response.usage:
                    usage = {
                        "input_tokens": chat_response.usage.input_tokens,
                        "output_tokens": chat_response.usage.output_tokens,
                        "total_tokens": chat_response.usage.total_tokens,
                        "cache_read_tokens": chat_response.usage.cache_read_tokens,
                        "cache_creation_tokens": chat_response.usage.cache_creation_tokens,
                        "iterations": chat_response.usage.iterations
                    }

                # Keep raw response and metadata for trace building
                response = chat_response.raw_response
                tool_calls_metadata = chat_response.tool_calls_metadata

            else:
                # Local development - use existing ChatHandler
                response = st.session_state.chat_handler.send_message(
                    messages=[{"role": msg["role"], "content": msg["content"]}
                             for msg in st.session_state.messages],
                    mcp_servers=mcp_servers,
                    model=model,
                    max_tokens=max_tokens,
                    uploaded_files=st.session_state.uploaded_files
                )

                # Calculate response time
                duration_ms = (time.time() - start_time) * 1000

                # Format response (existing logic)
                response_text = st.session_state.chat_handler.format_response(response)
                usage = st.session_state.chat_handler.get_usage_info(response)

                # No metadata for local development path
                tool_calls_metadata = None

            # Calculate query duration
            query_duration = (datetime.utcnow() - query_start_time).total_seconds()

            # Update session statistics
            st.session_state.total_queries += 1
            if usage:
                st.session_state.total_tokens += usage.get("total_tokens", 0)
                # Cache-aware cost estimation
                cache_read = usage.get("cache_read_tokens", 0)
                cache_creation = usage.get("cache_creation_tokens", 0)
                raw_input = usage.get("input_tokens", 0)

                if st.session_state.llm_provider == "gemini":
                    # Gemini Flash pricing: $0.15/M input, $0.60/M output
                    # Cached reads at 0.1x rate
                    uncached_input = max(0, raw_input - cache_read)
                    input_cost = (uncached_input * 0.00015 / 1000
                                  + cache_read * 0.000015 / 1000)
                    output_cost = usage.get("output_tokens", 0) * 0.0006 / 1000
                else:
                    # Claude Sonnet pricing: $3/M input, $15/M output
                    # Cached reads at 0.1x, cache writes at 1.25x
                    uncached_input = max(0, raw_input - cache_read - cache_creation)
                    input_cost = (uncached_input * 0.003 / 1000
                                  + cache_read * 0.0003 / 1000
                                  + cache_creation * 0.00375 / 1000)
                    output_cost = usage.get("output_tokens", 0) * 0.015 / 1000
                estimated_cost = input_cost + output_cost
                st.session_state.total_cost += estimated_cost

                # Log the response (audit trail)
                st.session_state.audit_logger.log_mcp_response(
                    user_email=user["email"],
                    user_id=user["user_id"],
                    servers=st.session_state.selected_servers,
                    response_length=len(response_text),
                    input_tokens=usage.get("input_tokens", 0),
                    output_tokens=usage.get("output_tokens", 0),
                    total_tokens=usage.get("total_tokens", 0),
                    estimated_cost=estimated_cost,
                    duration_seconds=query_duration,
                    session_id=st.session_state.session_id
                )

            # Build orchestration trace
            message_index = len(st.session_state.messages)  # Index where this message will be stored

            # Get provider name for trace
            if st.session_state.provider_instance:
                provider_name = st.session_state.provider_instance.get_provider_name()
            else:
                provider_name = "Claude"  # Fallback

            trace = build_orchestration_trace(
                query=prompt,
                response=response,
                messages=[{"role": msg["role"], "content": msg["content"]}
                         for msg in st.session_state.messages],
                duration_ms=duration_ms,
                tokens=usage.get("total_tokens", 0) if usage else 0,
                cost_usd=estimated_cost if usage else 0,
                tool_calls_metadata=tool_calls_metadata,
                provider_name=provider_name
            )

            # Store trace in session state
            st.session_state.traces[message_index] = trace

            # Add to history
            st.session_state.messages.append({
                "role": "assistant",
                "content": response_text,
                "usage": usage
            })

            # Rerun to display the new messages via render_chat_history
            st.rerun()

        except Exception as e:
            error_msg = f"❌ Error: {str(e)}"

            # Log error (audit trail)
            st.session_state.audit_logger.log_error(
                user_email=user["email"],
                user_id=user["user_id"],
                error_type=type(e).__name__,
                error_message=str(e),
                servers=st.session_state.selected_servers,
                session_id=st.session_state.session_id
            )

            st.session_state.messages.append({
                "role": "assistant",
                "content": error_msg
            })

            # Rerun to display the error
            st.rerun()


# Multi-step prompts used for the "repeat" benchmark phase
REPEAT_PROMPTS = [
    "Complete PatientOne Workflow",
    "Predict Treatment Response",
    "Quantum + GEARS Validation",
]


def _start_benchmark_phase(
    selected_prompts: List[str],
    benchmark_providers: List[tuple],
    phase: str
):
    """Initialize a single benchmark phase and start incremental execution.

    Args:
        selected_prompts: List of prompt names to benchmark
        benchmark_providers: List of (provider_key, model_id) tuples
        phase: One of "cold", "warm", "repeat"
    """
    # For repeat phase, only use multi-step prompts that are in the selection
    if phase == "repeat":
        prompts_to_run = [p for p in REPEAT_PROMPTS if p in selected_prompts]
    else:
        prompts_to_run = selected_prompts

    queue = []
    for prompt_name in prompts_to_run:
        prompt_text = EXAMPLE_PROMPTS.get(prompt_name)
        if not prompt_text:
            continue
        for provider_key, model_id in benchmark_providers:
            queue.append((prompt_name, prompt_text, provider_key, model_id, phase))

    st.session_state["benchmark_queue"] = queue
    st.session_state["benchmark_done"] = []
    st.session_state["benchmark_running"] = True
    st.session_state["benchmark_phase"] = phase
    st.session_state["benchmark_providers_config"] = benchmark_providers
    st.rerun()


def _run_next_benchmark_step():
    """Run a single benchmark prompt then rerun. Called once per Streamlit cycle."""
    if not st.session_state.get("benchmark_running"):
        return

    queue = st.session_state.get("benchmark_queue", [])
    done = st.session_state.get("benchmark_done", [])

    if not queue:
        # All done — finalize results
        _finalize_benchmark(done)
        return

    # Pop the next item
    prompt_name, prompt_text, provider_key, model_id, run_type = queue.pop(0)
    st.session_state["benchmark_queue"] = queue

    # Run this single prompt with a timeout to prevent websocket disconnects
    from utils.benchmark import BenchmarkResult
    import time
    import concurrent.futures

    BENCHMARK_PROMPT_TIMEOUT = 180  # seconds per prompt (safe with per-prompt rerun cycles)

    try:
        # Reuse existing provider instance to avoid memory accumulation
        provider_instance = st.session_state.provider_instance
        if provider_instance is None or provider_instance.get_provider_name().lower() != provider_key:
            provider_instance = get_provider(provider_name=provider_key)
        mcp_servers = get_server_config(st.session_state.selected_servers)

        messages = [ChatMessage(role="user", content=prompt_text)]
        start = time.time()

        # Run in a thread with timeout so long-running prompts don't block
        executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
        future = executor.submit(
            provider_instance.send_message,
            messages=messages,
            mcp_servers=mcp_servers,
            model=model_id,
            max_tokens=4096,
            uploaded_files=st.session_state.get("uploaded_files"),
        )
        try:
            chat_response = future.result(timeout=BENCHMARK_PROMPT_TIMEOUT)
        except concurrent.futures.TimeoutError:
            executor.shutdown(wait=False, cancel_futures=True)
            duration_ms = (time.time() - start) * 1000
            result = BenchmarkResult(
                prompt_name=prompt_name, provider=provider_key, model=model_id,
                run_type=run_type, duration_ms=duration_ms,
                error=f"Timeout: exceeded {BENCHMARK_PROMPT_TIMEOUT}s",
            )
            done.append(result)
            st.session_state["benchmark_done"] = done
            try:
                audit_logger = get_audit_logger()
                user = st.session_state.get("user", {"email": "dev@localhost", "user_id": "dev"})
                audit_logger.log_benchmark_prompt(
                    user_email=user["email"],
                    user_id=user["user_id"],
                    session_id=st.session_state.get("session_id", "unknown"),
                    result=result.to_dict()
                )
            except Exception:
                pass
            st.rerun()
            return
        finally:
            executor.shutdown(wait=False)

        duration_ms = (time.time() - start) * 1000
        usage = chat_response.usage

        if usage:
            # Cache-aware cost
            cache_read = usage.cache_read_tokens
            cache_creation = usage.cache_creation_tokens
            raw_input = usage.input_tokens
            if provider_key == "gemini":
                uncached = max(0, raw_input - cache_read)
                cost = (uncached * 0.00015 / 1000 + cache_read * 0.000015 / 1000
                        + usage.output_tokens * 0.0006 / 1000)
            else:
                uncached = max(0, raw_input - cache_read - cache_creation)
                cost = (uncached * 0.003 / 1000 + cache_read * 0.0003 / 1000
                        + cache_creation * 0.00375 / 1000
                        + usage.output_tokens * 0.015 / 1000)

            result = BenchmarkResult(
                prompt_name=prompt_name, provider=provider_key, model=model_id,
                run_type=run_type, input_tokens=usage.input_tokens,
                output_tokens=usage.output_tokens, total_tokens=usage.total_tokens,
                cache_read_tokens=cache_read, cache_creation_tokens=cache_creation,
                iterations=usage.iterations, duration_ms=duration_ms,
                estimated_cost=cost,
            )
        else:
            result = BenchmarkResult(
                prompt_name=prompt_name, provider=provider_key, model=model_id,
                run_type=run_type, duration_ms=duration_ms,
                error="No usage info returned",
            )

    except Exception as e:
        result = BenchmarkResult(
            prompt_name=prompt_name, provider=provider_key, model=model_id,
            run_type=run_type, error=str(e)[:200],
        )

    done.append(result)
    st.session_state["benchmark_done"] = done

    # Log this individual prompt result to audit log
    try:
        audit_logger = get_audit_logger()
        user = st.session_state.get("user", {"email": "dev@localhost", "user_id": "dev"})
        audit_logger.log_benchmark_prompt(
            user_email=user["email"],
            user_id=user["user_id"],
            session_id=st.session_state.get("session_id", "unknown"),
            result=result.to_dict()
        )
    except Exception:
        pass  # Don't fail benchmark on audit log errors

    # Rerun to process the next item (or finalize)
    st.rerun()


def _finalize_benchmark(results):
    """Finalize a benchmark phase: accumulate results, rebuild DataFrame, log.

    Args:
        results: List of BenchmarkResult objects from the current phase
    """
    from utils.benchmark import TokenBenchmark
    from utils.audit_logger import get_audit_logger
    import pandas as pd

    phase = st.session_state.get("benchmark_phase", "cold")
    st.session_state["benchmark_running"] = False
    st.session_state["benchmark_queue"] = []

    if not results:
        return

    # Mark this phase as done
    phases_done = st.session_state.get("benchmark_phases_done", set())
    phases_done.add(phase)
    st.session_state["benchmark_phases_done"] = phases_done

    # Accumulate results across phases
    all_results = st.session_state.get("benchmark_all_results", [])
    all_results.extend(results)
    st.session_state["benchmark_all_results"] = all_results

    # Rebuild DataFrame from all accumulated results
    rows = [r.to_dict() for r in all_results]
    df = pd.DataFrame(rows)

    # Compute real input tokens and cache hit rate
    # Claude reports input_tokens as uncached only; real total = input + cache_read + cache_creation
    # Gemini reports input_tokens as the full total (includes cached portion)
    if "cache_read_tokens" in df.columns and "input_tokens" in df.columns:
        cache_creation = df.get("cache_creation_tokens", 0)
        # Real input: for Claude = input + cache_read + cache_creation; for Gemini = input (already total)
        df["real_input_tokens"] = df["input_tokens"] + df["cache_read_tokens"] + cache_creation
        gemini_mask = df["provider"] == "gemini"
        df.loc[gemini_mask, "real_input_tokens"] = df.loc[gemini_mask, "input_tokens"]
        # Cache hit rate = cache_read / real_input
        df["cache_hit_rate"] = (
            df["cache_read_tokens"] / df["real_input_tokens"].replace(0, 1)
        ).round(3)

    st.session_state["benchmark_results"] = df
    st.session_state["benchmark_csv"] = TokenBenchmark.results_to_csv_string(all_results)

    # Log this phase to audit
    try:
        phase_rows = [r.to_dict() for r in results]
        audit_logger = get_audit_logger()
        user = st.session_state.get("user", {"email": "dev@localhost", "user_id": "dev"})
        audit_logger.log_benchmark_run(
            user_email=user["email"],
            user_id=user["user_id"],
            session_id=st.session_state.get("session_id", "unknown"),
            prompt_count=len(set(r.prompt_name for r in results)),
            provider_count=len(set(r.provider for r in results)),
            results=phase_rows
        )
    except Exception:
        pass  # Don't fail benchmark on audit log errors


def render_benchmark_results():
    """Render benchmark results if available."""
    # Show in-progress status
    if st.session_state.get("benchmark_running"):
        phase = st.session_state.get("benchmark_phase", "")
        done = st.session_state.get("benchmark_done", [])
        queue = st.session_state.get("benchmark_queue", [])
        total = len(done) + len(queue)
        st.markdown("---")
        st.subheader(f"Benchmark: {phase.title()} Phase")
        st.progress(len(done) / max(total, 1),
                     text=f"Completed {len(done)} of {total} runs...")
        if done:
            last = done[-1]
            status = f"Last: {last.prompt_name} ({last.provider})"
            if last.error:
                status += f" - Error: {last.error[:80]}"
            else:
                status += f" - {last.total_tokens:,} tokens, {last.duration_ms:.0f}ms"
            st.caption(status)
        return

    if "benchmark_results" not in st.session_state:
        return

    df = st.session_state["benchmark_results"]
    csv_data = st.session_state.get("benchmark_csv", "")
    phases_done = st.session_state.get("benchmark_phases_done", set())

    st.markdown("---")
    phases_label = ", ".join(sorted(phases_done)).title()
    st.subheader(f"Benchmark Results ({phases_label})")

    # Summary metrics per provider
    for provider_key in df["provider"].unique():
        prov_df = df[df["provider"] == provider_key]
        col1, col2, col3, col4 = st.columns(4)
        # Use real_input_tokens (accounts for Claude's split reporting)
        input_col = "real_input_tokens" if "real_input_tokens" in prov_df.columns else "input_tokens"
        col1.metric(f"{provider_key} Total Input", f"{prov_df[input_col].sum():,}")
        col2.metric(f"Total Output", f"{prov_df['output_tokens'].sum():,}")
        col3.metric(f"Total Cost", f"${prov_df['estimated_cost'].sum():.4f}")
        # Only average cache hit over runs that had cache reads
        if "cache_hit_rate" in prov_df.columns:
            cached_runs = prov_df[prov_df["cache_read_tokens"] > 0]
            if not cached_runs.empty:
                col4.metric(f"Avg Cache Hit", f"{cached_runs['cache_hit_rate'].mean():.1%}")
            else:
                col4.metric(f"Avg Cache Hit", "0%")
        else:
            col4.metric(f"Avg Cache Hit", "N/A")

    # Data table
    st.dataframe(df, use_container_width=True)

    # Bar chart: token usage by prompt, grouped by provider + run type
    try:
        import altair as alt
        chart_df = df[df["error"].isna() | (df["error"] == "None")].copy()
        if not chart_df.empty:
            chart_df["label"] = chart_df["provider"] + " (" + chart_df["run_type"] + ")"
            chart = alt.Chart(chart_df).mark_bar().encode(
                x=alt.X("prompt_name:N", title="Prompt", sort=None),
                y=alt.Y("total_tokens:Q", title="Total Tokens"),
                color="label:N",
                xOffset="label:N"
            ).properties(width=700, height=400, title="Token Usage by Prompt")
            st.altair_chart(chart, use_container_width=True)
    except ImportError:
        pass  # altair not available

    # CSV download
    if csv_data:
        st.download_button(
            "Export CSV",
            data=csv_data,
            file_name="benchmark_results.csv",
            mime="text/csv",
            key="download_benchmark"
        )


def main():
    """Main application."""
    # Require authentication (SSO via OAuth2 Proxy)
    user = require_authentication()

    # Initialize session state
    initialize_session_state()

    # Log session start (first time only)
    if "session_logged" not in st.session_state:
        st.session_state.audit_logger.log_session_start(
            user_email=user["email"],
            user_id=user["user_id"],
            session_id=st.session_state.session_id
        )
        st.session_state.session_logged = True

    # Render sidebar and get settings
    model, max_tokens, show_trace, trace_style = render_sidebar()

    # Main content area
    st.title("🧬 Precision Medicine MCP Chat")
    st.markdown("Chat interface for testing deployed MCP servers on GCP Cloud Run")

    # Check that at least one provider is available
    if not st.session_state.provider_instance and not st.session_state.chat_handler:
        st.error("❌ No LLM provider configured")
        st.info("""
        **Setup Instructions:**
        Set at least one API key:
        - `ANTHROPIC_API_KEY=your_key` for Claude
        - `GEMINI_API_KEY=your_key` for Gemini

        Or create a `.env` file with your key(s).
        """)
        return

    # Show server status
    render_server_status()

    st.markdown("---")

    # Render chat history
    render_chat_history(show_trace=show_trace, trace_style=trace_style)

    # Run next benchmark step if benchmark is in progress
    _run_next_benchmark_step()

    # Render benchmark results if available
    render_benchmark_results()

    # Chat input — handle pending example prompt or manual input
    if "pending_example" in st.session_state:
        pending = st.session_state.pending_example
        del st.session_state.pending_example
        handle_user_input(pending, model, max_tokens)
    elif prompt := st.chat_input("Ask me anything about precision medicine...", key="chat_input"):
        handle_user_input(prompt, model, max_tokens)


if __name__ == "__main__":
    main()
