"""
Company Analytics Dashboard Template

A company leaderboard dashboard demonstrating:
- Interactive dataframe with sparkline columns
- Segmented control for ranking (top spenders, gainers, shrinkers)
- Multi-select pills for account type filtering
- Time window filtering
- Growth score calculation
- Dialog popup for company details

This template uses synthetic data. Replace generate_company_data()
with your actual data source (e.g., Snowflake queries, CRM APIs, etc.)
"""

from datetime import date, timedelta
import numpy as np
import pandas as pd
import streamlit as st
import altair as alt

st.set_page_config(
    page_title="Company Analytics",
    page_icon=":material/business:",
    layout="wide",
)


# =============================================================================
# Synthetic Data Generation (Replace with your data source)
# =============================================================================

COMPANY_NAMES = [
    "Acme Corp", "TechFlow Inc", "DataDriven Co", "CloudFirst Ltd", 
    "InnovateTech", "ScaleUp Systems", "PrimeData Inc", "FutureStack",
    "ByteWise Corp", "StreamLine Co", "Quantum Labs", "NexGen Solutions",
    "AlphaMetrics", "BetaAnalytics", "GammaInsights", "DeltaData",
    "OmegaTech", "SigmaSoft", "ThetaCloud", "ZetaDigital",
]

ACCOUNT_TYPES = ["Enterprise", "Growth", "Startup", "Trial", "Internal"]
REGIONS = ["North America", "EMEA", "APAC", "LATAM"]
SEGMENTS = ["Technology", "Finance", "Healthcare", "Retail", "Manufacturing"]


@st.cache_data(ttl=3600)
def generate_company_data(days: int = 90) -> pd.DataFrame:
    """Generate synthetic company usage data.
    
    Replace this function with your actual data source.
    """
    np.random.seed(42)
    
    end_date = date.today() - timedelta(days=1)
    start_date = end_date - timedelta(days=days)
    dates = pd.date_range(start=start_date, end=end_date, freq="D")
    
    records = []
    
    for company in COMPANY_NAMES:
        # Assign static attributes
        account_type = np.random.choice(ACCOUNT_TYPES, p=[0.3, 0.25, 0.2, 0.15, 0.1])
        region = np.random.choice(REGIONS)
        segment = np.random.choice(SEGMENTS)
        
        # Generate usage pattern
        base_usage = np.random.randint(100, 10000)
        growth = np.random.uniform(-0.005, 0.01)  # Some companies shrink
        
        for i, dt in enumerate(dates):
            # Base trend
            trend = base_usage * (1 + growth) ** i
            
            # Weekly seasonality
            if dt.dayofweek >= 5:
                trend *= 0.3
            
            # Random noise
            daily_credits = max(0, trend * np.random.uniform(0.7, 1.3))
            
            records.append({
                "company_name": company,
                "date": dt,
                "daily_credits": daily_credits,
                "account_type": account_type,
                "region": region,
                "segment": segment,
            })
    
    return pd.DataFrame(records)


@st.cache_data(ttl=3600)
def load_company_data() -> pd.DataFrame:
    """Load all company data."""
    return generate_company_data(days=90)


def aggregate_companies(
    df: pd.DataFrame,
    days: int | None = None,
    account_types: list[str] | None = None,
    sort_by: str = "total_credits",
) -> pd.DataFrame:
    """Filter and aggregate company data."""
    result = df.copy()
    
    # Filter by time window
    if days:
        cutoff = pd.Timestamp.now() - pd.Timedelta(days=days)
        result = result[result["date"] >= cutoff]
    
    # Filter by account type
    if account_types:
        result = result[result["account_type"].isin(account_types)]
    
    if result.empty:
        return pd.DataFrame()
    
    # Aggregate to company level
    agg = result.groupby("company_name").agg(
        total_credits=("daily_credits", "sum"),
        active_days=("date", "nunique"),
        account_type=("account_type", "first"),
        region=("region", "first"),
        segment=("segment", "first"),
    ).reset_index()
    
    # Calculate daily average
    agg["daily_avg"] = agg["total_credits"] / agg["active_days"]
    
    # Build sparkline data (list of daily values)
    sparklines = (
        result.groupby("company_name")
        .apply(lambda x: x.sort_values("date")["daily_credits"].tolist())
        .reset_index()
    )
    sparklines.columns = ["company_name", "usage_trend"]
    agg = agg.merge(sparklines, on="company_name")
    
    # Calculate growth score (second half vs first half)
    def calc_growth(trend):
        if not trend or len(trend) < 2:
            return 0
        mid = len(trend) // 2
        first_half = sum(trend[:mid]) if mid > 0 else 0
        second_half = sum(trend[mid:])
        return second_half - first_half
    
    agg["growth_score"] = agg["usage_trend"].apply(calc_growth)
    
    # Sort
    if sort_by == "growth_asc":
        agg = agg.sort_values("growth_score", ascending=True)
    elif sort_by == "growth_desc":
        agg = agg.sort_values("growth_score", ascending=False)
    else:
        agg = agg.sort_values("total_credits", ascending=False)
    
    return agg


def render_company_dialog(company_name: str, company_row: pd.Series, df: pd.DataFrame):
    """Render company details inside a dialog."""
    company_data = df[df["company_name"] == company_name].sort_values("date")
    
    if company_data.empty:
        st.warning("No data available for this company.")
        return
    
    # Company info badges - extract from list format back to single value
    account_type = company_row["account_type"][0] if company_row["account_type"] else "Unknown"
    region = company_row["region"][0] if company_row["region"] else "Unknown"
    segment = company_row["segment"][0] if company_row["segment"] else "Unknown"
    total_credits = company_row["total_credits"]
    
    st.markdown(
        f":blue-badge[{account_type}] "
        f":violet-badge[{region}] "
        f":orange-badge[{segment}] "
        f":green-badge[{total_credits:,.0f} credits]"
    )
    
    # Summary metrics
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric("Total Credits", f"{total_credits:,.0f}")
    with col2:
        st.metric("Active Days", f"{company_row['active_days']:,}")
    with col3:
        growth = company_row["growth_score"]
        st.metric("Growth Score", f"{growth:+,.0f}")
    
    # Charts
    col1, col2 = st.columns(2)
    
    with col1:
        with st.container(border=True):
            st.markdown("**Daily usage**")
            st.line_chart(company_data, x="date", y="daily_credits", height=250)
    
    with col2:
        with st.container(border=True):
            st.markdown("**Cumulative usage**")
            chart_data = company_data.copy()
            chart_data["cumulative"] = chart_data["daily_credits"].cumsum()
            st.area_chart(chart_data, x="date", y="cumulative", height=250)


# =============================================================================
# Page Layout
# =============================================================================

# Load data
all_data = load_company_data()

st.markdown("# :material/business: Company Analytics")
st.caption("Track company adoption - usage, growth trends, and account details.")

# Filters
with st.container(border=True):
    st.markdown("**Filters**")
    
    # Company selection mode
    sort_mode = st.segmented_control(
        "Sort by",
        options=[
            "All companies",
            ":material/military_tech: Top spenders",
            ":material/trending_down: Top shrinkers",
            ":material/trending_up: Top gainers",
        ],
        default="All companies",
    )
    
    # Time window
    timeframe_options = {
        "All time": None,
        "Last 28 days": 28,
        "Last 7 days": 7,
    }
    timeframe = st.segmented_control(
        "Time window",
        options=list(timeframe_options.keys()),
        default="Last 28 days",
    )
    days_filter = timeframe_options.get(timeframe)
    
    # Account types
    account_types = st.pills(
        "Account types",
        options=ACCOUNT_TYPES,
        default=["Enterprise", "Growth", "Startup"],
        selection_mode="multi",
    )

# Determine sort order
if "Top shrinkers" in (sort_mode or ""):
    sort_by = "growth_asc"
elif "Top gainers" in (sort_mode or ""):
    sort_by = "growth_desc"
else:
    sort_by = "total_credits"

# Get filtered data
leaderboard = aggregate_companies(
    all_data,
    days=days_filter,
    account_types=account_types,
    sort_by=sort_by,
)

if leaderboard.empty:
    st.warning("No company data found for the selected filters.")
    st.stop()


def _to_list(val):
    """Convert a single value to a list for MultiselectColumn display."""
    return [val] if pd.notna(val) else []


# Convert columns to lists for MultiselectColumn display (shows nice colored chips)
for col in ["account_type", "region", "segment"]:
    leaderboard[col] = leaderboard[col].apply(_to_list)

# Companies dataframe
with st.container(border=True):
    timeframe_text = timeframe.lower() if timeframe != "All time" else "all time"
    st.markdown(f"**Companies — {timeframe_text}**")
    
    # Selection dataframe with cell-click support
    selection = st.dataframe(
        leaderboard,
        column_config={
            "company_name": st.column_config.TextColumn(
                "Company (👋 click to view details)",
                width="medium",
            ),
            "account_type": st.column_config.MultiselectColumn(
                "Type",
                options=ACCOUNT_TYPES,
                color="auto",
                width="small",
            ),
            "total_credits": st.column_config.NumberColumn(
                "Credits",
                format="%.0f",
            ),
            "growth_score": st.column_config.NumberColumn(
                "Growth",
                format="%+.0f",
                help="Credit change: second half vs first half of period",
            ),
            "usage_trend": st.column_config.LineChartColumn(
                "Trend",
                width="medium",
            ),
            "daily_avg": st.column_config.NumberColumn(
                "Daily Avg",
                format="%.1f",
            ),
            "active_days": st.column_config.NumberColumn(
                "Active Days",
                format="%d",
            ),
            "region": st.column_config.MultiselectColumn(
                "Region",
                options=REGIONS,
                color="auto",
            ),
            "segment": st.column_config.MultiselectColumn(
                "Segment",
                options=SEGMENTS,
                color="auto",
            ),
        },
        column_order=[
            "company_name", "account_type", "total_credits", "growth_score",
            "usage_trend", "daily_avg", "region", "segment",
        ],
        hide_index=True,
        use_container_width=True,
        on_select="rerun",
        selection_mode="single-cell",
        key="company_leaderboard",
    )

# Company drill-down via dialog when Company column cell is clicked
if selection.selection.cells:
    cell = selection.selection.cells[0]  # tuple: (row_index, column_name)
    row_idx, col_name = cell
    # Check if the clicked cell is in the company_name column
    if col_name == "company_name":
        selected_company = leaderboard.iloc[row_idx]["company_name"]
        company_row = leaderboard.iloc[row_idx]

        @st.dialog(f"{selected_company}", width="large")
        def show_company_dialog():
            render_company_dialog(
                selected_company,
                company_row=company_row,
                df=all_data,
            )

        show_company_dialog()
