src/streamlit_app.py

import streamlit as st
import pandas as pd
from pymongo import MongoClient
from datetime import datetime, date
import pytz
import os
from dotenv import load_dotenv

load_dotenv()

st.set_page_config(
    page_title="Agent Assistance Usage",
    page_icon=None,
    layout="wide",
    initial_sidebar_state="expanded"
)

st.markdown("""
<style>
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&family=IBM+Plex+Mono:wght@400;600&display=swap');

:root {
    --bg:        #f5f6fa;
    --surface:   #ffffff;
    --surface2:  #f0f2f8;
    --border:    #e2e5ef;
    --border2:   #c8cce0;
    --accent:    #2563eb;
    --accent-lt: #eff4ff;
    --accent2:   #0ea5e9;
    --green:     #16a34a;
    --amber:     #d97706;
    --red:       #dc2626;
    --text:      #111827;
    --text2:     #374151;
    --muted:     #6b7280;
    --muted2:    #9ca3af;
    --gap:       1.25rem;
}

html, body, [class*="css"] {
    font-family: 'Inter', sans-serif !important;
    background-color: var(--bg) !important;
    color: var(--text) !important;
}

#MainMenu, footer, header { visibility: hidden; }

.main .block-container {
    padding: 0 2.5rem 3rem;
    max-width: 1440px;
}

.topbar {
    background: var(--surface);
    border-bottom: 1px solid var(--border);
    margin: 0 -2.5rem 2.5rem -2.5rem;
    padding: 1.1rem 2.5rem;
    display: flex;
    align-items: center;
    justify-content: space-between;
}
.topbar-title {
    font-size: 1.1rem;
    font-weight: 700;
    color: var(--text);
    letter-spacing: -0.02em;
}
.topbar-sub {
    font-size: 0.73rem;
    color: var(--muted);
    margin-top: 2px;
}
.rate-badge {
    background: var(--accent-lt);
    border: 1px solid #bfdbfe;
    color: var(--accent);
    border-radius: 6px;
    font-size: 0.71rem;
    font-weight: 600;
    padding: 0.28rem 0.75rem;
    font-family: 'IBM Plex Mono', monospace;
}

section[data-testid="stSidebar"] {
    background: var(--surface) !important;
    border-right: 1px solid var(--border) !important;
}
section[data-testid="stSidebar"] .block-container {
    padding: 1.75rem 1.25rem;
}
.sb-brand {
    font-size: 0.7rem;
    font-weight: 700;
    text-transform: uppercase;
    letter-spacing: 0.14em;
    color: var(--accent);
    padding-bottom: 0.85rem;
    border-bottom: 2px solid var(--accent);
    margin-bottom: 1.5rem;
    display: block;
}
.sb-section {
    font-size: 0.66rem;
    font-weight: 700;
    text-transform: uppercase;
    letter-spacing: 0.12em;
    color: var(--muted2);
    margin: 1.5rem 0 0.6rem;
}

.stSelectbox > div > div,
.stDateInput > div > div {
    background: var(--surface2) !important;
    border: 1px solid var(--border2) !important;
    border-radius: 7px !important;
    color: var(--text) !important;
    font-size: 0.83rem !important;
}
.stSelectbox label, .stDateInput label {
    color: var(--text2) !important;
    font-size: 0.77rem !important;
    font-weight: 500 !important;
}

.stDateInput input {
    color: var(--text) !important;
    -webkit-text-fill-color: var(--text) !important;
}

.stButton > button {
    background: var(--accent) !important;
    color: white !important;
    border: none !important;
    border-radius: 8px !important;
    font-family: 'Inter', sans-serif !important;
    font-size: 0.84rem !important;
    font-weight: 600 !important;
    padding: 0.68rem 1.25rem !important;
    width: 100%;
    letter-spacing: 0.01em;
    margin-top: 0.5rem;
}
.stButton > button:hover { background: #1d4ed8 !important; }

.kpi-card {
    background: var(--surface);
    border: 1px solid var(--border);
    border-radius: 10px;
    padding: 1.25rem 1.4rem 1.1rem;
    position: relative;
    overflow: hidden;
    height: 100%;
}
.kpi-card::after {
    content: '';
    position: absolute;
    top: 0; left: 0; right: 0;
    height: 3px;
    border-radius: 10px 10px 0 0;
}
.kpi-card.c-blue::after  { background: var(--accent); }
.kpi-card.c-sky::after   { background: var(--accent2); }
.kpi-card.c-green::after { background: var(--green); }
.kpi-card.c-amber::after { background: var(--amber); }

.kpi-label {
    font-size: 0.68rem;
    font-weight: 700;
    text-transform: uppercase;
    letter-spacing: 0.1em;
    color: var(--muted);
    margin-bottom: 0.55rem;
}
.kpi-value {
    font-family: 'IBM Plex Mono', monospace;
    font-size: 1.55rem;
    font-weight: 600;
    color: var(--text);
    line-height: 1;
    margin-bottom: 0.35rem;
}
.kpi-sub {
    font-size: 0.71rem;
    color: var(--muted2);
    line-height: 1.45;
}

.sec-hdr {
    font-size: 0.68rem;
    font-weight: 700;
    text-transform: uppercase;
    letter-spacing: 0.1em;
    color: var(--muted);
    margin: 2rem 0 1rem;
    display: flex;
    align-items: center;
    gap: 0.6rem;
}
.sec-hdr::after {
    content: '';
    flex: 1;
    height: 1px;
    background: var(--border);
}

.col-hdr {
    display: grid;
    grid-template-columns: 1.8fr 1.1fr 120px 120px 130px;
    gap: 1rem;
    padding: 0 1.1rem 0.5rem;
    font-size: 0.66rem;
    font-weight: 700;
    text-transform: uppercase;
    letter-spacing: 0.08em;
    color: var(--muted2);
}
.data-row {
    display: grid;
    grid-template-columns: 1.8fr 1.1fr 120px 120px 130px;
    gap: 1rem;
    align-items: center;
    background: var(--surface);
    border: 1px solid var(--border);
    border-radius: 8px;
    padding: 0.9rem 1.1rem;
    margin-bottom: 0.5rem;
    transition: border-color 0.15s, box-shadow 0.15s;
}
.data-row:hover {
    border-color: var(--border2);
    box-shadow: 0 1px 6px rgba(0,0,0,0.05);
}
.row-name {
    font-size: 0.875rem;
    font-weight: 600;
    color: var(--text);
    margin-bottom: 4px;
}
.row-meta {
    font-size: 0.7rem;
    color: var(--muted);
    line-height: 1.45;
}
.token-val {
    font-family: 'IBM Plex Mono', monospace;
    font-size: 0.77rem;
    color: var(--text2);
    line-height: 1.6;
}
.cost-val {
    font-family: 'IBM Plex Mono', monospace;
    font-size: 0.8rem;
    color: var(--muted);
}
.cost-total {
    font-family: 'IBM Plex Mono', monospace;
    font-size: 0.88rem;
    font-weight: 600;
    color: var(--green);
    text-align: right;
}

.mbadge {
    display: inline-block;
    background: var(--accent-lt);
    border: 1px solid #bfdbfe;
    color: var(--accent);
    border-radius: 5px;
    padding: 0.12rem 0.5rem;
    font-size: 0.66rem;
    font-family: 'IBM Plex Mono', monospace;
    font-weight: 600;
}
.type-badge {
    display: inline-block;
    background: #f0fdf4;
    border: 1px solid #bbf7d0;
    color: var(--green);
    border-radius: 5px;
    padding: 0.12rem 0.5rem;
    font-size: 0.66rem;
    font-family: 'IBM Plex Mono', monospace;
    font-weight: 600;
    margin-left: 0.25rem;
}
.type-badge.email {
    background: #fff7ed;
    border-color: #fed7aa;
    color: var(--amber);
}

.stTabs [data-baseweb="tab-list"] {
    background: var(--surface2) !important;
    border-radius: 9px;
    border: 1px solid var(--border);
    padding: 0.2rem;
    gap: 0.15rem;
    margin-bottom: 0.5rem;
}
.stTabs [data-baseweb="tab"] {
    background: transparent !important;
    color: var(--muted) !important;
    border-radius: 7px !important;
    font-size: 0.8rem !important;
    font-weight: 500 !important;
    padding: 0.42rem 1rem !important;
}
.stTabs [aria-selected="true"] {
    background: var(--surface) !important;
    color: var(--text) !important;
    border: 1px solid var(--border) !important;
    font-weight: 600 !important;
    box-shadow: 0 1px 4px rgba(0,0,0,0.06);
}

.info-box {
    background: var(--accent-lt);
    border: 1px solid #bfdbfe;
    border-radius: 9px;
    padding: 1rem 1.3rem;
    font-size: 0.84rem;
    color: var(--text2);
    margin-bottom: 2rem;
    line-height: 1.65;
}

.block-spacer { margin-top: 2rem; }

.stDownloadButton > button {
    background: var(--surface) !important;
    color: var(--text2) !important;
    border: 1px solid var(--border2) !important;
    border-radius: 8px !important;
    font-size: 0.8rem !important;
    font-weight: 500 !important;
}
.stDownloadButton > button:hover { background: var(--surface2) !important; }

.stDataFrame { border-radius: 9px !important; overflow: hidden; }

.stProgress > div > div { background: var(--accent) !important; }

[data-testid="column"] { padding: 0 0.4rem !important; }
[data-testid="column"]:first-child { padding-left: 0 !important; }
[data-testid="column"]:last-child  { padding-right: 0 !important; }
</style>
""", unsafe_allow_html=True)


USD_TO_INR = 92.0

MODEL_PRICING = {
    "gpt-4.1-nano":    {"input": 0.10,  "output": 0.40},
    "gpt-4o-mini":     {"input": 0.15,  "output": 0.60},
    "gpt-4o":          {"input": 2.50,  "output": 10.00},
    "gpt-3.5-turbo":   {"input": 0.50,  "output": 1.50},
    "gpt-4-turbo":     {"input": 10.00, "output": 30.00},
    "claude-3-haiku":  {"input": 0.25,  "output": 1.25},
    "claude-3-sonnet": {"input": 3.00,  "output": 15.00},
    "claude-3-opus":   {"input": 15.00, "output": 75.00},
}

FEATURE_REGISTRY = [
    ("entity_extraction_logs",       "Entity Extraction",                  None),
    ("generate_nps_questions_logs",  "NPS Question Generation",            None),
    ("grammer_correction_logs",      "Grammar Correction",                 None),
    ("recommended_template_logs",    "Recommended Templates",              None),
    ("summarization_logs",           "Summarization",                      None),
    ("template_recommendation_logs", "Template Recommendation (Chat)",     "chat"),
    ("template_recommendation_logs", "Template Recommendation (Email)",    "email"),
    ("text_expansion_logs",          "Text Expansion",                     None),
    ("tone_change_logs",             "Tone Change",                        None),
]

FEATURE_DEFAULT_MODELS = {
    "Summarization": "gpt-4.1-nano",
    "Grammar Correction" : "gpt-4.1-nano"
}

NESTED_COLLECTIONS = {"template_recommendation_logs", "tone_change_logs"}


def fmt_inr(usd_val):
    return f"₹{usd_val * USD_TO_INR:,.2f}"


def build_excel(df):
    import io
    from openpyxl import Workbook
    from openpyxl.styles import Font, PatternFill, Alignment, Border, Side
    from openpyxl.utils import get_column_letter

    HEADER_FILL   = PatternFill("solid", fgColor="2563EB")
    HEADER_FONT   = Font(name="Calibri", bold=True, color="FFFFFF", size=10)
    TOTAL_FILL    = PatternFill("solid", fgColor="EFF4FF")
    TOTAL_FONT    = Font(name="Calibri", bold=True, color="111827", size=10)
    CELL_FONT     = Font(name="Calibri", size=10)
    BORDER_SIDE   = Side(style="thin", color="E2E5EF")
    THIN_BORDER   = Border(left=BORDER_SIDE, right=BORDER_SIDE,
                           top=BORDER_SIDE, bottom=BORDER_SIDE)

    COLUMNS = [
        ("Client ID",         "client_id",        18),
        ("Input Tokens",      "input_tokens",      14),
        ("Output Tokens",     "output_tokens",     14),
        ("Total Tokens",      "total_tokens",      14),
        ("Input Cost (INR)",  "input_cost_inr",    16),
        ("Output Cost (INR)", "output_cost_inr",   17),
        ("Total Cost (INR)",  "total_cost_inr",    16),
        ("Model",             "model",             16),
    ]

    wb = Workbook()
    wb.remove(wb.active)

    features = df["feature"].unique()

    for feature in features:
        sheet_name = feature[:31]
        ws = wb.create_sheet(title=sheet_name)

        ws.append([col[0] for col in COLUMNS])
        for col_idx, (_, _, width) in enumerate(COLUMNS, start=1):
            cell = ws.cell(row=1, column=col_idx)
            cell.fill      = HEADER_FILL
            cell.font      = HEADER_FONT
            cell.alignment = Alignment(horizontal="center", vertical="center")
            cell.border    = THIN_BORDER
            ws.column_dimensions[get_column_letter(col_idx)].width = width
        ws.row_dimensions[1].height = 20

        fdf = df[df["feature"] == feature].copy()
        client_data = fdf.groupby("client_id").agg(
            input_tokens=("input_tokens",     "sum"),
            output_tokens=("output_tokens",   "sum"),
            total_tokens=("total_tokens",     "sum"),
            input_cost_inr=("input_cost_inr", "sum"),
            output_cost_inr=("output_cost_inr","sum"),
            total_cost_inr=("total_cost_inr", "sum"),
            model=("model", "first"),
        ).reset_index().sort_values("total_cost_inr", ascending=False)

        for r_idx, row in enumerate(client_data.itertuples(index=False), start=2):
            values = [
                row.client_id,
                row.input_tokens,
                row.output_tokens,
                row.total_tokens,
                round(row.input_cost_inr,  4),
                round(row.output_cost_inr, 4),
                round(row.total_cost_inr,  4),
                row.model,
            ]
            for c_idx, val in enumerate(values, start=1):
                cell = ws.cell(row=r_idx, column=c_idx, value=val)
                cell.font      = CELL_FONT
                cell.border    = THIN_BORDER
                cell.alignment = Alignment(
                    horizontal="right" if c_idx > 1 else "left",
                    vertical="center"
                )

        total_row = len(client_data) + 2
        totals = [
            "TOTAL",
            int(client_data["input_tokens"].sum()),
            int(client_data["output_tokens"].sum()),
            int(client_data["total_tokens"].sum()),
            round(client_data["input_cost_inr"].sum(),  4),
            round(client_data["output_cost_inr"].sum(), 4),
            round(client_data["total_cost_inr"].sum(),  4),
            "",
        ]
        for c_idx, val in enumerate(totals, start=1):
            cell = ws.cell(row=total_row, column=c_idx, value=val)
            cell.fill      = TOTAL_FILL
            cell.font      = TOTAL_FONT
            cell.border    = THIN_BORDER
            cell.alignment = Alignment(
                horizontal="right" if c_idx > 1 else "left",
                vertical="center"
            )

        ws.freeze_panes = "A2"

    buf = io.BytesIO()
    wb.save(buf)
    buf.seek(0)
    return buf.getvalue()


def calc_cost(input_tokens, output_tokens, model):
    p = MODEL_PRICING[model]
    ic = (input_tokens  / 1_000_000) * p["input"]
    oc = (output_tokens / 1_000_000) * p["output"]
    return ic, oc, ic + oc


@st.cache_resource
def get_mongo_client():
    conn = os.getenv("DEV_MONGO_CONNECTION")
    return MongoClient(conn) if conn else None


def fetch_tokens(collection_name, label, ticket_filter, start_dt, end_dt, model):
    mongo = get_mongo_client()
    if mongo is None:
        return pd.DataFrame()

    col = mongo["agent_assistance"][collection_name]

    if collection_name == "template_recommendation_logs" and ticket_filter == "email":
        match = {
            "created_at": {"$gt": start_dt, "$lt": end_dt},
            "ticket_type": "email",
            "template": {"$ne": []},
        }
    elif collection_name == "template_recommendation_logs" and ticket_filter == "chat":
        match = {
            "created_at": {"$gt": start_dt, "$lt": end_dt},
            "ticket_type": {"$ne": "email"},
            "client_id": {"$exists": True, "$ne": None},
        }
    else:
        match = {
            "created_at": {"$gte": start_dt, "$lt": end_dt},
            "client_id":  {"$exists": True, "$ne": None},
        }
        if ticket_filter == "email":
            match["ticket_type"] = "email"
        elif ticket_filter == "chat":
            match["ticket_type"] = {"$ne": "email"}

    if collection_name in NESTED_COLLECTIONS:
        pipeline = [
            {"$match": match},
            {"$group": {"_id": "$client_id", "usage_arrays": {"$push": "$usage"}}},
            {"$unwind": "$usage_arrays"},
            {"$unwind": "$usage_arrays"},
            {"$group": {
                "_id": "$_id",
                "total_tokens":  {"$sum": "$usage_arrays.total_tokens"},
                "input_tokens":  {"$sum": "$usage_arrays.prompt_tokens"},
                "output_tokens": {"$sum": "$usage_arrays.completion_tokens"},
            }},
        ]
    else:
        pipeline = [
            {"$match": match},
            {"$unwind": "$usage"},
            {"$group": {
                "_id": "$client_id",
                "total_tokens":  {"$sum": "$usage.total_tokens"},
                "input_tokens":  {"$sum": "$usage.prompt_tokens"},
                "output_tokens": {"$sum": "$usage.completion_tokens"},
            }},
        ]

    pipeline.append({
        "$project": {
            "_id": 0,
            "client_id":     "$_id",
            "total_tokens":  1,
            "input_tokens":  1,
            "output_tokens": 1,
        }
    })

    try:
        results = list(col.aggregate(pipeline))
        if not results:
            return pd.DataFrame()
        df = pd.DataFrame(results)
        df["model"]          = model
        df["collection"]     = collection_name
        df["feature"]        = label
        df["ticket_filter"]  = ticket_filter or "all"
        ic, oc, tc = zip(*df.apply(
            lambda r: calc_cost(r["input_tokens"], r["output_tokens"], model), axis=1
        ))
        df["input_cost_usd"]  = ic
        df["output_cost_usd"] = oc
        df["total_cost_usd"]  = tc
        df["input_cost_inr"]  = df["input_cost_usd"]  * USD_TO_INR
        df["output_cost_inr"] = df["output_cost_usd"] * USD_TO_INR
        df["total_cost_inr"]  = df["total_cost_usd"]  * USD_TO_INR
        return df
    except Exception as e:
        st.error(f"Error fetching `{collection_name}` ({label}): {e}")
        return pd.DataFrame()


with st.sidebar:
    st.markdown('<span class="sb-brand">Agent Assistance</span>', unsafe_allow_html=True)

    st.markdown('<div class="sb-section">Date Range</div>', unsafe_allow_html=True)
    start_date = st.date_input("From", value=date(2026, 3, 29))
    end_date   = st.date_input("To",   value=date(2026, 4, 1))

    st.markdown('<div class="sb-section">Model Assignment</div>', unsafe_allow_html=True)
    st.caption("Assign a model to each feature.")

    feature_models = {}
    seen_labels = set()
    for (coll, label, _) in FEATURE_REGISTRY:
        if label not in seen_labels:
            default_model = FEATURE_DEFAULT_MODELS.get(label, "gpt-4o-mini")
            default_index = list(MODEL_PRICING.keys()).index(default_model)
            feature_models[label] = st.selectbox(
                label,
                options=list(MODEL_PRICING.keys()),
                index=default_index,
                key=f"model_{label}",
            )
            seen_labels.add(label)

    st.markdown("<div style='margin-top:1.5rem'></div>", unsafe_allow_html=True)
    run = st.button("Run Analysis", use_container_width=True)


if run:
    if not get_mongo_client():
        st.error("MongoDB connection not found. Set DEV_MONGO_CONNECTION in your .env file.")
        st.stop()

    start_dt = datetime(start_date.year, start_date.month, start_date.day, tzinfo=pytz.UTC)
    end_dt   = datetime(end_date.year,   end_date.month,   end_date.day,   tzinfo=pytz.UTC)

    all_dfs = []
    bar = st.progress(0, text="Fetching data…")

    for i, (coll, label, ticket_filter) in enumerate(FEATURE_REGISTRY):
        bar.progress(
            (i + 1) / len(FEATURE_REGISTRY),
            text=f"Fetching {label}…"
        )
        model = feature_models[label]
        df = fetch_tokens(coll, label, ticket_filter, start_dt, end_dt, model)
        if not df.empty:
            all_dfs.append(df)

    bar.empty()

    if not all_dfs:
        st.info("No data found for the selected date range.")
        st.stop()

    full_df = pd.concat(all_dfs, ignore_index=True)

    total_input  = int(full_df["input_tokens"].sum())
    total_output = int(full_df["output_tokens"].sum())
    total_tokens = int(full_df["total_tokens"].sum())
    pct_in  = total_input  / total_tokens * 100 if total_tokens else 0
    pct_out = total_output / total_tokens * 100 if total_tokens else 0

    st.markdown("<div style='margin-bottom: 0.25rem'></div>", unsafe_allow_html=True)
    c1, c2, c3, c4 = st.columns(4, gap="medium")

    with c1:
        st.markdown(f"""
        <div class="kpi-card c-blue">
            <div class="kpi-label">Total Tokens</div>
            <div class="kpi-value">{total_tokens:,}</div>
            <div class="kpi-sub">Input + Output combined</div>
        </div>""", unsafe_allow_html=True)

    with c2:
        st.markdown(f"""
        <div class="kpi-card c-sky">
            <div class="kpi-label">Input Tokens</div>
            <div class="kpi-value">{total_input:,}</div>
            <div class="kpi-sub">{pct_in:.1f}% of total &nbsp;&bull;&nbsp; {fmt_inr(full_df['input_cost_usd'].sum())}</div>
        </div>""", unsafe_allow_html=True)

    with c3:
        st.markdown(f"""
        <div class="kpi-card c-green">
            <div class="kpi-label">Output Tokens</div>
            <div class="kpi-value">{total_output:,}</div>
            <div class="kpi-sub">{pct_out:.1f}% of total &nbsp;&bull;&nbsp; {fmt_inr(full_df['output_cost_usd'].sum())}</div>
        </div>""", unsafe_allow_html=True)

    with c4:
        st.markdown(f"""
        <div class="kpi-card c-amber">
            <div class="kpi-label">Approximate Cost</div>
            <div class="kpi-value">{fmt_inr(full_df['total_cost_usd'].sum())}</div>
            <div class="kpi-sub">{len(FEATURE_REGISTRY)} features &nbsp;&bull;&nbsp; ${full_df['total_cost_usd'].sum():.4f} USD</div>
        </div>""", unsafe_allow_html=True)

    st.markdown("<div class='block-spacer'></div>", unsafe_allow_html=True)
    tab1, tab2 = st.tabs(["By Feature", "By Client"])

    with tab1:
        feat_df = full_df.groupby(["feature", "model", "ticket_filter"]).agg(
            input_tokens=("input_tokens",     "sum"),
            output_tokens=("output_tokens",   "sum"),
            total_tokens=("total_tokens",     "sum"),
            input_cost_inr=("input_cost_inr", "sum"),
            output_cost_inr=("output_cost_inr","sum"),
            total_cost_inr=("total_cost_inr", "sum"),
            total_cost_usd=("total_cost_usd", "sum"),
            client_count=("client_id",        "nunique"),
        ).reset_index().sort_values("total_cost_inr", ascending=False)

        st.markdown("""
        <div class="col-hdr">
            <div>Feature</div>
            <div>Tokens &nbsp;(Total / Input / Output)</div>
            <div>Input Cost</div>
            <div>Output Cost</div>
            <div style="text-align:right">Total Cost</div>
        </div>""", unsafe_allow_html=True)

        for _, row in feat_df.iterrows():
            pct = row["total_tokens"] / total_tokens * 100 if total_tokens else 0
            tfilter = row["ticket_filter"]
            type_badge = ""
            if tfilter == "email":
                type_badge = '<span class="type-badge email">email</span>'
            elif tfilter == "chat":
                type_badge = '<span class="type-badge">chat</span>'

            st.markdown(f"""
            <div class="data-row">
                <div>
                    <div class="row-name">{row['feature']}</div>
                    <div class="row-meta">
                        <span class="mbadge">{row['model']}</span>{type_badge}
                        &nbsp; {row['client_count']} client(s) &nbsp;&bull;&nbsp; {pct:.1f}% of tokens
                    </div>
                </div>
                <div class="token-val">
                    {row['total_tokens']:,}<br>
                    <span style="color:var(--muted2);font-size:0.7rem">
                        IN: {row['input_tokens']:,} | OUT: {row['output_tokens']:,}
                    </span>
                </div>
                <div class="cost-val">₹{row['input_cost_inr']:,.4f}</div>
                <div class="cost-val">₹{row['output_cost_inr']:,.4f}</div>
                <div class="cost-total">₹{row['total_cost_inr']:,.4f}</div>
            </div>""", unsafe_allow_html=True)

        st.markdown("<div class='block-spacer'></div>", unsafe_allow_html=True)
        st.markdown('<div class="sec-hdr">Cost Distribution by Feature</div>', unsafe_allow_html=True)
        chart_df = feat_df.set_index("feature")[["input_cost_inr", "output_cost_inr"]]
        chart_df.columns = ["Input Cost (INR)", "Output Cost (INR)"]
        st.bar_chart(chart_df, color=["#2563eb", "#0ea5e9"])

    with tab2:
        client_df = full_df.groupby("client_id").agg(
            input_tokens=("input_tokens",      "sum"),
            output_tokens=("output_tokens",    "sum"),
            total_tokens=("total_tokens",      "sum"),
            input_cost_inr=("input_cost_inr",  "sum"),
            output_cost_inr=("output_cost_inr","sum"),
            total_cost_inr=("total_cost_inr",  "sum"),
            features_used=("feature", lambda x: ", ".join(sorted(set(x)))),
        ).reset_index().sort_values("total_cost_inr", ascending=False)

        st.markdown("""
        <div class="col-hdr">
            <div>Client ID</div>
            <div>Tokens &nbsp;(Total / Input / Output)</div>
            <div>Input Cost</div>
            <div>Output Cost</div>
            <div style="text-align:right">Total Cost</div>
        </div>""", unsafe_allow_html=True)

        for _, row in client_df.iterrows():
            pct = row["total_tokens"] / total_tokens * 100 if total_tokens else 0
            st.markdown(f"""
            <div class="data-row">
                <div>
                    <div class="row-name" style="font-family:'IBM Plex Mono',monospace;font-size:0.84rem;">
                        {row['client_id']}
                    </div>
                    <div class="row-meta">
                        {row['features_used']} &nbsp;&bull;&nbsp; {pct:.1f}% of tokens
                    </div>
                </div>
                <div class="token-val">
                    {row['total_tokens']:,}<br>
                    <span style="color:var(--muted2);font-size:0.7rem">
                       IN: {row['input_tokens']:,} | OUT: {row['output_tokens']:,}
                    </span>
                </div>
                <div class="cost-val">₹{row['input_cost_inr']:,.4f}</div>
                <div class="cost-val">₹{row['output_cost_inr']:,.4f}</div>
                <div class="cost-total">₹{row['total_cost_inr']:,.4f}</div>
            </div>""", unsafe_allow_html=True)

        st.markdown("<div class='block-spacer'></div>", unsafe_allow_html=True)
        st.markdown('<div class="sec-hdr">Top 10 Clients by Cost</div>', unsafe_allow_html=True)
        top_c = client_df.head(10).set_index("client_id")[["input_cost_inr", "output_cost_inr"]]
        top_c.columns = ["Input Cost (INR)", "Output Cost (INR)"]
        st.bar_chart(top_c, color=["#16a34a", "#d97706"])

    st.markdown("<div class='block-spacer'></div>", unsafe_allow_html=True)
    st.markdown('<div class="sec-hdr">Export</div>', unsafe_allow_html=True)

    fname = f"token_usage_{start_date.strftime('%Y%m%d')}_to_{end_date.strftime('%Y%m%d')}.xlsx"
    excel_bytes = build_excel(full_df)
    st.download_button(
        label="⬇  Download",
        data=excel_bytes,
        file_name=fname,
        mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        use_container_width=False,
    )

else:
    st.markdown("""
    <div class="info-box">
        Set the <strong>date range</strong> and assign a <strong>model</strong> to each feature
        in the sidebar, then click <strong>Run Analysis</strong> to calculate token usage
        and estimated costs across all features in INR.
    </div>
    """, unsafe_allow_html=True)

    st.markdown('<div class="sec-hdr">Features Included in Analysis</div>', unsafe_allow_html=True)

    colors = ["c-blue", "c-sky", "c-green", "c-amber",
              "c-blue", "c-sky", "c-green", "c-amber", "c-blue"]
    cols = st.columns(4, gap="medium")

    for i, (coll, label, tfilter) in enumerate(FEATURE_REGISTRY):
        if tfilter == "email":
            type_tag_html = '<span class="type-badge email" style="margin-left:0">email</span>'
        elif tfilter == "chat":
            type_tag_html = '<span class="type-badge" style="margin-left:0">chat</span>'
        else:
            type_tag_html = ""

        default_model = FEATURE_DEFAULT_MODELS.get(label, "gpt-4o-mini")

        card_html = (
            f'<div class="kpi-card {colors[i]}" style="margin-bottom:1rem">'
            f'<div class="kpi-label">Feature</div>'
            f'<div style="font-size:0.855rem;font-weight:600;color:var(--text);'
            f'margin-bottom:0.5rem;line-height:1.3">{label}</div>'
            f'<div class="kpi-sub">'
            f'<span class="mbadge">{default_model}</span>'
            + (f'&nbsp;{type_tag_html}' if type_tag_html else "")
            + '</div>'
            '</div>'
        )

        with cols[i % 4]:
            st.markdown(card_html, unsafe_allow_html=True)