import streamlit as st
import time
from components.config_panel import render_config_panel
from components.pareto_chart import render_pareto_chart
from components.progress_chart import render_progress_chart
from components.model_cards import render_model_cards
from components.business_metrics import render_hero_section
from components.search_trajectory import render_search_trajectory
from components.objective_progress import render_objective_progress
from components.success_stories import render_success_stories

# Page configuration
st.set_page_config(
    page_title="PrimaLabs",
    page_icon="🔥",
    layout="wide",
    initial_sidebar_state="collapsed"
)

# Load custom CSS
st.markdown("""
<style>
    /* Hide Streamlit branding for cleaner look */
    #MainMenu {visibility: hidden;}
    footer {visibility: hidden;}

    /* Premium styling */
    .stTabs [data-baseweb="tab-list"] {
        gap: 8px;
    }

    .stTabs [data-baseweb="tab"] {
        height: 50px;
        padding: 10px 24px;
        background: rgba(255, 255, 255, 0.05);
        border-radius: 10px;
        color: rgba(255,255,255,0.8);
        font-weight: 500;
    }

    .stTabs [aria-selected="true"] {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
    }

    /* Metric containers */
    [data-testid="stMetricValue"] {
        font-size: 2rem;
        font-weight: 700;
    }

    /* Info boxes */
    .stAlert {
        background: rgba(99, 102, 241, 0.1);
        border: 1px solid rgba(99, 102, 241, 0.3);
        border-radius: 10px;
    }
</style>
""", unsafe_allow_html=True)

# Initialize session state
if 'optimization_running' not in st.session_state:
    st.session_state.optimization_running = False

if 'optimization_completed' not in st.session_state:
    st.session_state.optimization_completed = False

if 'active_tab' not in st.session_state:
    st.session_state.active_tab = 0  # Default to Optimization tab

# Animation mode - set to False for Hugging Face Spaces (instant results)
# Set to True for local demos with real-time progress animation
ENABLE_ANIMATION = True  # Set to False for HF Spaces deployment (instant results)

# Hero Section
render_hero_section()

# Render config panel and get config
config = render_config_panel()

# Store config in session state for collapsed panel access
st.session_state.last_config = config

# Keep reference as sidebar_config for backward compatibility
sidebar_config = config

# Handle start optimization button
if sidebar_config['start_button'] and not st.session_state.optimization_running:
    st.session_state.optimization_running = True
    st.session_state.optimization_completed = False
    st.session_state.active_tab = 0  # Switch to Optimization tab (index 0)
    st.session_state.just_started = True
    st.rerun()

# Handle reset button
if sidebar_config['reset_button']:
    st.session_state.optimization_running = False
    st.session_state.optimization_completed = False
    # Clear progress state
    if 'progress_step' in st.session_state:
        del st.session_state.progress_step
    if 'last_update' in st.session_state:
        del st.session_state.last_update
    st.session_state.active_tab = 0  # Return to Optimization tab
    st.rerun()

# Main Tabs
tabs = st.tabs(["Optimization", "Results", "Success Stories"])

# Optimization Tab
with tabs[0]:
    if not st.session_state.optimization_running and not st.session_state.optimization_completed:
        # Show selected objectives
        selected_objectives = [k.replace('_', ' ').title() for k, v in sidebar_config['objectives'].items() if v]
        objectives_str = ", ".join(selected_objectives)

        # Build stopping criteria message
        if sidebar_config['limiting_factor'] == 'models':
            stop_msg = f"Max {sidebar_config['num_models']:,} models (will complete in {sidebar_config['estimated_time']})"
        else:
            stop_msg = f"Max {sidebar_config['max_hours']}h wall time (will explore ~{sidebar_config['actual_models']:,} models)"

        # Get enabled techniques
        enabled_techniques = [k.replace('_', ' ').title() for k, v in sidebar_config.get('techniques', {}).items() if v]
        techniques_str = ", ".join(enabled_techniques) if enabled_techniques else "Standard optimization"

        st.info(f"**Ready to optimize {sidebar_config['selected_model']} for {sidebar_config['target_hardware']}**\n\n**Hardware:** {sidebar_config['num_units']} × {sidebar_config['target_hardware']} ({sidebar_config['hardware_specs']['vram']} GB VRAM)\n\n**Stop when:** {stop_msg}\n\n**Objectives:** {objectives_str}\n\n**Techniques:** {techniques_str}\n\nClick 'Start Optimization' in the sidebar to begin")

        # Add preview of what will happen
        st.markdown("### 🚀 What PrimaLabs Will Do")
        col1, col2, col3 = st.columns(3)
        with col1:
            st.markdown("**🔍 Explore**")
            st.caption(f"Apply {len(enabled_techniques)} optimization techniques across ~{sidebar_config['actual_models']:,} model variants, testing quantization (INT8/INT4), pruning, LoRA, and other compression methods")
        with col2:
            st.markdown("**📊 Optimize**")
            st.caption(f"Find Pareto-optimal configurations for {sidebar_config['target_hardware']} balancing {len([v for v in sidebar_config['objectives'].values() if v])} objectives: {objectives_str[:60]}...")
        with col3:
            st.markdown("**🎯 Deliver**")
            st.caption(f"Export production-ready LLMs optimized for {sidebar_config['hardware_specs']['category']} deployment with TensorRT/ONNX/PyTorch formats")
    else:
        # Handle progress step and animation logic
        if st.session_state.optimization_running:
            # Initialize progress_step if needed
            if 'progress_step' not in st.session_state:
                st.session_state.progress_step = 1  # Start at 1, not 0
                st.session_state.start_time = time.time()

            # Show initializing message only at start
            if st.session_state.progress_step == 1 and time.time() - st.session_state.start_time < 0.1:
                st.info("⚙️ **Initializing optimization...** Setting up GPU cluster and preparing model configurations")

            # Check if optimization should complete
            if st.session_state.progress_step >= 20:
                st.session_state.optimization_running = False
                st.session_state.optimization_completed = True
                # Calculate final savings for success message
                from mock_data.sample_data import estimate_inference_cost
                baseline_cost = estimate_inference_cost(32.5, "FP32")["cost_per_1M"]
                best_cost = estimate_inference_cost(6.2, "INT4")["cost_per_1M"]
                cost_reduction = ((baseline_cost - best_cost) / baseline_cost) * 100
                size_reduction = ((32.5 - 6.2) / 32.5) * 100

                st.success(f"🎉 **Optimization Complete!** Discovered 3 Pareto-optimal models • **{cost_reduction:.0f}% cost reduction** • **{size_reduction:.0f}% size reduction** • Switch to Results tab to explore →")
                st.rerun()

        # Show progress bar
        if st.session_state.optimization_running:
            progress_pct = min(st.session_state.get('progress_step', 0) / 20, 1.0)
            st.progress(progress_pct, text=f"Optimizing... {int(progress_pct * 100)}%")

        # Real-time metrics (show during and after optimization)
        if st.session_state.optimization_running or st.session_state.optimization_completed:
            from mock_data.sample_data import (estimate_inference_cost, estimate_latency,
                                              estimate_memory_footprint, estimate_energy_consumption)

            status_label = "Live Metrics" if st.session_state.optimization_running else "Final Metrics"
            st.markdown(f"#### {status_label}")

            # Calculate dynamic metrics based on progress (0-20 steps)
            progress_step = st.session_state.get('progress_step', 20) if st.session_state.optimization_completed else st.session_state.get('progress_step', 0)

            # Map progress to realistic metrics (show actual models being explored based on limiting factor)
            models_explored = min(int((progress_step / 20) * sidebar_config['actual_models']), sidebar_config['actual_models'])
            best_accuracy = min(72 + (progress_step / 20) * 13.2, 85.2)
            current_size = max(32.5 - (progress_step / 20) * 26.3, 6.2)

            # Calculate time elapsed (scaled to actual estimated time)
            time_elapsed_hours = (progress_step / 20) * sidebar_config['actual_time']
            # Format time display
            if time_elapsed_hours < 1:
                time_display = f"{int(time_elapsed_hours * 60)} min"
            elif time_elapsed_hours < 24:
                h = int(time_elapsed_hours)
                m = int((time_elapsed_hours - h) * 60)
                time_display = f"{h}h {m}m" if m > 0 else f"{h}h"
            else:
                d = int(time_elapsed_hours / 24)
                rh = int(time_elapsed_hours % 24)
                time_display = f"{d}d {rh}h"

            # Calculate all metrics
            baseline_metrics = estimate_inference_cost(32.5, "FP32")
            current_metrics = estimate_inference_cost(current_size, "INT4")

            current_latency = estimate_latency(current_size, "INT4")
            baseline_latency = estimate_latency(32.5, "FP32")

            current_memory = estimate_memory_footprint(current_size, "INT4")
            baseline_memory = estimate_memory_footprint(32.5, "FP32")

            current_energy = estimate_energy_consumption(current_size, "INT4")
            baseline_energy = estimate_energy_consumption(32.5, "FP32")

            # Build metrics list based on selected objectives
            metrics_to_show = []

            # Always show models explored (with progress)
            if st.session_state.optimization_running:
                models_display = f"{models_explored:,} / {sidebar_config['actual_models']:,}"
            else:
                models_display = f"{sidebar_config['actual_models']:,}"
            metrics_to_show.append(("Models Explored", models_display, None, "normal"))

            # Accuracy (always selected)
            if sidebar_config['objectives']['accuracy']:
                delta_acc = f"+{best_accuracy - 72:.1f}%" if progress_step > 0 else None
                metrics_to_show.append(("Best Accuracy", f"{best_accuracy:.1f}%", delta_acc, "normal"))

            # Model Size
            if sidebar_config['objectives']['size']:
                delta_size = f"-{32.5 - current_size:.1f} GB" if progress_step > 0 else None
                metrics_to_show.append(("Best Size", f"{current_size:.1f} GB", delta_size, "inverse"))

            # Inference Cost
            if sidebar_config['objectives']['cost']:
                cost_savings_pct = ((baseline_metrics['cost_per_1M'] - current_metrics['cost_per_1M']) / baseline_metrics['cost_per_1M']) * 100
                delta_cost = f"-{cost_savings_pct:.0f}%" if progress_step > 0 else None
                metrics_to_show.append(("Inference Cost", f"${current_metrics['cost_per_1M']:.2f}/1M", delta_cost, "inverse"))

            # Throughput
            if sidebar_config['objectives']['throughput']:
                delta_tput = f"+{int(current_metrics['throughput_qps'] - baseline_metrics['throughput_qps'])} QPS" if progress_step > 0 else None
                metrics_to_show.append(("Throughput", f"{int(current_metrics['throughput_qps'])} QPS", delta_tput, "normal"))

            # Latency
            if sidebar_config['objectives']['latency']:
                latency_reduction = baseline_latency - current_latency
                delta_lat = f"-{latency_reduction:.1f} ms" if progress_step > 0 else None
                metrics_to_show.append(("Latency", f"{current_latency:.1f} ms", delta_lat, "inverse"))

            # Memory Footprint
            if sidebar_config['objectives']['memory']:
                memory_reduction = baseline_memory - current_memory
                delta_mem = f"-{memory_reduction:.1f} GB" if progress_step > 0 else None
                metrics_to_show.append(("Memory", f"{current_memory:.1f} GB", delta_mem, "inverse"))

            # Energy Efficiency
            if sidebar_config['objectives']['energy']:
                energy_reduction = baseline_energy - current_energy
                delta_energy = f"-{energy_reduction:.0f} W" if progress_step > 0 else None
                metrics_to_show.append(("Energy", f"{current_energy:.0f} W", delta_energy, "inverse"))

            # Always show time elapsed (simulated)
            metrics_to_show.append(("Time Elapsed", time_display, None, "normal"))

            # Display metrics dynamically (4 per row)
            num_metrics = len(metrics_to_show)
            rows = (num_metrics + 3) // 4  # Ceiling division

            for row in range(rows):
                cols = st.columns(4)
                for i in range(4):
                    idx = row * 4 + i
                    if idx < num_metrics:
                        label, value, delta, delta_color = metrics_to_show[idx]
                        with cols[i]:
                            st.metric(label, value, delta=delta, delta_color=delta_color)

            st.markdown("---")

        # Run optimization animation loop using st.empty() for real-time updates
        if st.session_state.optimization_running:
            # Create empty containers for progressive updates
            status_container = st.empty()
            objective_container = st.empty()
            divider1_container = st.empty()
            progress_container = st.empty()
            divider2_container = st.empty()
            trajectory_container = st.empty()

            # Animate through all 20 steps
            for step in range(st.session_state.progress_step, 21):
                # Update status
                with status_container:
                    st.info(f"⚙️ **Optimizing...** Step {step}/20 ({int(step/20*100)}%)")

                # Update objective progress charts
                with objective_container.container():
                    render_objective_progress(
                        budget_hours=sidebar_config['estimated_hours'],
                        is_running=True,
                        selected_objectives=sidebar_config['objectives'],
                        current_step=step
                    )

                with divider1_container:
                    st.markdown("---")

                # Update progress chart
                with progress_container.container():
                    render_progress_chart(
                        budget_hours=sidebar_config['estimated_hours'],
                        is_running=True,
                        enable_animation=ENABLE_ANIMATION,
                        current_step=step
                    )

                with divider2_container:
                    st.markdown("---")

                # Update search trajectory
                with trajectory_container.container():
                    render_search_trajectory(
                        is_running=True,
                        current_step=step
                    )

                # Update session state
                st.session_state.progress_step = step

                # Sleep to create animation effect (allows browser to render)
                if ENABLE_ANIMATION and step < 20:
                    time.sleep(0.5)

            # Mark as complete
            st.session_state.progress_step = 20
            st.session_state.optimization_running = False
            st.session_state.optimization_completed = True

            # Clear status and show completion
            status_container.success("🎉 **Optimization Complete!**")
            time.sleep(1)
            st.rerun()

        elif st.session_state.optimization_completed:
            # Show final results
            render_objective_progress(
                budget_hours=sidebar_config['estimated_hours'],
                is_running=False,
                selected_objectives=sidebar_config['objectives']
            )

            st.markdown("---")

            render_progress_chart(
                budget_hours=sidebar_config['estimated_hours'],
                is_running=False,
                enable_animation=ENABLE_ANIMATION
            )

            st.markdown("---")

            render_search_trajectory(
                is_running=False
            )

    # Footer for Optimization tab
    st.markdown("---")
    if sidebar_config['limiting_factor'] == 'models':
        footer_msg = f"PrimaLabs • {sidebar_config['selected_model']} • {sidebar_config['target_hardware']} • {sidebar_config['num_units']} units • {sidebar_config['models_per_hour']} models/hour"
    else:
        footer_msg = f"PrimaLabs • {sidebar_config['selected_model']} • {sidebar_config['target_hardware']} • {sidebar_config['num_units']} units • ~{sidebar_config['actual_models']:,} models in {sidebar_config['max_hours']}h"
    st.caption(footer_msg)

# Results Tab
with tabs[1]:
    if st.session_state.optimization_running or st.session_state.optimization_completed:
        # Add key highlights at the top
        if st.session_state.optimization_completed:
            st.markdown("### 🎯 Optimization Results")

            # Key metrics in highlight boxes
            from mock_data.sample_data import estimate_inference_cost
            baseline_cost = estimate_inference_cost(32.5, "FP32")["cost_per_1M"]
            best_cost = estimate_inference_cost(6.2, "INT4")["cost_per_1M"]
            cost_reduction = ((baseline_cost - best_cost) / baseline_cost) * 100
            size_reduction = ((32.5 - 6.2) / 32.5) * 100

            col1, col2, col3, col4 = st.columns(4)
            with col1:
                st.markdown(f"""
                <div style='background: linear-gradient(135deg, #10b981 0%, #059669 100%); padding: 1.5rem; border-radius: 12px; text-align: center;'>
                    <div style='font-size: 2rem; font-weight: bold; color: white;'>{sidebar_config['actual_models']:,}</div>
                    <div style='color: rgba(255,255,255,0.9); font-size: 0.9rem;'>Models Explored</div>
                </div>
                """, unsafe_allow_html=True)
            with col2:
                st.markdown(f"""
                <div style='background: linear-gradient(135deg, #6366f1 0%, #4f46e5 100%); padding: 1.5rem; border-radius: 12px; text-align: center;'>
                    <div style='font-size: 2rem; font-weight: bold; color: white;'>3</div>
                    <div style='color: rgba(255,255,255,0.9); font-size: 0.9rem;'>Pareto-Optimal Models</div>
                </div>
                """, unsafe_allow_html=True)
            with col3:
                st.markdown(f"""
                <div style='background: linear-gradient(135deg, #f59e0b 0%, #d97706 100%); padding: 1.5rem; border-radius: 12px; text-align: center;'>
                    <div style='font-size: 2rem; font-weight: bold; color: white;'>{cost_reduction:.0f}%</div>
                    <div style='color: rgba(255,255,255,0.9); font-size: 0.9rem;'>Cost Reduction</div>
                </div>
                """, unsafe_allow_html=True)
            with col4:
                st.markdown(f"""
                <div style='background: linear-gradient(135deg, #ec4899 0%, #db2777 100%); padding: 1.5rem; border-radius: 12px; text-align: center;'>
                    <div style='font-size: 2rem; font-weight: bold; color: white;'>{size_reduction:.0f}%</div>
                    <div style='color: rgba(255,255,255,0.9); font-size: 0.9rem;'>Size Reduction</div>
                </div>
                """, unsafe_allow_html=True)

            st.markdown("---")

        # Show optimized models
        st.markdown("### 🏆 Discovered Models")
        st.caption("Six pareto-optimal models, each optimized for different trade-offs")
        render_model_cards()

        if st.session_state.optimization_completed:
            # Add comparison table
            st.markdown("---")
            st.markdown("### 📊 Model Comparison")

            from mock_data.sample_data import generate_discovered_models, estimate_inference_cost, estimate_latency, estimate_memory_footprint, estimate_energy_consumption

            models = generate_discovered_models()

            # Build comparison data
            import pandas as pd
            comparison_data = []

            for model in models:
                cost_info = estimate_inference_cost(model['size_gb'], model['quantization'])
                baseline_cost = estimate_inference_cost(32.5, "FP32")["cost_per_1M"]
                cost_savings = ((baseline_cost - cost_info['cost_per_1M']) / baseline_cost) * 100

                comparison_data.append({
                    'Model': model['name'],
                    'Accuracy': f"{model['accuracy']:.1f}%",
                    'Size': f"{model['size_gb']:.1f} GB",
                    'Cost/1M': f"${cost_info['cost_per_1M']:.2f}",
                    'Savings': f"{cost_savings:.0f}%",
                    'Throughput': f"{int(cost_info['throughput_qps'])} QPS",
                    'Latency': f"{estimate_latency(model['size_gb'], model['quantization']):.1f} ms",
                    'GPU': cost_info['gpu_tier']
                })

            df_comparison = pd.DataFrame(comparison_data)
            st.dataframe(df_comparison, use_container_width=True, hide_index=True)

            st.markdown("---")
            st.markdown("### Performance Summary")

            # Calculate cost metrics for performance summary
            baseline_cost = estimate_inference_cost(32.5, "FP32")["cost_per_1M"]
            best_cost = estimate_inference_cost(4.8, "INT8")["cost_per_1M"]
            cost_reduction = ((baseline_cost - best_cost) / baseline_cost) * 100

            col1, col2, col3, col4 = st.columns(4)

            with col1:
                st.metric("Best Accuracy", "85.2%", delta="Optimized-7B-Q8")

            with col2:
                st.metric("Smallest Model", "4.8 GB", delta="-85% vs baseline")

            with col3:
                st.metric("Lowest Cost", f"${best_cost:.2f}/1M", delta=f"-{cost_reduction:.0f}%", delta_color="inverse")

            with col4:
                annual_savings = (baseline_cost - best_cost) * 1000  # Assuming 1B requests/year
                st.metric("Est. Annual Savings", f"${annual_savings:,.0f}", delta="vs baseline")

            # Add Pareto Front Analysis
            st.markdown("---")
            render_pareto_chart(selected_objectives=sidebar_config['objectives'])
    else:
        st.info("Start optimization to discover optimized model variants")

    # Footer for Results tab
    st.markdown("---")
    if sidebar_config['limiting_factor'] == 'models':
        footer_msg = f"PrimaLabs • {sidebar_config['selected_model']} • {sidebar_config['target_hardware']} • {sidebar_config['num_units']} units • {sidebar_config['models_per_hour']} models/hour"
    else:
        footer_msg = f"PrimaLabs • {sidebar_config['selected_model']} • {sidebar_config['target_hardware']} • {sidebar_config['num_units']} units • ~{sidebar_config['actual_models']:,} models in {sidebar_config['max_hours']}h"
    st.caption(footer_msg)

# Success Stories Tab
with tabs[2]:
    render_success_stories()