"""Explore the actual data shape: session index keys, message fields, source distribution.

Run: /home/sagar/trace/.venv/bin/python3 analysis-14022026/session_view_scripts/research/00_data_shape.py
"""
import json
import os
import statistics
from collections import Counter

DIR = os.path.join(os.path.dirname(__file__), "../viewer/public/data/sessions")
IDX = os.path.join(os.path.dirname(__file__), "../viewer/public/data/index.json")

with open(IDX) as f:
    index = json.load(f)

# Show index structure
devs = list(index["developers"].keys())
print("=== DEVELOPERS:", devs)
for dev in index["developers"].values():
    print("Session keys:", list(dev["sessions"][0].keys()))
    break

# Aggregate stats across ALL sessions
total_sessions = 0
source_counts = Counter()
total_msgs_global = 0
type_counts_global = Counter()
sessions_with_tools = 0
sessions_without_tools = 0
tool_name_counts = Counter()
user_msg_lengths = []
sessions_with_interruption = 0
result_status_counts = Counter()

for dev_name, dev in index["developers"].items():
    for s in dev["sessions"]:
        total_sessions += 1
        source_counts[s["source"]] += 1
        if s.get("interruptions", 0) > 0:
            sessions_with_interruption += 1

        fpath = os.path.join(DIR, f'{s["id"]}.json')
        if not os.path.exists(fpath):
            continue
        with open(fpath) as f2:
            msgs = json.load(f2)

        total_msgs_global += len(msgs)
        has_tools = False
        for m in msgs:
            type_counts_global[m["type"]] += 1
            if m["type"] == "user" and m.get("content"):
                user_msg_lengths.append(len(m["content"]))
            if m["type"] == "tool_call" and m.get("tool"):
                has_tools = True
                tool_name_counts[m["tool"]["tool_name"]] += 1
            if m["type"] == "tool_result" and m.get("result"):
                result_status_counts[m["result"]["status"]] += 1

        if has_tools:
            sessions_with_tools += 1
        else:
            sessions_without_tools += 1

print(f"\nTotal sessions: {total_sessions}")
print(f"Source distribution: {dict(source_counts)}")
print(f"Total messages: {total_msgs_global}")
print(f"Message type distribution: {dict(type_counts_global)}")
print(f"Sessions with tools: {sessions_with_tools}, without: {sessions_without_tools}")
print(f"Tool name distribution (top 25): {tool_name_counts.most_common(25)}")
print(f"Result status distribution: {dict(result_status_counts)}")
print(f"Sessions with interruptions: {sessions_with_interruption}")

if user_msg_lengths:
    sorted_lens = sorted(user_msg_lengths)
    print(
        f"User msg length: min={min(user_msg_lengths)}, "
        f"median={statistics.median(user_msg_lengths):.0f}, "
        f"mean={statistics.mean(user_msg_lengths):.0f}, "
        f"max={max(user_msg_lengths)}, "
        f"p90={sorted_lens[int(0.9 * len(sorted_lens))]:.0f}"
    )
    print(f"Total user messages: {len(user_msg_lengths)}")
