"""Export session data for the datewise viewer.

Exports per-developer sessions with message timestamps.
Uses message timestamps (not session first_seen) for accurate dates.

Output: session_viewer_data.json → used by the Vite viewer app.
"""
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from utils.connection import init, query_df

OUTPUT = Path(__file__).resolve().parent / "viewer" / "public" / "data"
OUTPUT.mkdir(parents=True, exist_ok=True)

import argparse
_parser = argparse.ArgumentParser()
_parser.add_argument("--prod", action="store_true", help="Use production database")
_args = _parser.parse_args()

conn = init(env_filename=".prod.env" if _args.prod else ".local.env")
conn.execute("SET client_encoding TO 'UTF8'")

# Get all sessions with start/end from message timestamps
sessions = query_df(conn, """
    SELECT s.id, s.user_email, s.source, s.model, s.repo_name, s.cwd,
           MIN(m.timestamp) as session_start,
           MAX(m.timestamp) as session_end,
           COUNT(*) as total_messages,
           COUNT(CASE WHEN m.msg_type = 'user' THEN 1 END) as user_messages,
           COUNT(CASE WHEN m.msg_type = 'assistant' THEN 1 END) as assistant_messages,
           COUNT(CASE WHEN m.msg_type = 'tool_call' THEN 1 END) as tool_calls,
           COUNT(CASE WHEN m.msg_type = 'tool_result' THEN 1 END) as tool_results,
           COUNT(CASE WHEN m.msg_type = 'user' AND m.content ILIKE 'This session is being continued from a previous conversation%%' THEN 1 END) as compactions,
           COUNT(CASE WHEN m.msg_type = 'user' AND m.content ILIKE '<turn_aborted>%%' THEN 1 END) as interruptions_codex
    FROM sessions s
    JOIN messages m ON m.session_id = s.id
    WHERE (s.org ILIKE 'pratilipi%%' OR s.org ILIKE 'sagar-%%' OR s.org ILIKE 'staple-ai%%' OR s.org ILIKE 'vertexcover%%')
      AND s.source NOT IN ('test', 'qc_trace_install')
      AND s.user_email IS NOT NULL
      AND m.timestamp IS NOT NULL
    GROUP BY s.id, s.user_email, s.source, s.model, s.repo_name, s.cwd
    ORDER BY MIN(m.timestamp)
""")

# Get all messages with timestamps (content fetched separately to handle encoding)
messages = query_df(conn, """
    WITH ranked AS (
        SELECT m.id, m.session_id, m.msg_type, m.timestamp, m.model, m.content,
               LAG(m.content) OVER (PARTITION BY m.session_id, m.msg_type ORDER BY m.timestamp, m.id) AS prev_content
        FROM messages m
        JOIN sessions s ON s.id = m.session_id
        WHERE (s.org ILIKE 'pratilipi%%' OR s.org ILIKE 'sagar-%%' OR s.org ILIKE 'staple-ai%%' OR s.org ILIKE 'vertexcover%%')
          AND s.source NOT IN ('test', 'qc_trace_install')
          AND s.user_email IS NOT NULL
          AND m.timestamp IS NOT NULL
    )
    SELECT id, session_id, msg_type, timestamp, model
    FROM ranked
    WHERE content IS DISTINCT FROM prev_content
       OR prev_content IS NULL
    ORDER BY timestamp
""")

# Fetch content row by row to handle encoding errors
print("Fetching message content...")
content_map = {}
with conn.cursor("content_cur") as cur:
    cur.itersize = 2000
    cur.execute("""
        SELECT m.id, m.content
        FROM messages m
        JOIN sessions s ON s.id = m.session_id
        WHERE (s.org ILIKE 'pratilipi%%' OR s.org ILIKE 'sagar-%%' OR s.org ILIKE 'staple-ai%%' OR s.org ILIKE 'vertexcover%%')
          AND s.source NOT IN ('test', 'qc_trace_install')
          AND s.user_email IS NOT NULL
          AND m.timestamp IS NOT NULL
    """)
    for row in cur:
        mid, content = row
        if content and isinstance(content, str):
            content_map[mid] = content
        else:
            content_map[mid] = ""

# Fetch tool_calls data
print("Fetching tool calls...")
tool_calls_df = query_df(conn, """
    SELECT tc.message_id, tc.tool_name, tc.tool_input
    FROM tool_calls tc
    JOIN messages m ON m.id = tc.message_id
    JOIN sessions s ON s.id = m.session_id
    WHERE (s.org ILIKE 'pratilipi%%' OR s.org ILIKE 'sagar-%%' OR s.org ILIKE 'staple-ai%%' OR s.org ILIKE 'vertexcover%%')
      AND s.source NOT IN ('test', 'qc_trace_install')
      AND s.user_email IS NOT NULL
""")
tool_call_map = {}
for _, row in tool_calls_df.iterrows():
    ti = row["tool_input"]
    if isinstance(ti, dict):
        ti_str = json.dumps(ti, indent=2, default=str)
    elif ti:
        ti_str = str(ti)
    else:
        ti_str = ""
    tool_call_map[row["message_id"]] = {
        "tool_name": row["tool_name"],
        "tool_input": ti_str,
    }

# Fetch tool_results data
print("Fetching tool results...")
tool_results_map = {}
with conn.cursor("tr_cur") as cur:
    cur.itersize = 2000
    cur.execute("""
        SELECT tr.message_id, tr.status, tr.output
        FROM tool_results tr
        JOIN messages m ON m.id = tr.message_id
        JOIN sessions s ON s.id = m.session_id
        WHERE (s.org ILIKE 'pratilipi%%' OR s.org ILIKE 'sagar-%%' OR s.org ILIKE 'staple-ai%%' OR s.org ILIKE 'vertexcover%%')
          AND s.source NOT IN ('test', 'qc_trace_install')
          AND s.user_email IS NOT NULL
    """)
    for row in cur:
        mid, status, output = row
        tool_results_map[mid] = {
            "status": status or "",
            "output": (output[:2000] if output and isinstance(output, str) else "") ,
        }

conn.close()

# Count Claude Code interruptions per session from tool_results
CLAUDE_INTERRUPT_SIG = "The user doesn't want to proceed with this tool use"
interruptions_claude = {}
for _, row in messages.iterrows():
    mid = row["id"]
    if row["msg_type"] == "tool_result" and mid in tool_results_map:
        output = tool_results_map[mid].get("output", "")
        if output and CLAUDE_INTERRUPT_SIG in output:
            sid = row["session_id"]
            interruptions_claude[sid] = interruptions_claude.get(sid, 0) + 1

# Build the data structure
developers = {}
for _, row in sessions.iterrows():
    email = row["user_email"]
    name = email.split("@")[0]
    if name not in developers:
        developers[name] = {"email": email, "sessions": []}

    developers[name]["sessions"].append({
        "id": row["id"],
        "source": row["source"],
        "model": row["model"],
        "repo_name": row["repo_name"],
        "cwd": row["cwd"],
        "start": str(row["session_start"]),
        "end": str(row["session_end"]),
        "total_messages": int(row["total_messages"]),
        "user_messages": int(row["user_messages"]),
        "assistant_messages": int(row["assistant_messages"]),
        "tool_calls": int(row["tool_calls"]),
        "tool_results": int(row["tool_results"]),
        "compactions": int(row["compactions"]),
        "interruptions": int(row["interruptions_codex"]) + interruptions_claude.get(row["id"], 0),
    })

# Build messages by session
messages_by_session = {}
for _, row in messages.iterrows():
    sid = row["session_id"]
    if sid not in messages_by_session:
        messages_by_session[sid] = []
    msg = {
        "id": row["id"],
        "type": row["msg_type"],
        "timestamp": str(row["timestamp"]),
        "model": row["model"],
        "content": content_map.get(row["id"], ""),
    }
    mid = row["id"]
    if mid in tool_call_map:
        msg["tool"] = tool_call_map[mid]
    if mid in tool_results_map:
        msg["result"] = tool_results_map[mid]
    messages_by_session[sid].append(msg)

import math

def clean(obj):
    if isinstance(obj, float) and (math.isnan(obj) or math.isinf(obj)):
        return None
    if isinstance(obj, dict):
        return {k: clean(v) for k, v in obj.items()}
    if isinstance(obj, list):
        return [clean(v) for v in obj]
    return obj

# Write index file (developers + session metadata only, no messages)
index_data = clean({"developers": developers})
index_file = OUTPUT / "index.json"
index_file.write_text(json.dumps(index_data, default=str))
print(f"Wrote {index_file} ({index_file.stat().st_size / 1024:.0f} KB)")

# Write per-session message files (lazy loaded)
sessions_dir = OUTPUT / "sessions"
sessions_dir.mkdir(parents=True, exist_ok=True)
for sid, msgs in messages_by_session.items():
    msgs = clean(msgs)
    sf = sessions_dir / f"{sid}.json"
    sf.write_text(json.dumps(msgs, default=str))

print(f"Wrote {len(messages_by_session)} session files to {sessions_dir}/")
print(f"  {len(developers)} developers, {len(sessions)} sessions, {len(messages)} messages")
