# %% [markdown]
# # Notebook 1: Quantitative Baselines
# Session counts, durations, message distributions, token usage, and tool frequency.
# Establishes the numbers everything else builds on.

# %%
import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import pandas as pd
import numpy as np
from utils.connection import init, query_df
from utils import sessions, messages, tokens, tools
from IPython.display import display, Markdown

pd.set_option("display.max_columns", 30)
pd.set_option("display.max_rows", 100)
pd.set_option("display.max_colwidth", 80)

conn = init()
display(Markdown("**Connected.**"))

# %% [markdown]
# ## 1 — High-level counts

# %%
display(Markdown("### Users"))
users = sessions.list_users(conn)
display(users)

display(Markdown("### Orgs"))
orgs = sessions.list_orgs(conn)
display(orgs)

display(Markdown("### Repos"))
repos = sessions.list_repos(conn)
display(repos)

# %%
display(Markdown("### Message type distribution"))
msg_counts = messages.count_by_type(conn)
display(msg_counts)

# %% [markdown]
# ## 2 — Sessions by user and CLI source

# %%
all_sessions = sessions.list_all(conn, limit=500)
display(Markdown(f"**Total sessions:** {len(all_sessions)}"))

# Sessions per user
display(Markdown("### Sessions per user"))
display(all_sessions.groupby("user_email").agg(
    session_count=("id", "count"),
    sources=("source", lambda x: ", ".join(sorted(x.unique()))),
    repos=("repo_name", lambda x: ", ".join(sorted(x.dropna().unique()))),
).reset_index())

# Sessions per CLI source
display(Markdown("### Sessions per CLI source"))
display(all_sessions.groupby("source").agg(
    session_count=("id", "count"),
    users=("user_email", "nunique"),
).reset_index().sort_values("session_count", ascending=False))

# %% [markdown]
# ## 3 — Session duration distribution

# %%
all_sessions["duration_min"] = (
    (all_sessions["last_updated"] - all_sessions["first_seen"])
    .dt.total_seconds() / 60
)

display(Markdown("### Duration stats (minutes)"))
display(all_sessions["duration_min"].describe().round(1).to_frame())

display(Markdown("### Duration by user"))
display(all_sessions.groupby("user_email")["duration_min"].describe().round(1))

display(Markdown("### Duration by source"))
display(all_sessions.groupby("source")["duration_min"].describe().round(1))

# %% [markdown]
# ## 4 — Messages per session distribution

# %%
msg_per_session = query_df(conn, """
    SELECT m.session_id, s.user_email, s.source, s.repo_name,
           COUNT(*) AS total_msgs,
           COUNT(*) FILTER (WHERE m.msg_type = 'user') AS user_msgs,
           COUNT(*) FILTER (WHERE m.msg_type = 'assistant') AS assistant_msgs,
           COUNT(*) FILTER (WHERE m.msg_type = 'tool_call') AS tool_call_msgs,
           COUNT(*) FILTER (WHERE m.msg_type = 'tool_result') AS tool_result_msgs
    FROM   messages m
    JOIN   sessions s ON s.id = m.session_id
    GROUP  BY m.session_id, s.user_email, s.source, s.repo_name
    ORDER  BY total_msgs DESC
""")

display(Markdown("### Messages per session — percentiles"))
display(msg_per_session[["total_msgs", "user_msgs", "assistant_msgs", "tool_call_msgs", "tool_result_msgs"]].describe().round(1))

display(Markdown("### Messages per session — by user"))
display(msg_per_session.groupby("user_email")[["total_msgs", "user_msgs", "tool_call_msgs"]].describe().round(1))

# %% [markdown]
# ## 5 — Token usage

# %%
display(Markdown("### Token totals by user"))
tok_by_user = tokens.summary_by_user(conn)
display(tok_by_user)

display(Markdown("### Token totals by model"))
tok_by_model = tokens.summary_by_model(conn)
display(tok_by_model)

display(Markdown("### Token totals by day"))
tok_by_day = tokens.summary_by_day(conn)
display(tok_by_day)

# %% [markdown]
# ## 6 — Tool frequency (normalized)

# %%
# Overall tool frequency
display(Markdown("### Overall tool frequency"))
tf = tools.call_frequency(conn)
display(tf)

# Per-user tool frequency
display(Markdown("### Tool frequency — per user"))
for _, row in users.iterrows():
    email = row["user_email"]
    display(Markdown(f"#### {email}"))
    tf_user = tools.call_frequency(conn, email=email)
    display(tf_user.head(15))

# %% [markdown]
# ## 7 — Tool normalization: shell_command / Bash / exec_command

# %%
# These are equivalent across CLIs — show the mapping
shell_tools = tf[tf["tool_name"].isin(["shell_command", "Bash", "exec_command"])]
display(Markdown("### Shell tool variants (same concept, different CLIs)"))
display(shell_tools)

# Normalized view
display(Markdown("### Normalized tool frequency (shell variants merged)"))
tf_norm = tf.copy()
tf_norm["tool_name"] = tf_norm["tool_name"].replace({
    "shell_command": "Shell (normalized)",
    "Bash": "Shell (normalized)",
    "exec_command": "Shell (normalized)",
})
tf_norm = tf_norm.groupby("tool_name").agg(
    call_count=("call_count", "sum"),
    session_count=("session_count", "sum"),
).reset_index().sort_values("call_count", ascending=False)
display(tf_norm.head(15))

# %% [markdown]
# ## 8 — Tool failure rates

# %%
display(Markdown("### Tool failure rates"))
fr = tools.failure_rate(conn)
fr["failure_pct"] = (fr["failures"] / fr["total"] * 100).round(1)
display(fr.sort_values("failure_pct", ascending=False))

# %% [markdown]
# ## 9 — ajaysingh cross-CLI comparison

# %%
ajay_sessions = all_sessions[all_sessions["user_email"].str.contains("ajay", case=False, na=False)]
if not ajay_sessions.empty:
    ajay_email = ajay_sessions["user_email"].iloc[0]
    display(Markdown(f"### {ajay_email} — sessions by CLI"))
    display(ajay_sessions.groupby("source").agg(
        sessions=("id", "count"),
        avg_duration_min=("duration_min", "mean"),
        repos=("repo_name", lambda x: ", ".join(sorted(x.dropna().unique()))),
    ).reset_index().round(1))

    # Tool frequency per CLI for ajay
    for src in ajay_sessions["source"].unique():
        src_session_ids = ajay_sessions[ajay_sessions["source"] == src]["id"].tolist()
        display(Markdown(f"#### Tools used in {src}"))
        for sid in src_session_ids[:3]:
            tf_s = tools.call_frequency(conn, session_id=sid)
            if not tf_s.empty:
                display(Markdown(f"Session `{sid[:12]}...`"))
                display(tf_s.head(10))

# %% [markdown]
# ---
# *End of Notebook 1.*
