# %% [markdown]
# # Notebook 5: Developer Workflow Comparison
# Head-to-head comparison of the 2 developers. Cross-CLI comparison for ajaysingh.

# %%
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", 100)

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

# %% [markdown]
# ## 1 — Load all sessions and messages

# %%
all_sessions = sessions.list_all(conn, limit=500)
all_sessions["duration_min"] = (
    (all_sessions["last_updated"] - all_sessions["first_seen"])
    .dt.total_seconds() / 60
)

msg_per_session = query_df(conn, """
    SELECT m.session_id, s.user_email, s.source,
           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_calls,
           COUNT(*) FILTER (WHERE m.msg_type = 'tool_result') AS tool_results
    FROM   messages m
    JOIN   sessions s ON s.id = m.session_id
    GROUP  BY m.session_id, s.user_email, s.source
""")

sess_merged = all_sessions.merge(msg_per_session, left_on="id", right_on="session_id", suffixes=("", "_msg"))

# %% [markdown]
# ## 2 — Head-to-head session stats

# %%
display(Markdown("### Session-level comparison"))
comparison = sess_merged.groupby("user_email").agg(
    sessions=("id", "count"),
    avg_duration_min=("duration_min", "mean"),
    median_duration_min=("duration_min", "median"),
    avg_total_msgs=("total_msgs", "mean"),
    median_total_msgs=("total_msgs", "median"),
    avg_user_msgs=("user_msgs", "mean"),
    avg_tool_calls=("tool_calls", "mean"),
    avg_tool_results=("tool_results", "mean"),
).round(1)
display(comparison)

# %% [markdown]
# ## 3 — Tool preferences

# %%
display(Markdown("### Tool usage comparison"))
users_list = all_sessions["user_email"].unique()

tool_pref = []
for email in users_list:
    tf = tools.call_frequency(conn, email=email)
    tf["user_email"] = email
    tool_pref.append(tf)
tool_pref = pd.concat(tool_pref, ignore_index=True)

# Pivot: tool_name as rows, users as columns
pivot = tool_pref.pivot_table(
    index="tool_name", columns="user_email",
    values="call_count", fill_value=0
)
pivot["total"] = pivot.sum(axis=1)
pivot = pivot.sort_values("total", ascending=False)
display(pivot.head(20))

# %% [markdown]
# ## 4 — Read-heavy vs Edit-heavy

# %%
display(Markdown("### Exploration vs editing ratio per user"))
for email in users_list:
    tf = tools.call_frequency(conn, email=email)
    explore = tf[tf["tool_name"].isin(["Read", "Grep", "Glob", "View", "ReadFile", "read_file", "file_search", "grep_search", "codebase_search", "list_dir"])]["call_count"].sum()
    edit = tf[tf["tool_name"].isin(["Edit", "Write", "MultiEdit", "MultiEditTool"])]["call_count"].sum()
    shell = tf[tf["tool_name"].isin(["Bash", "shell_command", "exec_command"])]["call_count"].sum()
    total = tf["call_count"].sum()
    display(Markdown(f"**{email}**: explore={explore} ({explore/total*100:.0f}%), "
                     f"edit={edit} ({edit/total*100:.0f}%), "
                     f"shell={shell} ({shell/total*100:.0f}%), "
                     f"total={total}"))

# %% [markdown]
# ## 5 — Planning behavior (TodoWrite, Task, EnterPlanMode)

# %%
PLANNING_TOOLS = ["TodoWrite", "TaskCreate", "TaskUpdate", "TaskList", "Task", "EnterPlanMode"]

display(Markdown("### Planning tool usage per user"))
planning = tool_pref[tool_pref["tool_name"].isin(PLANNING_TOOLS)]
if not planning.empty:
    display(planning.pivot_table(
        index="tool_name", columns="user_email",
        values="call_count", fill_value=0
    ))
else:
    display(Markdown("No planning tools used."))

# %% [markdown]
# ## 6 — Session structure: front-loaded context vs iterative

# %%
display(Markdown("### First message length vs total messages (front-loading signal)"))

first_user_msgs = query_df(conn, """
    WITH ranked AS (
        SELECT m.session_id, m.content, s.user_email,
               LENGTH(m.content) AS content_len,
               ROW_NUMBER() OVER (PARTITION BY m.session_id ORDER BY m.timestamp) AS rn
        FROM   messages m
        JOIN   sessions s ON s.id = m.session_id
        WHERE  m.msg_type = 'user'
    )
    SELECT session_id, user_email, content_len AS first_msg_len
    FROM   ranked WHERE rn = 1
""")

sess_with_first = sess_merged.merge(first_user_msgs, left_on="id", right_on="session_id", suffixes=("", "_first"))

display(Markdown("### Avg first message length vs avg session messages"))
display(sess_with_first.groupby("user_email").agg(
    avg_first_msg_chars=("first_msg_len", "mean"),
    avg_user_msgs=("user_msgs", "mean"),
    avg_total_msgs=("total_msgs", "mean"),
).round(0))

# %% [markdown]
# ## 7 — Exploration before first edit

# %%
display(Markdown("### Tool calls before first edit per session"))

tool_sequence = query_df(conn, """
    SELECT tc.tool_name, m.session_id, m.timestamp, s.user_email
    FROM   tool_calls tc
    JOIN   messages m ON m.id = tc.message_id
    JOIN   sessions s ON s.id = m.session_id
    ORDER  BY m.session_id, m.timestamp
""")

EDIT_SET = {"Edit", "Write", "MultiEdit", "MultiEditTool"}

pre_edit_counts = []
for sid, group in tool_sequence.groupby("session_id"):
    email = group["user_email"].iloc[0]
    first_edit_idx = None
    for i, (_, row) in enumerate(group.iterrows()):
        if row["tool_name"] in EDIT_SET:
            first_edit_idx = i
            break
    if first_edit_idx is not None and first_edit_idx > 0:
        pre_edit = group.iloc[:first_edit_idx]
        pre_edit_counts.append({
            "session_id": sid,
            "user_email": email,
            "calls_before_first_edit": first_edit_idx,
            "reads_before_edit": (pre_edit["tool_name"].isin(["Read", "ReadFile", "read_file"])).sum(),
            "greps_before_edit": (pre_edit["tool_name"].isin(["Grep", "grep_search", "codebase_search"])).sum(),
        })

if pre_edit_counts:
    pre_edit_df = pd.DataFrame(pre_edit_counts)
    display(pre_edit_df.groupby("user_email")[["calls_before_first_edit", "reads_before_edit", "greps_before_edit"]].describe().round(1))

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

# %%
ajay_mask = all_sessions["user_email"].str.contains("ajay", case=False, na=False)
ajay_sessions = all_sessions[ajay_mask]

if ajay_sessions["source"].nunique() > 1:
    ajay_email = ajay_sessions["user_email"].iloc[0]
    display(Markdown(f"### {ajay_email} — behavior across CLIs"))

    ajay_merged = sess_merged[sess_merged["user_email"] == ajay_email]
    display(ajay_merged.groupby("source").agg(
        sessions=("id", "count"),
        avg_duration_min=("duration_min", "mean"),
        avg_total_msgs=("total_msgs", "mean"),
        avg_user_msgs=("user_msgs", "mean"),
        avg_tool_calls=("tool_calls", "mean"),
    ).round(1))

    # Tool frequency per CLI
    for src in ajay_sessions["source"].unique():
        src_sids = ajay_sessions[ajay_sessions["source"] == src]["id"].tolist()
        display(Markdown(f"#### Top tools in {src}"))
        tf_src = query_df(conn, """
            SELECT tc.tool_name, COUNT(*) AS call_count
            FROM   tool_calls tc
            JOIN   messages m ON m.id = tc.message_id
            WHERE  m.session_id = ANY(%s)
            GROUP  BY tc.tool_name
            ORDER  BY call_count DESC
            LIMIT  10
        """, (src_sids,))
        display(tf_src)

# %% [markdown]
# ## 9 — Developer archetype classification

# %%
display(Markdown("### Developer archetype signals"))

for email in users_list:
    mask = sess_merged["user_email"] == email
    user_sess = sess_merged[mask]
    avg_user_msgs = user_sess["user_msgs"].mean()
    avg_tool_calls = user_sess["tool_calls"].mean()
    avg_duration = user_sess["duration_min"].mean()
    short_sessions = (user_sess["user_msgs"] <= 3).mean() * 100

    display(Markdown(f"#### {email}"))
    signals = []
    if avg_user_msgs <= 3 and avg_tool_calls > 10:
        signals.append("One-shotter (few msgs, many tool calls)")
    if avg_user_msgs > 10:
        signals.append("Iterator (many back-and-forth messages)")
    if short_sessions > 50:
        signals.append(f"Many short sessions ({short_sessions:.0f}% have <=3 user msgs)")

    tf = tools.call_frequency(conn, email=email)
    explore_count = tf[tf["tool_name"].isin(["Read", "Grep", "Glob"])]["call_count"].sum()
    edit_count = tf[tf["tool_name"].isin(["Edit", "Write"])]["call_count"].sum()
    if explore_count > edit_count * 2:
        signals.append(f"Explorer (explore:edit = {explore_count}:{edit_count})")

    display(Markdown(f"avg_user_msgs={avg_user_msgs:.1f}, avg_tool_calls={avg_tool_calls:.1f}, "
                     f"avg_duration={avg_duration:.1f}min, short_sessions={short_sessions:.0f}%"))
    for s in signals:
        print(f"  → {s}")

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