# %% [markdown]
# # Claude Code Compacted Sessions — Parent/Child Analysis
#
# Identifies compacted Claude Code sessions and maps parent→child relationships.
#
# **Detection:** A compaction event is a `user` message containing
# "This session is being continued from a previous conversation".
#
# **Two scenarios:**
# 1. In-place compaction (`/compact`) — summary injected mid-session, same session ID
# 2. Resume compaction (`--resume <id>`) — new session created, first message is the summary
#
# If the compaction message is the **first** message in a session, it's a child
# that was resumed from a parent. We find the parent by looking at the user's
# prior session that ended just before this one started.

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

from utils.connection import init, query_df
import pandas as pd

pd.set_option("display.max_colwidth", 300)
pd.set_option("display.max_columns", None)
pd.set_option("display.width", 240)
conn = init()
print("Connected.\n")

ORG_FILTER = "sagar-%"

# %% [markdown]
# ## 1 — Find all compaction events

# %%
compaction_msgs = query_df(conn, """
    SELECT m.session_id,
           m.id AS message_id,
           m.msg_type,
           m.timestamp AS compaction_ts,
           LEFT(m.content, 1000) AS content_preview,
           s.user_email,
           s.repo_name,
           s.first_seen,
           s.last_updated
    FROM   messages m
    JOIN   sessions s ON s.id = m.session_id
    WHERE  s.source = 'claude_code'
      AND  s.org LIKE %s
      AND  m.content ILIKE '%%This session is being continued from a previous conversation%%'
    ORDER  BY m.timestamp ASC
""", (ORG_FILTER,))

print(f"Total compaction events: {len(compaction_msgs)}")
print(f"Unique sessions with compaction: {compaction_msgs['session_id'].nunique()}\n")

# %% [markdown]
# ## 2 — Classify: in-place vs resumed (child) sessions
#
# Check if the compaction message is the first message in the session.

# %%
# Get the first message timestamp for each session that has a compaction event
session_ids = compaction_msgs["session_id"].unique().tolist()

first_msgs = query_df(conn, """
    WITH ranked AS (
        SELECT session_id, id AS message_id, msg_type, timestamp,
               ROW_NUMBER() OVER (PARTITION BY session_id ORDER BY timestamp ASC) AS rn
        FROM   messages
        WHERE  session_id = ANY(%s)
    )
    SELECT session_id, message_id AS first_msg_id, msg_type AS first_msg_type, timestamp AS first_msg_ts
    FROM   ranked
    WHERE  rn = 1
""", (session_ids,))

# Merge to check if compaction msg is the first message
compaction_analysis = compaction_msgs.merge(
    first_msgs, on="session_id", how="left"
)

# A compaction is "resumed" (child session) if the compaction message is the first message
# We compare message_ids since timestamps might not be exact
compaction_analysis["is_first_message"] = (
    compaction_analysis["message_id"] == compaction_analysis["first_msg_id"]
)

# For in-place compactions, count how many times per session
compaction_counts = compaction_msgs.groupby("session_id").size().reset_index(name="compaction_count")

# Deduplicate to one row per session, keeping first compaction event
session_level = compaction_analysis.sort_values("compaction_ts").drop_duplicates(
    subset="session_id", keep="first"
).merge(compaction_counts, on="session_id")

session_level["type"] = session_level["is_first_message"].map(
    {True: "RESUMED (child)", False: "IN-PLACE (/compact)"}
)

print("=== Compacted sessions breakdown ===\n")
print(session_level["type"].value_counts().to_string())
print()

print("=== Per-session details ===\n")
display_cols = ["session_id", "type", "compaction_count", "user_email", "repo_name", "first_seen", "last_updated"]
print(session_level[display_cols].to_string(index=False))

# %% [markdown]
# ## 3 — Find parent sessions for resumed (child) sessions
#
# Strategy: For each child session, find the same user's most recent prior session
# on the same repo that ended before the child started.

# %%
children = session_level[session_level["type"] == "RESUMED (child)"].copy()
print(f"\n=== Resumed (child) sessions to find parents for: {len(children)} ===\n")

if not children.empty:
    # For each child, find parent = same user + same repo + ended just before child started
    parent_matches = []
    for _, child in children.iterrows():
        candidates = query_df(conn, """
            SELECT id AS parent_session_id,
                   user_email,
                   repo_name,
                   first_seen AS parent_first_seen,
                   last_updated AS parent_last_updated
            FROM   sessions
            WHERE  source = 'claude_code'
              AND  user_email = %s
              AND  (repo_name = %s OR (%s IS NULL AND repo_name IS NULL))
              AND  last_updated <= %s
              AND  id != %s
            ORDER  BY last_updated DESC
            LIMIT  1
        """, (child["user_email"], child["repo_name"], child["repo_name"],
              child["first_seen"], child["session_id"]))

        if not candidates.empty:
            row = candidates.iloc[0]
            parent_matches.append({
                "child_session_id": child["session_id"],
                "parent_session_id": row["parent_session_id"],
                "child_first_seen": child["first_seen"],
                "parent_last_updated": row["parent_last_updated"],
                "user_email": child["user_email"],
                "repo_name": child["repo_name"],
                "compaction_count": child["compaction_count"],
            })
        else:
            parent_matches.append({
                "child_session_id": child["session_id"],
                "parent_session_id": None,
                "child_first_seen": child["first_seen"],
                "parent_last_updated": None,
                "user_email": child["user_email"],
                "repo_name": child["repo_name"],
                "compaction_count": child["compaction_count"],
            })

    parents_df = pd.DataFrame(parent_matches)

    # Check if any parent is itself a child (multi-hop chains)
    parent_ids_found = parents_df["parent_session_id"].dropna().tolist()
    parents_that_are_children = set(parent_ids_found) & set(children["session_id"].tolist())

    print("=== Parent → Child Mapping ===\n")
    print(parents_df.to_string(index=False))
    print()

    found = parents_df["parent_session_id"].notna().sum()
    print(f"Parents found: {found} / {len(children)}")
    if parents_that_are_children:
        print(f"Multi-hop chains detected (parent is also a child): {parents_that_are_children}")

# %% [markdown]
# ## 4 — Build full session chains
#
# Walk backward from each child to build chains: grandparent → parent → child

# %%
if not children.empty and not parents_df.empty:
    # Build adjacency: child -> parent
    child_to_parent = dict(zip(
        parents_df["child_session_id"],
        parents_df["parent_session_id"]
    ))

    # Also check if any parent was itself resumed
    all_compacted_ids = set(session_level["session_id"].tolist())

    def build_chain(session_id):
        """Walk backward from session to root."""
        chain = [session_id]
        current = session_id
        seen = set()
        while current in child_to_parent and child_to_parent[current] and current not in seen:
            seen.add(current)
            current = child_to_parent[current]
            chain.append(current)
        chain.reverse()  # root first
        return chain

    # Build chains for all children
    chains = []
    for child_id in children["session_id"]:
        chain = build_chain(child_id)
        chains.append({
            "chain_length": len(chain),
            "root_session": chain[0],
            "chain": " → ".join(chain[-3:]) if len(chain) <= 3 else f"...({len(chain)-2} more) → " + " → ".join(chain[-2:]),
            "full_chain": " → ".join(chain),
            "leaf_session": chain[-1],
            "user_email": children[children["session_id"] == child_id].iloc[0]["user_email"],
            "repo_name": children[children["session_id"] == child_id].iloc[0]["repo_name"],
        })

    chains_df = pd.DataFrame(chains).sort_values("chain_length", ascending=False)

    print("\n=== Session Chains (root → ... → leaf) ===\n")
    print(chains_df[["chain_length", "user_email", "repo_name", "full_chain"]].to_string(index=False))

# %% [markdown]
# ## 5 — Summary statistics

# %%
print("\n=== Summary ===\n")
print(f"Total sessions with compaction events: {len(session_level)}")
print(f"  In-place (/compact):    {(session_level['type'] == 'IN-PLACE (/compact)').sum()}")
print(f"  Resumed (child):        {(session_level['type'] == 'RESUMED (child)').sum()}")
print(f"  Max compactions in one session: {session_level['compaction_count'].max()}")
print()

if not children.empty:
    print(f"Parent→child pairs found: {parents_df['parent_session_id'].notna().sum()}")
    print(f"Orphan children (no parent in DB): {parents_df['parent_session_id'].isna().sum()}")
    if len(chains_df):
        print(f"Longest chain: {chains_df['chain_length'].max()} sessions")
print()

# Per-user breakdown
per_user = session_level.groupby("user_email").agg(
    total_compacted=("session_id", "count"),
    in_place=("type", lambda x: (x == "IN-PLACE (/compact)").sum()),
    resumed=("type", lambda x: (x == "RESUMED (child)").sum()),
    total_compactions=("compaction_count", "sum"),
).reset_index()
print("=== Per-user breakdown ===\n")
print(per_user.to_string(index=False))

# %%
conn.close()
print("\nDone.")
