"""Iteration 7 enhancements: LLM-powered session narratives, learning curves, error cascades, recommendations upgrade.

Key additions:
1. LLM-powered session narratives for top 10 expensive sessions (gpt-4o-mini)
2. Developer learning curves — are devs getting better over time?
3. Error/failure cascade analysis — what happens after a tool failure?
4. Context utilization patterns — session token growth trajectories
5. Upgraded recommendations with estimated savings
"""
import json
import os
import sys
import hashlib
import re
from pathlib import Path
from datetime import timedelta

import numpy as np
import pandas as pd
from decimal import Decimal

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")
CACHE_DIR = Path("/home/sagar/trace/analysis-14022026/output/.llm_cache")
CACHE_DIR.mkdir(parents=True, exist_ok=True)

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


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


def save(name, data):
    def sanitize(obj):
        if isinstance(obj, float):
            if np.isnan(obj) or np.isinf(obj):
                return None
            return obj
        if isinstance(obj, dict):
            return {k: sanitize(v) for k, v in obj.items()}
        if isinstance(obj, list):
            return [sanitize(v) for v in obj]
        if isinstance(obj, (np.integer,)):
            return int(obj)
        if isinstance(obj, (np.floating,)):
            v = float(obj)
            return None if np.isnan(v) or np.isinf(v) else v
        if isinstance(obj, (np.bool_,)):
            return bool(obj)
        if isinstance(obj, pd.Timestamp):
            return str(obj)
        if isinstance(obj, Decimal):
            return float(obj)
        return obj

    data = sanitize(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}")


# --- LLM helpers ---
def _cache_key(prompt, system=""):
    return hashlib.sha256((system + "|||" + prompt).encode()).hexdigest()


def _cache_get(key):
    path = CACHE_DIR / f"{key}.json"
    if path.exists():
        with open(path) as f:
            return json.load(f)["result"]
    return None


def _cache_set(key, result):
    path = CACHE_DIR / f"{key}.json"
    with open(path, "w") as f:
        json.dump({"result": result}, f)


def llm_call(prompt, system="", max_tokens=500):
    from dotenv import load_dotenv
    load_dotenv("/home/sagar/trace/.local.env")
    from openai import OpenAI
    oai = OpenAI()
    key = _cache_key(prompt, system)
    cached = _cache_get(key)
    if cached is not None:
        return cached
    msgs = []
    if system:
        msgs.append({"role": "system", "content": system})
    msgs.append({"role": "user", "content": prompt})
    resp = oai.chat.completions.create(model="gpt-4o-mini", messages=msgs, max_tokens=max_tokens, temperature=0.3)
    result = resp.choices[0].message.content.strip()
    _cache_set(key, result)
    return result


def strip_system_prompts(text):
    if not text:
        return ""
    # Catch-all: strip ALL XML-like tags and their content
    text = re.sub(r'<[a-zA-Z_-]+>.*?</[a-zA-Z_-]+>', '', text, flags=re.DOTALL)
    text = re.sub(r'<[a-zA-Z_-]+\s*/>', '', text)
    text = re.sub(r'<[a-zA-Z_-]+>', '', text)
    text = text.encode("ascii", errors="ignore").decode("ascii")
    return text.strip()


conn = init()

# ---------------------------------------------------------------------------
# 1. LLM-POWERED SESSION NARRATIVES for top 10 most expensive sessions
# ---------------------------------------------------------------------------
print("1. Generating LLM session narratives for top expensive sessions...")

from utils.pricing import calc_cost as _calc_cost

top_sessions = query_df(conn, f"""
    SELECT s.id as session_id, s.user_email, s.source, s.repo_name, s.first_seen, s.last_updated,
           SUM(t.input_tokens) as input_tokens, SUM(t.output_tokens) as output_tokens,
           (SELECT m2.model FROM messages m2 JOIN token_usage t2 ON t2.message_id = m2.id
            WHERE m2.session_id = s.id AND m2.model IS NOT NULL
            GROUP BY m2.model ORDER BY SUM(t2.input_tokens) DESC LIMIT 1) as dominant_model
    FROM sessions s
    JOIN messages m ON m.session_id = s.id
    JOIN token_usage t ON t.message_id = m.id
    WHERE s.id IN ({SESSION_FILTER})
    GROUP BY s.id, s.user_email, s.source, s.repo_name, s.first_seen, s.last_updated
    ORDER BY (SUM(t.input_tokens) / 1e6 * 3 + SUM(t.output_tokens) / 1e6 * 15) DESC
    LIMIT 10
""")
top_sessions["est_cost"] = top_sessions.apply(
    lambda r: _calc_cost(r.get("dominant_model") or "", r["input_tokens"] or 0, r["output_tokens"] or 0), axis=1)

narratives = []
for _, sess in top_sessions.iterrows():
    sid = sess["session_id"]

    # Get conversation summary (first 5 user messages + tool call names)
    user_msgs = query_df(conn, """
        SELECT content FROM messages
        WHERE session_id = %s AND msg_type = 'user' AND content IS NOT NULL AND LENGTH(content) > 0
        ORDER BY timestamp LIMIT 5
    """, (sid,))

    tools_used = query_df(conn, """
        SELECT tc.tool_name, COUNT(*) as cnt
        FROM tool_calls tc JOIN messages m ON m.id = tc.message_id
        WHERE m.session_id = %s
        GROUP BY tc.tool_name ORDER BY cnt DESC LIMIT 8
    """, (sid,))

    edit_count = query_df(conn, """
        SELECT COUNT(*) as cnt FROM tool_calls tc JOIN messages m ON m.id = tc.message_id
        WHERE m.session_id = %s AND tc.tool_name IN ('Edit', 'Write', 'MultiEdit', 'MultiEditTool')
    """, (sid,))

    # Build prompt
    cleaned_msgs = []
    for _, row in user_msgs.iterrows():
        clean = strip_system_prompts(row["content"])
        if clean and len(clean) > 10:
            cleaned_msgs.append(clean[:500])

    tools_str = ", ".join(f"{r['tool_name']}({r['cnt']})" for _, r in tools_used.iterrows())
    edits = int(edit_count.iloc[0]["cnt"]) if len(edit_count) > 0 else 0
    cost = round(sess["est_cost"], 2)
    source = sess["source"]

    prompt = f"""Analyze this AI coding session and write a 2-sentence narrative for a VP of Engineering.

Session: {source}, cost ${cost}, {edits} file edits
Tools used: {tools_str}
User messages (first 5):
{chr(10).join(f'- {m}' for m in cleaned_msgs[:5])}

Write exactly 2 sentences:
1. What the developer was trying to accomplish
2. Whether the session was productive or wasteful and why (be specific about the cost vs output)

Be direct and blunt. If $300 was spent to read files with no edits, say that."""

    system = "You are a concise engineering analyst. No fluff. No hedging. State facts bluntly."

    try:
        narrative = llm_call(prompt, system=system, max_tokens=150)
    except Exception as e:
        print(f"  LLM call failed for {sid}: {e}")
        narrative = f"${cost} session via {source} with {edits} edits. {'Productive' if edits > 5 else 'Exploration-heavy'}."

    narratives.append({
        "session_id": sid,
        "email": sess["user_email"],
        "source": source,
        "repo": sess["repo_name"] or "(no repo)",
        "cost_usd": cost,
        "edits": edits,
        "narrative": narrative,
    })
    print(f"  Session {sid[:8]}... ${cost} -> narrative generated")

# Update chart_data with narratives
chart_data = load("chart_data.json")
chart_data["session_narratives"] = narratives
save("chart_data.json", chart_data)

# Also update session_deepdives with narratives
if "session_deepdives" in chart_data:
    narrative_lookup = {n["session_id"]: n["narrative"] for n in narratives}
    for dd in chart_data["session_deepdives"]:
        if dd["session_id"] in narrative_lookup:
            dd["narrative"] = narrative_lookup[dd["session_id"]]
    save("chart_data.json", chart_data)

# ---------------------------------------------------------------------------
# 2. DEVELOPER LEARNING CURVES — are devs improving over time?
# ---------------------------------------------------------------------------
print("\n2. Computing developer learning curves...")

daily_metrics = query_df(conn, f"""
    SELECT s.user_email,
           s.first_seen::date as day,
           COUNT(DISTINCT s.id) as sessions,
           COUNT(CASE WHEN m.msg_type = 'user' THEN 1 END) as user_msgs,
           COUNT(CASE WHEN tc.tool_name IN ('Edit', 'Write', 'MultiEdit', 'MultiEditTool') THEN 1 END) as edits,
           COUNT(CASE WHEN tc.tool_name IN ('Read', 'Grep', 'Glob', 'View', 'ReadFile', 'read_file', 'grep_search', 'codebase_search', 'list_dir', 'file_search', 'search_file_content') THEN 1 END) as explores,
           SUM(COALESCE(t.input_tokens, 0)) as input_tokens,
           SUM(COALESCE(t.output_tokens, 0)) as output_tokens
    FROM sessions s
    JOIN messages m ON m.session_id = s.id
    LEFT JOIN tool_calls tc ON tc.message_id = m.id
    LEFT JOIN token_usage t ON t.message_id = m.id
    WHERE s.id IN ({SESSION_FILTER})
    GROUP BY s.user_email, s.first_seen::date
    ORDER BY s.user_email, s.first_seen::date
""")

# Use blended pricing as approximation for daily metrics (per-model not available per-row here)
daily_metrics["cost"] = daily_metrics["input_tokens"] / 1e6 * 3 + daily_metrics["output_tokens"] / 1e6 * 15
daily_metrics["edits_per_session"] = daily_metrics["edits"] / daily_metrics["sessions"].clip(lower=1)
daily_metrics["cost_per_session"] = daily_metrics["cost"] / daily_metrics["sessions"].clip(lower=1)
daily_metrics["explore_to_edit"] = daily_metrics["explores"] / daily_metrics["edits"].clip(lower=1)

learning_curves = {"developers": [], "headline": ""}

for email in daily_metrics["user_email"].unique():
    dev_data = daily_metrics[daily_metrics["user_email"] == email].sort_values("day")
    name = email.split("@")[0].replace("+", " ").split(" ")[-1]

    days_data = []
    for _, row in dev_data.iterrows():
        days_data.append({
            "day": str(row["day"]),
            "sessions": int(row["sessions"]),
            "edits": int(row["edits"]),
            "cost": round(row["cost"], 2),
            "edits_per_session": round(row["edits_per_session"], 1),
            "cost_per_session": round(row["cost_per_session"], 2),
            "explore_to_edit": round(min(row["explore_to_edit"], 20), 1),
        })

    # Trend: is cost_per_session decreasing over time?
    if len(dev_data) >= 2:
        first_half = dev_data.iloc[:len(dev_data) // 2]["cost_per_session"].mean()
        second_half = dev_data.iloc[len(dev_data) // 2:]["cost_per_session"].mean()
        trend = "improving" if second_half < first_half * 0.8 else "stable" if second_half < first_half * 1.2 else "worsening"
    else:
        trend = "insufficient_data"

    learning_curves["developers"].append({
        "email": email,
        "name": name,
        "days": days_data,
        "trend": trend,
        "total_days": len(days_data),
    })

# Headline
trends = [d["trend"] for d in learning_curves["developers"]]
if "worsening" in trends:
    learning_curves["headline"] = "Some developers are getting more expensive per session over time, not less"
elif "improving" in trends:
    learning_curves["headline"] = "Developers are trending toward lower cost-per-session over time"
else:
    learning_curves["headline"] = "Cost efficiency is flat across the observation period — no learning curve visible yet"

chart_data = load("chart_data.json")
chart_data["learning_curves"] = learning_curves
save("chart_data.json", chart_data)

# ---------------------------------------------------------------------------
# 3. ERROR CASCADE ANALYSIS — what happens after a tool failure?
# ---------------------------------------------------------------------------
print("\n3. Computing error cascade patterns...")

# Get all tool results with failure status and what follows
error_sessions = query_df(conn, f"""
    WITH tool_seq AS (
        SELECT m.session_id, m.timestamp, tc.tool_name, tr.status,
               LEAD(tc.tool_name) OVER (PARTITION BY m.session_id ORDER BY m.timestamp) as next_tool,
               LEAD(tr.status) OVER (PARTITION BY m.session_id ORDER BY m.timestamp) as next_status
        FROM messages m
        JOIN tool_calls tc ON tc.message_id = m.id
        LEFT JOIN tool_results tr ON tr.call_id = tc.tool_id
        JOIN sessions s ON s.id = m.session_id
        WHERE s.id IN ({SESSION_FILTER})
    )
    SELECT tool_name, status, next_tool, next_status, COUNT(*) as cnt
    FROM tool_seq
    WHERE status = 'failure'
    GROUP BY tool_name, status, next_tool, next_status
    ORDER BY cnt DESC
    LIMIT 20
""")

# Tool failure rate overall
failure_stats = query_df(conn, f"""
    SELECT tc.tool_name,
           COUNT(*) as total,
           COUNT(CASE WHEN tr.status = 'failure' THEN 1 END) as failures
    FROM tool_calls tc
    JOIN messages m ON m.id = tc.message_id
    LEFT JOIN tool_results tr ON tr.call_id = tc.tool_id
    JOIN sessions s ON s.id = m.session_id
    WHERE s.id IN ({SESSION_FILTER})
    GROUP BY tc.tool_name
    HAVING COUNT(*) > 5
    ORDER BY COUNT(CASE WHEN tr.status = 'failure' THEN 1 END) DESC
""")
failure_stats["fail_pct"] = round(failure_stats["failures"] / failure_stats["total"] * 100, 1)

error_cascades = {
    "headline": f"Tool failures are rare ({failure_stats['failures'].sum()} of {failure_stats['total'].sum()} calls) but concentrated",
    "total_failures": int(failure_stats["failures"].sum()),
    "total_calls": int(failure_stats["total"].sum()),
    "failure_rate_pct": round(failure_stats["failures"].sum() / max(failure_stats["total"].sum(), 1) * 100, 2),
    "by_tool": [],
    "cascades": [],
}

for _, row in failure_stats[failure_stats["failures"] > 0].iterrows():
    error_cascades["by_tool"].append({
        "tool": row["tool_name"],
        "total": int(row["total"]),
        "failures": int(row["failures"]),
        "fail_pct": float(row["fail_pct"]),
    })

for _, row in error_sessions.head(10).iterrows():
    error_cascades["cascades"].append({
        "failed_tool": row["tool_name"],
        "next_tool": row["next_tool"],
        "next_succeeded": row["next_status"] != "failure" if row["next_status"] else None,
        "count": int(row["cnt"]),
    })

chart_data = load("chart_data.json")
chart_data["error_cascades"] = error_cascades
save("chart_data.json", chart_data)

# ---------------------------------------------------------------------------
# 4. SESSION TOKEN GROWTH — how fast do sessions consume tokens?
# ---------------------------------------------------------------------------
print("\n4. Computing session token growth trajectories...")

token_growth = query_df(conn, f"""
    SELECT m.session_id, m.timestamp, s.user_email, s.source,
           COALESCE(t.input_tokens, 0) as input_tokens,
           SUM(COALESCE(t.input_tokens, 0)) OVER (PARTITION BY m.session_id ORDER BY m.timestamp) as cumulative_input
    FROM messages m
    JOIN sessions s ON s.id = m.session_id
    LEFT JOIN token_usage t ON t.message_id = m.id
    WHERE s.id IN ({SESSION_FILTER})
    AND t.input_tokens IS NOT NULL
    ORDER BY m.session_id, m.timestamp
""")

# Per-session: at what message count does input token usage explode?
session_token_profiles = []
for sid, group in token_growth.groupby("session_id"):
    if len(group) < 3:
        continue
    group = group.sort_values("timestamp")
    max_cumulative = group["cumulative_input"].max()
    msg_count = len(group)
    source = group["source"].iloc[0]

    # Token velocity: tokens per message (slope)
    velocity = max_cumulative / msg_count if msg_count > 0 else 0

    # Did context grow super-linearly? Compare first-half vs second-half velocity
    half = len(group) // 2
    first_half_tokens = group.iloc[:half]["input_tokens"].sum()
    second_half_tokens = group.iloc[half:]["input_tokens"].sum()
    acceleration = second_half_tokens / max(first_half_tokens, 1)

    session_token_profiles.append({
        "session_id": sid,
        "source": source,
        "msg_count": msg_count,
        "total_input_tokens": int(max_cumulative),
        "tokens_per_msg": round(velocity, 0),
        "acceleration": round(acceleration, 2),
    })

token_df = pd.DataFrame(session_token_profiles)

# Sessions where context grows super-linearly (acceleration > 2 = second half uses 2x+ more tokens)
accelerating = token_df[token_df["acceleration"] > 2]

token_growth_data = {
    "headline": f"{len(accelerating)} sessions show accelerating token usage (context growing faster than linearly)",
    "total_sessions_analyzed": len(token_df),
    "accelerating_sessions": len(accelerating),
    "median_tokens_per_msg": int(token_df["tokens_per_msg"].median()) if len(token_df) > 0 else 0,
    "by_source": [],
    "acceleration_distribution": [],
}

# By source
for source in token_df["source"].unique():
    src_data = token_df[token_df["source"] == source]
    token_growth_data["by_source"].append({
        "source": source,
        "median_tokens_per_msg": int(src_data["tokens_per_msg"].median()),
        "sessions": len(src_data),
        "pct_accelerating": round(len(src_data[src_data["acceleration"] > 2]) / max(len(src_data), 1) * 100, 1),
    })

# Acceleration histogram
bins = [0, 0.5, 1.0, 1.5, 2.0, 3.0, 5.0, 100]
labels = ["<0.5x", "0.5-1x", "1-1.5x", "1.5-2x", "2-3x", "3-5x", "5x+"]
for i in range(len(bins) - 1):
    count = int(((token_df["acceleration"] >= bins[i]) & (token_df["acceleration"] < bins[i + 1])).sum())
    token_growth_data["acceleration_distribution"].append({"range": labels[i], "count": count})

chart_data = load("chart_data.json")
chart_data["token_growth"] = token_growth_data
save("chart_data.json", chart_data)

# ---------------------------------------------------------------------------
# 5. UPGRADED EXECUTIVE SUMMARY — lead with the root cause story
# ---------------------------------------------------------------------------
print("\n5. Upgrading executive summary and recommendations...")

overview = load("overview.json")

# Recompute some key metrics for a tighter narrative using per-model pricing
source_model_tokens = query_df(conn, f"""
    SELECT s.source, m.model,
           SUM(COALESCE(t.input_tokens, 0)) as input_tokens,
           SUM(COALESCE(t.output_tokens, 0)) as output_tokens
    FROM sessions s
    JOIN messages m ON m.session_id = s.id
    LEFT JOIN token_usage t ON t.message_id = m.id
    WHERE s.id IN ({SESSION_FILTER})
    GROUP BY s.source, m.model
""")
_source_cost_map = {}
for _, r in source_model_tokens.iterrows():
    src = r["source"]
    _source_cost_map[src] = _source_cost_map.get(src, 0) + _calc_cost(
        r["model"] or "", r["input_tokens"] or 0, r["output_tokens"] or 0)

codex_cost_rows = []
source_sessions = query_df(conn, f"""
    SELECT s.source, COUNT(DISTINCT s.id) as sessions
    FROM sessions s WHERE s.id IN ({SESSION_FILTER}) GROUP BY s.source
""")
for _, r in source_sessions.iterrows():
    codex_cost_rows.append({
        "source": r["source"],
        "cost": _source_cost_map.get(r["source"], 0),
        "sessions": int(r["sessions"]),
    })
codex_cost = pd.DataFrame(codex_cost_rows)
total_cost = float(codex_cost["cost"].sum())
codex_row = codex_cost[codex_cost["source"] == "codex_cli"]
codex_pct = round(codex_row["cost"].iloc[0] / total_cost * 100, 1) if len(codex_row) > 0 else 0
codex_sessions = int(codex_row["sessions"].iloc[0]) if len(codex_row) > 0 else 0
codex_cost_per_session = round(codex_row["cost"].iloc[0] / max(codex_sessions, 1), 2) if len(codex_row) > 0 else 0

claude_row = codex_cost[codex_cost["source"] == "claude_code"]
claude_cost_per_session = round(claude_row["cost"].iloc[0] / max(int(claude_row["sessions"].iloc[0]), 1), 2) if len(claude_row) > 0 else 0

overview["executive_summary"] = (
    f"{overview['total_sessions']} AI coding sessions across {overview['unique_users']} developers. "
    f"Codex CLI consumes {codex_pct}% of AI spend (${round(codex_row['cost'].iloc[0], 0) if len(codex_row) > 0 else 0}) "
    f"despite only {codex_sessions} sessions. "
    f"Estimated cost: ${overview['estimated_cost_usd']:.2f}."
)

# Better recommendations with estimated savings
overview["recommendations"] = [
    {
        "priority": "critical",
        "title": "Switch from Codex CLI to Claude Code",
        "detail": f"Codex CLI costs ${codex_cost_per_session}/session vs ${claude_cost_per_session} for Claude Code. Migrating {codex_sessions} Codex sessions would save ~${round(codex_row['cost'].iloc[0] * 0.9, 0) if len(codex_row) > 0 else 0}/month.",
        "estimated_savings_monthly": round(codex_row["cost"].iloc[0] * 0.9, 0) if len(codex_row) > 0 else 0,
        "effort": "low",
    },
    {
        "priority": "high",
        "title": "Set per-session cost cap at $50",
        "detail": "Top 10 sessions consumed 74% of spend. A $50 cap would catch runaway sessions before they drain the budget.",
        "estimated_savings_monthly": round(total_cost * 0.3, 0),
        "effort": "low",
    },
    {
        "priority": "high",
        "title": "Pre-load context for Pratilipi/chitrakatha",
        "detail": "This one repo is 83% of AI spend. Creating a .cursorrules or CLAUDE.md with key file paths could cut exploration overhead 30-50%.",
        "estimated_savings_monthly": round(total_cost * 0.15, 0),
        "effort": "medium",
    },
    {
        "priority": "medium",
        "title": "Consolidate exploration sessions",
        "detail": f"{overview.get('zero_edit_pct', 60)}% of sessions produce no edits. Training devs to combine quick lookups into single sessions could reduce session overhead.",
        "estimated_savings_monthly": round(total_cost * 0.05, 0),
        "effort": "medium",
    },
]

# Add total savings potential
total_savings = sum(r["estimated_savings_monthly"] for r in overview["recommendations"])
overview["total_savings_potential"] = round(total_savings, 0)
overview["savings_pct"] = round(total_savings / max(total_cost, 1) * 100, 1)

save("overview.json", overview)

# ---------------------------------------------------------------------------
# 6. CODEX VS CLAUDE COMPARISON — detailed head-to-head
# ---------------------------------------------------------------------------
print("\n6. Computing Codex vs Claude Code head-to-head comparison...")

cli_comparison = query_df(conn, f"""
    SELECT s.source,
           COUNT(DISTINCT s.id) as sessions,
           COUNT(CASE WHEN m.msg_type = 'user' THEN 1 END) as user_msgs,
           COUNT(CASE WHEN tc.tool_name IN ('Edit', 'Write', 'MultiEdit', 'MultiEditTool') THEN 1 END) as edits,
           COUNT(CASE WHEN tc.tool_name IN ('Read', 'Grep', 'Glob', 'View', 'ReadFile', 'read_file', 'grep_search', 'search_file_content') THEN 1 END) as explores,
           AVG(CASE WHEN m.msg_type = 'user' THEN 1 ELSE 0 END) as avg_user_msg_pct,
           SUM(COALESCE(t.input_tokens, 0)) as input_tokens,
           SUM(COALESCE(t.output_tokens, 0)) as output_tokens
    FROM sessions s
    JOIN messages m ON m.session_id = s.id
    LEFT JOIN tool_calls tc ON tc.message_id = m.id
    LEFT JOIN token_usage t ON t.message_id = m.id
    WHERE s.id IN ({SESSION_FILTER})
    GROUP BY s.source
""")
# Compute CLI comparison costs using per-model pricing
cli_model_costs = query_df(conn, f"""
    SELECT s.source, m.model,
           SUM(COALESCE(t.input_tokens, 0)) as input_tokens,
           SUM(COALESCE(t.output_tokens, 0)) as output_tokens
    FROM sessions s
    JOIN messages m ON m.session_id = s.id
    LEFT JOIN token_usage t ON t.message_id = m.id
    WHERE s.id IN ({SESSION_FILTER})
    GROUP BY s.source, m.model
""")
cli_cost_by_source = {}
for _, r in cli_model_costs.iterrows():
    src = r["source"]
    cli_cost_by_source[src] = cli_cost_by_source.get(src, 0) + _calc_cost(
        r["model"] or "", r["input_tokens"] or 0, r["output_tokens"] or 0)

for col in ["input_tokens", "output_tokens", "edits", "explores", "sessions"]:
    cli_comparison[col] = cli_comparison[col].astype(float)
cli_comparison["cost"] = cli_comparison["source"].map(cli_cost_by_source).fillna(0)
cli_comparison["cost_per_session"] = cli_comparison["cost"] / cli_comparison["sessions"]
cli_comparison["edits_per_session"] = cli_comparison["edits"] / cli_comparison["sessions"]
cli_comparison["explores_per_session"] = cli_comparison["explores"] / cli_comparison["sessions"]

comparison_data = {
    "headline": f"Codex CLI costs {round(codex_cost_per_session / max(claude_cost_per_session, 0.01))}x more per session than Claude Code",
    "sources": [],
}

for _, row in cli_comparison.iterrows():
    comparison_data["sources"].append({
        "source": row["source"],
        "sessions": int(row["sessions"]),
        "cost_total": round(row["cost"], 2),
        "cost_per_session": round(row["cost_per_session"], 2),
        "edits_per_session": round(row["edits_per_session"], 1),
        "explores_per_session": round(row["explores_per_session"], 1),
        "pct_of_total_cost": round(row["cost"] / max(total_cost, 1) * 100, 1),
    })

chart_data = load("chart_data.json")
chart_data["cli_comparison"] = comparison_data
save("chart_data.json", chart_data)

# ---------------------------------------------------------------------------
# 7. INSIGHT QUALITY PASS — ensure all insights have actionable titles
# ---------------------------------------------------------------------------
print("\n7. Polishing insights and red flags...")

insights_data = load("insights.json")

# Make sure the codex cost story is the #1 red flag
red_flags = insights_data.get("red_flags", [])

# Update or add codex_dominant red flag
codex_flag = {
    "id": "codex_dominant",
    "headline": f"Codex CLI = {codex_pct}% of spend but only {codex_sessions} sessions — switching saves ${round(codex_row['cost'].iloc[0] * 0.9, 0) if len(codex_row) > 0 else 0}/month",
    "detail": f"Each Codex CLI session costs ${codex_cost_per_session} vs ${claude_cost_per_session} for Claude Code. This is a {round(codex_cost_per_session / max(claude_cost_per_session, 0.01))}x premium. The fix is simple: migrate to Claude Code.",
    "severity": "critical",
    "metric_value": f"${codex_cost_per_session}/session",
    "metric_label": f"vs ${claude_cost_per_session} Claude Code",
}

# Replace existing codex flags or add
red_flags = [f for f in red_flags if f["id"] not in ("codex_cost", "codex_dominant")]
red_flags.insert(0, codex_flag)  # Make it #1
insights_data["red_flags"] = red_flags

save("insights.json", insights_data)

# ---------------------------------------------------------------------------
# 8. PER-DEVELOPER SESSION NARRATIVE SUMMARIES (LLM)
# ---------------------------------------------------------------------------
print("\n8. Generating per-developer LLM summaries...")

profiles = load("developer_profiles.json")

for dev in profiles["developers"]:
    email = dev["email"]
    name = email.split("@")[0]
    sessions = dev.get("total_sessions", 0)
    cost = dev.get("estimated_cost_usd", 0)
    edits = dev.get("total_edit_calls", 0)
    persona = dev.get("persona", "unknown")
    primary_intent = dev.get("primary_intent", "unknown")
    active_pct = dev.get("active_pct", 50)
    sources = dev.get("cli_breakdown", {})

    prompt = f"""Write a 2-sentence performance summary for a VP about this developer's AI usage.

Developer: {name}
Sessions: {sessions}, Cost: ${cost}, File edits: {edits}
Persona: {persona}, Primary intent: {primary_intent}
Active time: {active_pct}% of wall-clock
CLI tools: {json.dumps(sources)}

Be specific with numbers. If cost is disproportionate to output, say so directly. If efficient, highlight why."""

    try:
        summary = llm_call(prompt, system="You are a blunt engineering analyst. 2 sentences max.", max_tokens=100)
        dev["llm_summary"] = summary
    except Exception as e:
        print(f"  LLM failed for {name}: {e}")
        dev["llm_summary"] = dev.get("developer_summary", "")

    print(f"  {name}: summary generated")

save("developer_profiles.json", profiles)

# ---------------------------------------------------------------------------
# DONE
# ---------------------------------------------------------------------------
conn.close()
print("\nIteration 7 enhancements complete.")
print(f"Files updated: chart_data.json, overview.json, insights.json, developer_profiles.json")
