# %% [markdown]
# # Notebook 6: Session Resolution Estimation
# Classify sessions as completed / abandoned / errored using proxy signals.

# %%
import sys, os, re
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, 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", 120)

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

# %% [markdown]
# ## 1 — Load session data with signal features

# %%
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_counts = query_df(conn, """
    SELECT m.session_id,
           COUNT(*) AS total_msgs,
           COUNT(*) FILTER (WHERE m.msg_type = 'user') AS user_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
    GROUP  BY m.session_id
""")

all_sessions = all_sessions.merge(msg_counts, left_on="id", right_on="session_id", how="left")

# %% [markdown]
# ## 2 — Signal: git commit in shell tool calls

# %%
# Check for git commit in shell-like tool calls
git_commits = query_df(conn, """
    SELECT DISTINCT m.session_id
    FROM   tool_calls tc
    JOIN   messages m ON m.id = tc.message_id
    WHERE  tc.tool_name IN ('Bash', 'shell_command', 'exec_command')
    AND    (
        tc.tool_input::text ILIKE %s
        OR tc.tool_input::text ILIKE %s
    )
""", ('%git commit%', '%git push%'))

all_sessions["has_git_commit"] = all_sessions["id"].isin(git_commits["session_id"].tolist())
display(Markdown(f"**Sessions with git commit/push:** {all_sessions['has_git_commit'].sum()} / {len(all_sessions)}"))

# %% [markdown]
# ## 3 — Signal: last message type

# %%
last_msg_type = query_df(conn, """
    WITH ranked AS (
        SELECT m.session_id, m.msg_type,
               ROW_NUMBER() OVER (PARTITION BY m.session_id ORDER BY m.timestamp DESC) AS rn
        FROM messages m
    )
    SELECT session_id, msg_type AS last_msg_type FROM ranked WHERE rn = 1
""")

all_sessions = all_sessions.merge(last_msg_type, left_on="id", right_on="session_id", how="left", suffixes=("", "_last"))

display(Markdown("### Last message type distribution"))
display(all_sessions["last_msg_type"].value_counts().to_frame("count"))

# %% [markdown]
# ## 4 — Signal: tool failure rate per session

# %%
session_failures = query_df(conn, """
    SELECT m.session_id,
           COUNT(*) AS total_results,
           COUNT(*) FILTER (WHERE tr.status = 'failure') AS failures
    FROM   tool_results tr
    JOIN   messages m ON m.id = tr.message_id
    GROUP  BY m.session_id
""")
session_failures["failure_rate"] = (session_failures["failures"] / session_failures["total_results"]).round(3)

all_sessions = all_sessions.merge(session_failures, left_on="id", right_on="session_id", how="left", suffixes=("", "_fail"))
all_sessions["failure_rate"] = all_sessions["failure_rate"].fillna(0)

display(Markdown("### Tool failure rate distribution"))
display(all_sessions["failure_rate"].describe().round(3).to_frame())

# %% [markdown]
# ## 5 — Signal: repeated edits to same file (struggle)

# %%
repeated_edits = query_df(conn, """
    SELECT m.session_id, COUNT(*) AS repeat_edit_files
    FROM (
        SELECT m.session_id,
               tc.tool_input->>'file_path' AS file_path,
               COUNT(*) AS edit_count
        FROM   tool_calls tc
        JOIN   messages m ON m.id = tc.message_id
        WHERE  tc.tool_name IN ('Edit', 'Write', 'MultiEdit')
        AND    tc.tool_input->>'file_path' IS NOT NULL
        GROUP  BY m.session_id, tc.tool_input->>'file_path'
        HAVING COUNT(*) >= 5
    ) sub
    JOIN messages m ON m.session_id = sub.session_id
    GROUP BY m.session_id
""")

all_sessions["has_repeated_edits"] = all_sessions["id"].isin(
    repeated_edits["session_id"].tolist() if not repeated_edits.empty else []
)

# %% [markdown]
# ## 6 — Classify sessions

# %%
def classify_session(row):
    """Heuristic session outcome classification."""
    # Abandoned: very few messages, no tool calls
    if row.get("total_msgs", 0) <= 2 and row.get("tool_calls", 0) == 0:
        return "abandoned"

    # Errored: high failure rate
    if row.get("failure_rate", 0) > 0.3 and row.get("total_results", 0) > 5:
        return "errored"

    # Completed: has git commit
    if row.get("has_git_commit", False):
        return "completed"

    # Completed: reasonable session with assistant ending
    if row.get("last_msg_type") == "assistant" and row.get("tool_calls", 0) > 3:
        return "likely_completed"

    # Short but functional
    if row.get("tool_calls", 0) > 0 and row.get("user_msgs", 0) >= 1:
        return "likely_completed"

    return "unclear"

all_sessions["outcome"] = all_sessions.apply(classify_session, axis=1)

display(Markdown("### Session outcome distribution"))
display(all_sessions["outcome"].value_counts().to_frame("count"))

display(Markdown("### Outcome by user"))
display(pd.crosstab(all_sessions["user_email"], all_sessions["outcome"]))

display(Markdown("### Outcome by source"))
display(pd.crosstab(all_sessions["source"], all_sessions["outcome"]))

# %% [markdown]
# ## 7 — Per-user estimated success rate

# %%
display(Markdown("### Estimated success rate per user"))
for email in all_sessions["user_email"].unique():
    mask = all_sessions["user_email"] == email
    user_sess = all_sessions[mask]
    n = len(user_sess)
    completed = (user_sess["outcome"].isin(["completed", "likely_completed"])).sum()
    display(Markdown(f"**{email}**: {completed}/{n} ({completed/n*100:.0f}%) estimated successful"))

# %% [markdown]
# ## 8 — Correlation: session duration vs outcome

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

display(Markdown("### Tool calls by outcome"))
display(all_sessions.groupby("outcome")["tool_calls"].describe().round(1))

# %% [markdown]
# ## 9 — Correlation: failure rate vs outcome

# %%
display(Markdown("### Failure rate by outcome"))
display(all_sessions.groupby("outcome")["failure_rate"].describe().round(3))

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