ui_components.py

"""
UI components for Streamlit application
"""

import streamlit as st
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from typing import Dict, Optional
from config import (
    AVAILABLE_MODELS,
    DEFAULT_MODEL,
    DEFAULT_TEMPERATURE,
    DEFAULT_MAX_TOKENS,
    DEFAULT_SAMPLE_COUNT,
    SAMPLE_COUNT_OPTIONS,
    SIDEBAR_TITLE
)


def render_sidebar() -> Dict:
    """
    Render the sidebar with configuration options

    Returns:
        Dictionary with configuration values
    """
    st.sidebar.title(SIDEBAR_TITLE)

    # API Key section
    st.sidebar.subheader("🔑 OpenAI API Key")

    # Initialize session state for API key if not exists
    if "api_key" not in st.session_state:
        st.session_state.api_key = ""

    # Initialize available models in session state
    if "available_models" not in st.session_state:
        st.session_state.available_models = AVAILABLE_MODELS

    api_key_input = st.sidebar.text_input(
        "API Key",
        value=st.session_state.api_key,
        type="password",
        help="Enter your OpenAI API key",
        label_visibility="collapsed"
    )

    # Save/Validate button
    col1, col2 = st.sidebar.columns([1, 1])
    with col1:
        if st.button("💾 Save & Validate API Key", use_container_width=True):
            st.session_state.api_key = api_key_input
            st.session_state.api_key_validated = False  # Reset validation
            st.rerun()

    st.sidebar.divider()

    # Model and Parameters section
    st.sidebar.subheader("🤖 Model & Parameters")

    # Use models from session state
    available_models = st.session_state.available_models

    # Model selection
    model = st.sidebar.selectbox(
        "Model",
        options=available_models,
        index=available_models.index(DEFAULT_MODEL) if DEFAULT_MODEL in available_models else 0,
        help="Select the OpenAI model to use"
    )

    # Refresh models button
    if st.sidebar.button("🔄 Refresh Models", help="Fetch available models from your API key"):
        st.session_state.refresh_models = True
        st.rerun()

    # Temperature
    temperature = st.sidebar.slider(
        "Temperature",
        min_value=0.0,
        max_value=2.0,
        value=DEFAULT_TEMPERATURE,
        step=0.1,
        help="Higher values make output more random, lower values more deterministic"
    )

    # Max Tokens
    max_tokens = st.sidebar.number_input(
        "Max Tokens",
        min_value=100,
        max_value=4000,
        value=DEFAULT_MAX_TOKENS,
        step=100,
        help="Maximum number of tokens in the response"
    )

    # Sample count
    sample_count = st.sidebar.selectbox(
        "Number of Samples (responses)",
        options=SAMPLE_COUNT_OPTIONS,
        index=SAMPLE_COUNT_OPTIONS.index(DEFAULT_SAMPLE_COUNT) if DEFAULT_SAMPLE_COUNT in SAMPLE_COUNT_OPTIONS else 2,
        help="Number of LLM responses to generate for analysis"
    )

    return {
        "api_key": st.session_state.api_key,
        "model": model,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "sample_count": sample_count,
    }


def render_summary_stats(stats: Dict):
    """Render summary statistics in metric cards"""
    col1, col2, col3, col4 = st.columns(4)

    with col1:
        st.metric("Total Responses", stats["total_responses"])

    with col2:
        st.metric("Total Mentions", stats["total_mentions"])

    with col3:
        st.metric("Unique Brands", stats["unique_brands"])

    with col4:
        st.metric("Avg Brands/Response", f"{stats['avg_brands_per_response']:.1f}")


def render_brands_ranking_table(df: pd.DataFrame):
    """Render the brands ranking table"""
    st.subheader("📊 Brand Ranking")

    # Display dataframe with full width
    st.dataframe(
        df,
        use_container_width=True,
        hide_index=True,
        column_config={
            "Brand": st.column_config.TextColumn("Brand", width="medium"),
            "Visibility Rate": st.column_config.TextColumn("Visibility %", width="small"),
            "Top-1 Share": st.column_config.TextColumn("Top-1 %", width="small"),
            "Avg Position": st.column_config.TextColumn("Avg Pos", width="small"),
            "Mention Share": st.column_config.TextColumn("Mention %", width="small"),
            "Total Mentions": st.column_config.NumberColumn("Mentions", width="small"),
            "Appearances": st.column_config.NumberColumn("Appears", width="small"),
        }
    )


def render_visibility_chart(df: pd.DataFrame, top_n: int = 10):
    """Render visibility rate bar chart"""
    st.subheader(f"📈 Top {top_n} Brands by Visibility")

    # Take top N brands
    df_top = df.head(top_n).copy()

    # Convert percentage strings to floats for plotting
    df_top["Visibility_Value"] = df_top["Visibility Rate"].str.rstrip("%").astype(float)

    # Create bar chart
    fig = px.bar(
        df_top,
        x="Visibility_Value",
        y="Brand",
        orientation="h",
        labels={"Visibility_Value": "Visibility Rate (%)", "Brand": "Brand"},
        text="Visibility Rate",
        color="Visibility_Value",
        color_continuous_scale="Blues"
    )

    fig.update_layout(
        showlegend=False,
        height=400,
        yaxis={"categoryorder": "total ascending"}
    )

    fig.update_traces(textposition="outside")

    st.plotly_chart(fig, use_container_width=True)


def render_comparison_chart(df: pd.DataFrame, top_n: int = 10):
    """Render comparison chart with multiple metrics"""
    st.subheader(f"🔍 Multi-Metric Comparison (Top {top_n})")

    # Take top N brands
    df_top = df.head(top_n).copy()

    # Convert percentage strings to floats
    df_top["Visibility_Value"] = df_top["Visibility Rate"].str.rstrip("%").astype(float)
    df_top["Top1_Value"] = df_top["Top-1 Share"].str.rstrip("%").astype(float)
    df_top["Mention_Value"] = df_top["Mention Share"].str.rstrip("%").astype(float)

    # Create grouped bar chart
    fig = go.Figure()

    fig.add_trace(go.Bar(
        name="Visibility Rate",
        x=df_top["Brand"],
        y=df_top["Visibility_Value"],
        text=df_top["Visibility Rate"],
        textposition="outside",
    ))

    fig.add_trace(go.Bar(
        name="Top-1 Share",
        x=df_top["Brand"],
        y=df_top["Top1_Value"],
        text=df_top["Top-1 Share"],
        textposition="outside",
    ))

    fig.add_trace(go.Bar(
        name="Mention Share",
        x=df_top["Brand"],
        y=df_top["Mention_Value"],
        text=df_top["Mention Share"],
        textposition="outside",
    ))

    fig.update_layout(
        barmode="group",
        xaxis_title="Brand",
        yaxis_title="Percentage (%)",
        height=500,
        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1)
    )

    st.plotly_chart(fig, use_container_width=True)


def export_results_csv(df: pd.DataFrame, stats: Dict, filename: str = "brand_analysis") -> str:
    """
    Export results to CSV format

    Returns:
        CSV string
    """
    # Add summary stats as header comments
    csv_content = f"# LLM Brand Visibility Analysis\n"
    csv_content += f"# Total Responses: {stats['total_responses']}\n"
    csv_content += f"# Total Mentions: {stats['total_mentions']}\n"
    csv_content += f"# Unique Brands: {stats['unique_brands']}\n"
    csv_content += f"# Avg Brands per Response: {stats['avg_brands_per_response']:.2f}\n"
    csv_content += "#\n"

    # Add dataframe
    csv_content += df.to_csv(index=False)

    return csv_content


def export_results_txt(df: pd.DataFrame, stats: Dict) -> str:
    """
    Export results to TXT format

    Returns:
        TXT string
    """
    txt_content = "=" * 60 + "\n"
    txt_content += "LLM BRAND VISIBILITY ANALYSIS REPORT\n"
    txt_content += "=" * 60 + "\n\n"

    # Summary statistics
    txt_content += "SUMMARY STATISTICS:\n"
    txt_content += "-" * 60 + "\n"
    txt_content += f"Total Responses:          {stats['total_responses']}\n"
    txt_content += f"Total Mentions:           {stats['total_mentions']}\n"
    txt_content += f"Unique Brands:            {stats['unique_brands']}\n"
    txt_content += f"Avg Brands per Response:  {stats['avg_brands_per_response']:.2f}\n"
    txt_content += "\n"

    # Brand ranking
    txt_content += "BRAND RANKING:\n"
    txt_content += "-" * 60 + "\n"
    txt_content += df.to_string(index=False)
    txt_content += "\n\n"

    txt_content += "=" * 60 + "\n"
    txt_content += "End of Report\n"
    txt_content += "=" * 60 + "\n"

    return txt_content