"""Iteration 5 enhancements: developer summaries, session arc per-dev, token waste analysis.

Reads existing JSON files, enriches them, and writes back.
"""
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd

sys.path.insert(0, "/home/sagar/trace/analysis-14022026")
from utils.connection import init, query_df

OUTPUT_DIR = Path("/home/sagar/trace/analysis-14022026/output")
DASHBOARD_DIR = Path("/home/sagar/trace/analysis-14022026/dashboard/public/data")

def load(name):
    with open(OUTPUT_DIR / name) as f:
        return json.load(f)

def save(name, data):
    text = json.dumps(data, indent=2, default=str)
    for p in [OUTPUT_DIR / name, DASHBOARD_DIR / name]:
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(text)
    print(f"  wrote {name}")

conn = init()

# ---------------------------------------------------------------------------
# 1. DEVELOPER SUMMARIES — one-sentence narrative per developer
# ---------------------------------------------------------------------------
print("Generating developer summaries...")

profiles = load("developer_profiles.json")
overview = load("overview.json")
total_cost = overview.get("estimated_cost_usd", 1)

for dev in profiles["developers"]:
    email = dev["email"]
    name = email.split("@")[0].replace("+", " ").split(" ")[-1]
    cost = dev["estimated_cost_usd"]
    cost_pct = round(cost / max(total_cost, 1) * 100, 1)
    sessions = dev["total_sessions"]
    commits = dev.get("commit_count", 0)
    persona = dev["persona"]
    eff = dev.get("avg_efficiency_score", 0)
    primary_intent = dev.get("primary_intent", "unknown")
    edit_calls = dev["total_edit_calls"]
    explore_calls = dev["total_explore_calls"]
    hours = dev["total_hours"]

    # Build narrative
    parts = []
    parts.append(f"{name} is a {persona.lower()}")

    if cost_pct > 50:
        parts.append(f"who consumed {cost_pct}% of total AI spend (${cost:,.0f})")
    else:
        parts.append(f"who spent ${cost:,.0f} ({cost_pct}% of budget)")

    if commits == 0:
        parts.append("with zero commits detected via CLI")
    elif commits > 0:
        parts.append(f"shipping {commits} commit{'s' if commits != 1 else ''} at ${dev.get('cost_per_commit', 0):,.0f}/commit")

    if eff < 25:
        parts.append(f"— average efficiency score of just {eff}/100")
    elif eff >= 50:
        parts.append(f"with a solid {eff}/100 efficiency score")

    if primary_intent == "understand":
        parts.append(". Primarily uses AI to understand code rather than build.")
    elif primary_intent == "build_feature":
        parts.append(". Primarily uses AI to build new features.")
    elif primary_intent == "fix_bug":
        parts.append(". Primarily uses AI for bug fixing and debugging.")
    else:
        parts.append(".")

    dev["developer_summary"] = " ".join(parts)

save("developer_profiles.json", profiles)

# ---------------------------------------------------------------------------
# 2. PER-DEVELOPER SESSION ARC BREAKDOWN
# ---------------------------------------------------------------------------
print("Computing per-developer session arcs...")

session_ids_sql = """
    SELECT id FROM sessions
    WHERE source NOT IN ('test', 'qc_trace_install')
    AND user_email IS NOT NULL
    AND org ILIKE 'pratilipi%%'
"""

TOOL_CATEGORY = {
    "Bash": "shell", "shell_command": "shell", "exec_command": "shell",
    "Read": "explore", "Grep": "explore", "Glob": "explore", "View": "explore",
    "ReadFile": "explore", "read_file": "explore", "grep_search": "explore",
    "codebase_search": "explore", "list_dir": "explore", "file_search": "explore",
    "Edit": "edit", "Write": "edit", "MultiEdit": "edit", "MultiEditTool": "edit",
    "TodoRead": "planning", "TodoWrite": "planning",
    "Task": "delegation", "WebFetch": "web", "WebSearch": "web",
}

workflow_df = query_df(conn, f"""
    SELECT m.session_id, tc.tool_name, 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
    WHERE m.session_id IN ({session_ids_sql})
    ORDER BY m.session_id, m.timestamp
""")
workflow_df["category"] = workflow_df["tool_name"].map(lambda x: TOOL_CATEGORY.get(x, "other"))

# Classify arcs per session
def classify_arc(seq):
    if len(seq) == 0:
        return "empty"
    has_edit = "edit" in seq
    has_shell = "shell" in seq
    has_explore = "explore" in seq
    if len(seq) <= 2 and not has_edit:
        return "quick_lookup"
    if has_explore and not has_edit and not has_shell:
        return "explore_only"
    if has_explore and has_edit and not has_shell:
        return "explore_then_edit"
    if has_edit and has_shell:
        return "full_cycle"
    if has_shell and not has_edit:
        return "shell_only"
    if has_edit and not has_explore:
        return "direct_edit"
    return "mixed"

session_arcs = workflow_df.groupby(["session_id", "user_email"])["category"].apply(list).reset_index()
session_arcs["arc"] = session_arcs["category"].apply(classify_arc)

# Per-developer arc breakdown
dev_arcs = session_arcs.groupby(["user_email", "arc"]).size().unstack(fill_value=0)
dev_arcs_dict = {}
for email in dev_arcs.index:
    dev_arcs_dict[email] = dev_arcs.loc[email].to_dict()

for dev in profiles["developers"]:
    dev["session_arcs"] = dev_arcs_dict.get(dev["email"], {})

save("developer_profiles.json", profiles)

# ---------------------------------------------------------------------------
# 3. TOKEN WASTE ANALYSIS — tokens spent on sessions with no output
# ---------------------------------------------------------------------------
print("Computing token waste analysis...")

sess_tokens = query_df(conn, f"""
    SELECT m.session_id, s.user_email,
           SUM(t.input_tokens) as input_tokens,
           SUM(t.output_tokens) as output_tokens
    FROM token_usage t
    JOIN messages m ON m.id = t.message_id
    JOIN sessions s ON s.id = m.session_id
    WHERE s.source NOT IN ('test', 'qc_trace_install')
    AND s.user_email IS NOT NULL
    AND s.org ILIKE 'pratilipi%%'
    GROUP BY m.session_id, s.user_email
""")

# Check which sessions had edits
edit_sessions = query_df(conn, f"""
    SELECT DISTINCT m.session_id
    FROM tool_calls tc JOIN messages m ON m.id = tc.message_id
    WHERE tc.tool_name IN ('Edit', 'Write', 'MultiEdit', 'MultiEditTool')
    AND m.session_id IN ({session_ids_sql})
""")

sess_tokens["had_edits"] = sess_tokens["session_id"].isin(edit_sessions["session_id"])
from utils.pricing import calc_cost as _calc_cost

# Get dominant model per session for cost calculation
sess_model = query_df(conn, f"""
    SELECT m.session_id, m.model, SUM(t.input_tokens) as tok
    FROM messages m
    JOIN token_usage t ON t.message_id = m.id
    JOIN sessions s ON s.id = m.session_id
    WHERE s.source NOT IN ('test', 'qc_trace_install')
    AND s.user_email IS NOT NULL AND s.org ILIKE 'pratilipi%%'
    GROUP BY m.session_id, m.model
""")
# Pick the model with most tokens per session
sess_dominant_model = sess_model.sort_values("tok", ascending=False).drop_duplicates("session_id")[["session_id", "model"]]
sess_tokens = sess_tokens.merge(sess_dominant_model, on="session_id", how="left")

def calc_cost(row):
    return _calc_cost(row.get("model") or "", row["input_tokens"] or 0, row["output_tokens"] or 0)

sess_tokens["est_cost"] = sess_tokens.apply(calc_cost, axis=1)

no_edit_tokens = sess_tokens[~sess_tokens["had_edits"]]
with_edit_tokens = sess_tokens[sess_tokens["had_edits"]]

token_waste = {
    "sessions_without_edits": {
        "count": len(no_edit_tokens),
        "total_cost_usd": round(no_edit_tokens["est_cost"].sum(), 2),
        "pct_of_total_cost": round(no_edit_tokens["est_cost"].sum() / max(sess_tokens["est_cost"].sum(), 1) * 100, 1),
        "avg_cost_per_session": round(no_edit_tokens["est_cost"].mean(), 2),
    },
    "sessions_with_edits": {
        "count": len(with_edit_tokens),
        "total_cost_usd": round(with_edit_tokens["est_cost"].sum(), 2),
        "pct_of_total_cost": round(with_edit_tokens["est_cost"].sum() / max(sess_tokens["est_cost"].sum(), 1) * 100, 1),
        "avg_cost_per_session": round(with_edit_tokens["est_cost"].mean(), 2),
    },
    "by_developer": [],
}

for email in sess_tokens["user_email"].unique():
    dev_data = sess_tokens[sess_tokens["user_email"] == email]
    dev_no_edit = dev_data[~dev_data["had_edits"]]
    dev_with_edit = dev_data[dev_data["had_edits"]]
    token_waste["by_developer"].append({
        "email": email,
        "no_edit_sessions": len(dev_no_edit),
        "no_edit_cost": round(dev_no_edit["est_cost"].sum(), 2),
        "edit_sessions": len(dev_with_edit),
        "edit_cost": round(dev_with_edit["est_cost"].sum(), 2),
        "waste_pct": round(dev_no_edit["est_cost"].sum() / max(dev_data["est_cost"].sum(), 0.01) * 100, 1),
    })

# Add to chart_data
chart_data = load("chart_data.json")
chart_data["token_waste"] = token_waste
chart_data["token_waste"]["headline"] = f"${token_waste['sessions_without_edits']['total_cost_usd']:,.0f} ({token_waste['sessions_without_edits']['pct_of_total_cost']}%) spent on sessions that produced no code changes"

save("chart_data.json", chart_data)

# ---------------------------------------------------------------------------
# 4. Add token waste insight
# ---------------------------------------------------------------------------
print("Adding token waste insight...")

insights_data = load("insights.json")
insights = insights_data["insights"]

# Check if already exists
if not any(i["id"] == "token_waste" for i in insights):
    insights.append({
        "id": "token_waste",
        "title": f"${token_waste['sessions_without_edits']['total_cost_usd']:,.0f} spent on sessions with no code output",
        "severity": "critical",
        "metric": f"{token_waste['sessions_without_edits']['pct_of_total_cost']}% of spend → 0 edits",
        "detail": f"Out of ${round(sess_tokens['est_cost'].sum()):,} total spend, ${token_waste['sessions_without_edits']['total_cost_usd']:,.0f} ({token_waste['sessions_without_edits']['pct_of_total_cost']}%) went to {token_waste['sessions_without_edits']['count']} sessions that never edited a file. Average cost: ${token_waste['sessions_without_edits']['avg_cost_per_session']:.2f}/session (no-edit) vs ${token_waste['sessions_with_edits']['avg_cost_per_session']:.2f}/session (with edits).",
        "recommendation": "Not all no-edit sessions are waste — some are legitimate exploration. But at this cost ratio, even a 20% reduction would save significant budget.",
        "data": token_waste,
    })

insights_data["insights"] = insights
save("insights.json", insights_data)

# ---------------------------------------------------------------------------
# 5. PER-DEVELOPER PROMPT COMPLEXITY (avg first-message length)
# ---------------------------------------------------------------------------
print("Computing prompt complexity per developer...")

first_msgs = query_df(conn, f"""
    WITH ranked AS (
        SELECT m.content, s.user_email, m.session_id,
               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'
        AND m.content IS NOT NULL AND LENGTH(m.content) > 0
        AND m.session_id IN ({session_ids_sql})
    )
    SELECT content, user_email, session_id FROM ranked WHERE rn = 1
""")

first_msgs["length"] = first_msgs["content"].str.len()
dev_prompt_stats = first_msgs.groupby("user_email")["length"].agg(["mean", "median"]).reset_index()

for dev in profiles["developers"]:
    row = dev_prompt_stats[dev_prompt_stats["user_email"] == dev["email"]]
    if len(row) > 0:
        dev["avg_first_prompt_length"] = round(row.iloc[0]["mean"])
        dev["median_first_prompt_length"] = round(row.iloc[0]["median"])

save("developer_profiles.json", profiles)

conn.close()
print("\nIteration 5 enhancements complete.")
