"""Streamlit UI components for the hill climber dashboard.

This module contains all Streamlit-specific UI rendering logic,
keeping UI concerns separate from data and plotting.
"""

import time
from datetime import datetime
from pathlib import Path
from typing import Optional, List, Dict, Any, Set, Tuple
import pandas as pd

try:
    import streamlit as st

except ImportError:
    st = None


def apply_custom_css() -> None:
    """Apply custom CSS styling to the dashboard.
    
    Adjusts padding, font sizes, and text wrapping for optimal dashboard appearance.
    """

    if st is None:
        return
        
    st.markdown("""
        <style>
        .main { padding-top: 1.5rem !important; }
        .main .block-container { padding-top: 1.5rem !important; }
        .main h2 { font-size: 1.5rem !important; margin-top: 0.5rem !important; }
        .main h3 { font-size: 1.1rem !important; }
        .main h2:first-of-type { margin-top: 0 !important; padding-top: 0 !important; }
        
        /* Prevent main content from going under header */
        [data-testid="stAppViewContainer"] > section:first-child {
            padding-top: 1.5rem !important;
        }
        .stMainBlockContainer {
            padding-top: 1.5rem !important;
        }
        
        /* Widen sidebar to fit logo and give text more space */
        [data-testid="stSidebar"] {
            width: 275px !important;
            min-width: 275px !important;
        }
        [data-testid="stSidebar"] > div:first-child {
            width: 275px !important;
        }
        
        /* Prevent text wrapping in sidebar - use ellipsis instead */
        [data-testid="stSidebar"] p {
            white-space: nowrap !important;
            overflow: hidden !important;
            text-overflow: ellipsis !important;
        }
        
        /* Consistent horizontal rule spacing in sidebar */
        [data-testid="stSidebar"] hr {
            margin-top: 0.5rem !important;
            margin-bottom: 1rem !important;
        }
        
        /* Consistent horizontal rule spacing in main content */
        .main hr {
            margin-top: 0.5rem !important;
            margin-bottom: 1rem !important;
        }
        </style>
    """, unsafe_allow_html=True)


def render_sidebar_title() -> None:
    """Render the sidebar logo.
    
    Displays the Hill climber logo in the sidebar.
    """

    if st is None:
        return
    
    import os
    logo_path = os.path.join(os.path.dirname(__file__), 'assets', 'logo.svg')
    
    if os.path.exists(logo_path):
        with open(logo_path, 'r') as f:
            svg_content = f.read()
        st.sidebar.markdown(
            f"<div style='margin-top: -3rem; padding-top: 0.25rem;'>{svg_content}</div>",
            unsafe_allow_html=True
        )
    else:
        # Fallback to text if logo not found
        st.sidebar.markdown(
            "<h1 style='margin-top: -3rem; padding-top: 0.25rem; font-size: 2.8rem; line-height: 1.2; color: #ff4b4b;'>"
            "Hill<br>climber</h1>",
            unsafe_allow_html=True
        )
    
    st.sidebar.markdown("---")


def render_database_selector(session_state: Any, db_files: List[Path], project_root: Path) -> Optional[str]:
    """Render database selection UI in sidebar.
    
    Args:
        session_state (Any): Streamlit session state object.
        db_files (List[Path]): List of available database files in project.
        project_root (Path): Project root directory for relative path display.
        
    Returns:
        str: Selected database path, or None if no database selected.
    """

    if st is None:
        return None
        
    st.sidebar.header("Run data")
    
    db_path = session_state.db_path
    
    if not db_files:
        st.sidebar.info("No .db files found in project")
        return db_path
    
    # Create display labels with relative paths
    db_labels = []
    for db_file in db_files:
        try:
            rel_path = db_file.relative_to(project_root)
            db_labels.append(str(rel_path))
        except ValueError:
            # File is outside project root
            db_labels.append(str(db_file))
    
    # Find current selection index
    current_idx = 0
    if session_state.db_user_selected and db_path:
        try:
            current_idx = next(i for i, f in enumerate(db_files) if str(f) == db_path)
        except StopIteration:
            pass
    
    # Database file selector
    selected_idx = st.sidebar.selectbox(
        "Select database",
        options=list(range(len(db_labels))),
        format_func=lambda i: db_labels[i],
        index=current_idx,
        label_visibility="collapsed"
    )
    
    selected_file = db_files[selected_idx]
    
    # Update selection if changed
    if str(selected_file) != session_state.db_path:
        session_state.db_path = str(selected_file)
        session_state.db_user_selected = True
        st.rerun()
    
    # Warn if path doesn't exist
    if session_state.db_user_selected and not Path(db_path).exists():
        st.sidebar.warning(f"⚠ Not found: `{db_path}`")
    
    return db_path


def render_auto_refresh_controls() -> Tuple[bool, float]:
    """Render auto-refresh controls in sidebar.
    
    Returns:
        Tuple[bool, float]: Tuple of (auto_refresh_enabled, refresh_interval_seconds).
    """

    if st is None:
        return False, 60.0
        
    auto_refresh = st.sidebar.checkbox("Auto-refresh", key="auto_refresh")

    refresh_interval_minutes = st.sidebar.slider(
        "Refresh interval (minutes)",
        min_value=0.5, max_value=5.0, step=0.5,
        key="refresh_interval"
    )

    if st.sidebar.button("Refresh now", key="refresh_now"):

        # Increment refresh key to force clean plot re-rendering
        st.session_state.plot_refresh_key = st.session_state.get('plot_refresh_key', 0) + 1

        # Save current plot options before refresh
        st.session_state.saved_history_type = st.session_state.get('history_type', 'Improvements (best)')
        st.session_state.saved_additional_base_metrics = st.session_state.get('additional_base_metrics', [])
        st.session_state.saved_normalize_metrics = st.session_state.get('normalize_metrics', False)
        st.session_state.saved_show_exchanges = st.session_state.get('show_exchanges', False)
        st.session_state.saved_max_points = st.session_state.get('max_points', 1000)
        st.session_state.saved_plot_columns = st.session_state.get('plot_columns', 'Two columns')
        st.rerun()
    
    return auto_refresh, refresh_interval_minutes * 60


def render_plot_options(available_metrics: List[str]) -> Dict[str, Any]:
    """Render plot configuration options in sidebar.
    
    Args:
        available_metrics (List[str]): List of available metric names from database.
        
    Returns:
        Dict[str, Any]: Dictionary with plot configuration options including history_type,
            objective_metric, additional_metrics, normalize_metrics, show_exchanges,
            max_points, and n_cols.
    """

    if st is None:
        return {}
        
    st.sidebar.markdown("<hr style='margin-top: 0.5rem; margin-bottom: 1rem;'>", unsafe_allow_html=True)
    st.sidebar.subheader("Plot options")
    
    # Restore saved values if they exist (from manual refresh)
    if 'saved_history_type' in st.session_state and 'history_type' not in st.session_state:
        st.session_state.history_type = st.session_state.saved_history_type
    if 'saved_additional_base_metrics' in st.session_state and 'additional_base_metrics' not in st.session_state:
        st.session_state.additional_base_metrics = st.session_state.saved_additional_base_metrics
    if 'saved_normalize_metrics' in st.session_state and 'normalize_metrics' not in st.session_state:
        st.session_state.normalize_metrics = st.session_state.saved_normalize_metrics
    if 'saved_show_exchanges' in st.session_state and 'show_exchanges' not in st.session_state:
        st.session_state.show_exchanges = st.session_state.saved_show_exchanges
    if 'saved_max_points' in st.session_state and 'max_points' not in st.session_state:
        st.session_state.max_points = st.session_state.saved_max_points
    if 'saved_plot_columns' in st.session_state and 'plot_columns' not in st.session_state:
        st.session_state.plot_columns = st.session_state.saved_plot_columns
    
    # Set initial defaults for widgets if not in session state
    if 'max_points' not in st.session_state:
        st.session_state.max_points = 1000
    if 'plot_columns' not in st.session_state:
        st.session_state.plot_columns = 'Two columns'
    if 'auto_refresh' not in st.session_state:
        st.session_state.auto_refresh = True
    if 'refresh_interval' not in st.session_state:
        st.session_state.refresh_interval = 1.0
    
    # Extract non-objective metrics
    base_metrics = [m for m in available_metrics if "Objective" not in m]
    
    # Additional metrics - widget value automatically preserved via key
    st.sidebar.markdown("Additional metrics")
    additional_metrics = st.sidebar.multiselect(
        "Additional metrics",
        options=sorted(base_metrics),
        key="additional_base_metrics",
        help="Note: 'All perturbations' shows metrics only for accepted perturbations (rejected ones have no metrics)",
        label_visibility="collapsed"
    )
    
    # History type selector - choose which event type to display
    st.sidebar.markdown("History type")
    history_type = st.sidebar.radio(
        "History type",
        options=["Improvements only", "Accepted steps", "All perturbations"],
        key="history_type",
        help="Improvements: Only new best values (monotonically improving)\n"
             "Accepted steps: All SA acceptances (includes exploration)\n"
             "All perturbations: Sampled perturbations (at db_step_interval)\n"
             "  - Includes both accepted and rejected\n"
             "  - Metrics shown only for accepted perturbations",
        label_visibility="collapsed"
    )
    
    # Map display name to internal key
    history_type_map = {
        "Improvements only": "improvements",
        "Accepted steps": "accepted",
        "All perturbations": "perturbations"
    }
    history_key = history_type_map[history_type]
    
    # Objective is always "Objective value"
    objective_metric = "Objective value"
    
    # Layout - widget value automatically preserved via key
    st.sidebar.markdown("Plot layout")

    plot_columns = st.sidebar.radio(
        "Plot layout",
        options=["One column", "Two columns"],
        key="plot_columns",
        help="Switch between two-column or single-column plot layout",
        label_visibility="collapsed"
    )
    
    # Other options subsection
    st.sidebar.markdown("Other")
    normalize_metrics = st.sidebar.checkbox(
        "Normalize", 
        key="normalize_metrics",
        help="Scale all metrics to [0, 1] range for easier comparison when they have different scales"
    )
    show_exchanges = st.sidebar.checkbox(
        "Show exchanges",
        key="show_exchanges",
        help="Draw vertical markers at replica exchange events"
    )

    n_cols = 2 if plot_columns == "Two columns" else 1
    
    # Set max points to constant value of 1000 (not exposed to user)
    max_points = 1000
    
    return {
        'history_type': history_key,
        'objective_metric': objective_metric,
        'additional_metrics': additional_metrics,
        'normalize_metrics': normalize_metrics,
        'show_exchanges': show_exchanges,
        'max_points': max_points,
        'n_cols': n_cols
    }


def render_run_information(metadata: Dict[str, Any]) -> None:
    """Render run information section in sidebar.
    
    Args:
        metadata (Dict[str, Any]): Dictionary with run metadata including start_time,
            hyperparameters, n_replicas, checkpoint_file, objective_function_name,
            and dataset_size.
    """

    if st is None:
        return
        
    st.sidebar.markdown("---")
    st.sidebar.subheader("Run information")
    
    hyperparams = metadata['hyperparameters']
    
    # Format display values
    max_time_minutes = hyperparams.get('max_time', 0)

    if max_time_minutes:
        if max_time_minutes >= 60:  # 60 minutes or more
            max_time_display = f"{max_time_minutes / 60:.2f} hr"
        else:
            max_time_display = f"{max_time_minutes:.0f} min"
    else:
        max_time_display = 'N/A'
    
    checkpoint_file = metadata.get('checkpoint_file')
    checkpoint_display = Path(checkpoint_file).name if checkpoint_file else 'None'
    
    objective_func_name = metadata.get('objective_function_name', 'N/A') or 'N/A'
    
    dataset_size = metadata.get('dataset_size')
    dataset_size_display = f"{dataset_size:,}" if dataset_size else 'N/A'
    
    run_info = f"""**Started:** {datetime.fromtimestamp(metadata['start_time']).strftime('%Y-%m-%d %H:%M')}  
**Max time:** {max_time_display}  
**Dataset size:** {dataset_size_display}  
**Replicas:** {metadata['n_replicas']}  
**Objective:** {objective_func_name}  
**Checkpoint:** {checkpoint_display}"""
    
    st.sidebar.markdown(run_info)


def render_hyperparameters(metadata: Dict[str, Any]) -> None:
    """Render hyperparameters section in sidebar.
    
    Args:
        metadata (Dict[str, Any]): Dictionary with run metadata containing hyperparameters
            and exchange_interval.
    """

    if st is None:
        return
        
    st.sidebar.markdown("---")
    st.sidebar.subheader("Hyperparameters")
    
    hyperparams = metadata['hyperparameters']
    
    # Format temperatures in scientific notation
    t_min = hyperparams.get('T_min', 'N/A')
    t_max = hyperparams.get('T_max', 'N/A')
    t_min_str = f"{t_min:.1e}" if isinstance(t_min, (int, float)) else t_min
    t_max_str = f"{t_max:.1e}" if isinstance(t_max, (int, float)) else t_max
    
    initial_step_spread = hyperparams.get('initial_step_spread', hyperparams.get('step_spread', 'N/A'))
    final_step_spread = hyperparams.get('final_step_spread', 'N/A')
    step_spread_text = f"**Initial step spread:** {initial_step_spread}  \n"

    if final_step_spread != 'N/A':
        step_spread_text += f"**Final step spread:** {final_step_spread}  \n"
    
    hyperparams_text = f"""**Mode:** {hyperparams.get('mode', 'N/A')}  
{step_spread_text}**Perturb fraction:** {hyperparams.get('perturb_fraction', 'N/A')}  
**Cooling rate:** {hyperparams.get('cooling_rate', 'N/A')}  
**Exchange interval:** {metadata['exchange_interval']}  
**Exchange strategy:** {hyperparams.get('exchange_strategy', 'N/A')}  
**T_min:** {t_min_str}  
**T_max:** {t_max_str}"""
    
    st.sidebar.markdown(hyperparams_text)


def render_temperature_ladder(temp_ladder_df: pd.DataFrame) -> None:
    """Render temperature ladder section in sidebar.
    
    Args:
        temp_ladder_df (pd.DataFrame): DataFrame with replica_id and temperature columns.
    """

    if st is None:
        return
        
    st.sidebar.markdown("---")
    st.sidebar.subheader("Temperature ladder")
    
    if not temp_ladder_df.empty:
        temp_lines = []

        for _, row in temp_ladder_df.iterrows():
            temp_lines.append(f"**Replica {int(row['replica_id'])}:** {row['temperature']:.1e}")

        st.sidebar.markdown("  \n".join(temp_lines))


def render_leaderboard(leaderboard_df: pd.DataFrame) -> None:
    """Render replica leaderboard in main content area.
    
    Args:
        leaderboard_df (pd.DataFrame): DataFrame with replica stats including replica_id,
            best_objective, current_perturbation_num, and temperature.
    """
    if st is None:
        return
        
    st.header("Replica leaderboard")
    
    if not leaderboard_df.empty:

        cols = st.columns(3)
        medals = ['1<sup>st</sup>:', '2<sup>nd</sup>:', '3<sup>rd</sup>:']

        for idx, (_, row) in enumerate(leaderboard_df.iterrows()):
            with cols[idx]:
                st.markdown(f"### {medals[idx]} Replica {int(row['replica_id'])}", unsafe_allow_html=True)
                st.markdown(f"**Objective:** {row['best_objective']:.4f}")
    else:
        st.info("No replica data available yet")


def render_progress_stats(stats: Dict[str, Any], metadata: Dict[str, Any]) -> None:
    """Render progress statistics in main content area.
    
    Args:
        stats (Dict[str, Any]): Dictionary with total_perturbations and total_accepted.
        metadata (Dict[str, Any]): Dictionary with run metadata including start_time.
    """

    if st is None:
        return
    
    total_perturbations = stats.get('total_perturbations', 0)
    total_accepted = stats.get('total_accepted', 0)
    
    # Calculate elapsed time - use end_time if run is complete, otherwise current time
    end_time = metadata.get('end_time')

    if end_time is not None:
        elapsed_time = end_time - metadata['start_time']

    else:
        elapsed_time = time.time() - metadata['start_time']
    
    # Format elapsed time
    if elapsed_time > 3600:  # More than 60 minutes
        elapsed_display = f"{elapsed_time / 3600:.2f} hr"

    elif elapsed_time > 60:  # More than 60 seconds
        elapsed_display = f"{int(elapsed_time / 60)} min"

    else:
        elapsed_display = f"{int(elapsed_time)} sec"
    
    st.subheader(f"Optimization progress ({elapsed_display})")
    
    if elapsed_time > 0 and total_perturbations:
        exploration_rate = total_perturbations / elapsed_time
        progress_rate = total_accepted / elapsed_time
        acceptance_rate = (total_accepted / total_perturbations * 100) if total_perturbations > 0 else 0
        
        st.markdown(
            f"**Exploration rate:** {exploration_rate:,.1f} perturbations/sec | "
            f"**Progress rate:** {progress_rate:,.1f} accepted/sec | "
            f"**Acceptance rate:** {acceptance_rate:.1f}% ({total_accepted:,} / {total_perturbations:,})"
        )

    else:
        st.markdown("**Exploration rate:** N/A")
