"""Enrichment pass: additional analyses on top of base data.

Adds:
1. Per-developer work schedule heatmap (hour x day-of-week)
2. LLM-powered qualitative analysis (tacit knowledge, intent, abandonment, session narratives)
3. Monthly trends and longitudinal data
4. Prompt effectiveness (count vs success, length vs success)
5. Struggle files (5+ edits to same file per session)
6. Day-of-week patterns
7. Developer learning curves (monthly zero-edit trends)
8. Inter-session cadence
9. Developer efficiency (edits per dollar)
10. Zero-edit rate by hour
11. Cost per productive session
"""
import json
import os
import sys
import hashlib
import re
from collections import Counter
from decimal import Decimal
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")
CACHE_DIR = Path("/home/sagar/trace/analysis-14022026/output/.llm_cache")
CACHE_DIR.mkdir(parents=True, exist_ok=True)

conn = init()

from dotenv import load_dotenv
load_dotenv("/home/sagar/trace/.local.env")

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

SESSION_FILTER_SQL = """
    s.source NOT IN ('test', 'qc_trace_install')
    AND s.user_email IS NOT NULL
    AND s.org ILIKE 'pratilipi%%'
"""

IDLE_THRESHOLD_SEC = 30 * 60

from utils.pricing import calc_cost as _calc_cost

EDIT_TOOLS = {"Edit", "Write", "MultiEdit", "MultiEditTool"}
EXPLORE_TOOLS = {"Read", "Grep", "Glob", "View", "ReadFile", "read_file",
                 "grep_search", "codebase_search", "list_dir", "file_search"}


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


def sanitize(obj):
    if isinstance(obj, float):
        if np.isnan(obj) or np.isinf(obj) or abs(obj) > 1e12:
            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) or abs(v) > 1e12) 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


def strip_xml_tags(text):
    if not isinstance(text, str):
        return text
    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 = re.sub(r'\n\s*\n\s*\n', '\n\n', text)
    return text.strip()


def strip_xml_recursive(obj):
    if isinstance(obj, str):
        return strip_xml_tags(obj)
    if isinstance(obj, list):
        return [strip_xml_recursive(item) for item in obj]
    if isinstance(obj, dict):
        return {k: strip_xml_recursive(v) for k, v in obj.items()}
    return obj


def save(name, data):
    data = strip_xml_recursive(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}")


def strip_system_prompts(text):
    if not text:
        return ""
    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()


def is_empty_prompt(text):
    if not text:
        return True
    clean = text.strip().lower()
    return clean in ('', '/clear', 'warmup', '/warmup', 'clear')


# --- LLM helpers (with disk cache) ---
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=1000):
    from openai import OpenAI
    key = _cache_key(prompt, system)
    cached = _cache_get(key)
    if cached is not None:
        return cached
    oai = OpenAI()
    messages = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": prompt})
    resp = oai.chat.completions.create(
        model="gpt-4o-mini", messages=messages,
        max_tokens=max_tokens, temperature=0.3,
    )
    result = resp.choices[0].message.content.strip()
    _cache_set(key, result)
    return result


def llm_call_json(prompt, system="", max_tokens=1000):
    result = llm_call(prompt, system, max_tokens)
    json_match = re.search(r"```(?:json)?\s*([\s\S]*?)```", result)
    if json_match:
        result = json_match.group(1).strip()
    return json.loads(result)


# ===========================================================================
# 1. PER-DEVELOPER WORK SCHEDULE HEATMAP (hour x day-of-week)
# ===========================================================================
print("1. Computing per-developer work schedule heatmaps...")

work_patterns = query_df(conn, f"""
    SELECT s.user_email,
           EXTRACT(DOW FROM m.timestamp) as dow,
           EXTRACT(HOUR FROM m.timestamp) as hour_utc,
           COUNT(*) as msg_count
    FROM messages m
    JOIN sessions s ON s.id = m.session_id
    WHERE s.id IN ({SESSION_FILTER})
    AND m.msg_type = 'user'
    GROUP BY s.user_email, EXTRACT(DOW FROM m.timestamp), EXTRACT(HOUR FROM m.timestamp)
""")

work_schedule = {"developers": []}
for email in work_patterns["user_email"].unique():
    dev = work_patterns[work_patterns["user_email"] == email]
    name = email.split("@")[0]
    heatmap = []
    for _, row in dev.iterrows():
        heatmap.append({
            "dow": int(row["dow"]),
            "hour_utc": int(row["hour_utc"]),
            "hour_ist": int((float(row["hour_utc"]) + 5.5) % 24),
            "count": int(row["msg_count"]),
        })
    peak = dev.nlargest(3, "msg_count")
    peak_hours_ist = [int((float(h) + 5.5) % 24) for h in peak["hour_utc"]]
    work_schedule["developers"].append({
        "email": email,
        "name": name,
        "heatmap": heatmap,
        "peak_hours_ist": peak_hours_ist,
        "total_interactions": int(dev["msg_count"].sum()),
    })

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


# ===========================================================================
# 2. MONTHLY TRENDS
# ===========================================================================
print("\n2. Computing monthly trends...")

date_range_df = query_df(conn, """
    SELECT MIN(m.timestamp)::date as first_msg, MAX(m.timestamp)::date as last_msg,
           MAX(m.timestamp)::date - MIN(m.timestamp)::date as span_days
    FROM messages m JOIN sessions s ON s.id = m.session_id
    WHERE s.org ILIKE 'pratilipi%%' AND m.timestamp IS NOT NULL
""")
first_date = str(date_range_df.iloc[0]["first_msg"])
last_date = str(date_range_df.iloc[0]["last_msg"])
span_days = int(date_range_df.iloc[0]["span_days"])
span_months = round(span_days / 30.4, 1)

monthly = query_df(conn, """
    SELECT DATE_TRUNC('month', m.timestamp)::date as month,
           COUNT(DISTINCT m.session_id) as sessions,
           COUNT(*) as messages,
           COUNT(CASE WHEN m.msg_type = 'user' THEN 1 END) as user_msgs
    FROM messages m JOIN sessions s ON s.id = m.session_id
    WHERE s.org ILIKE 'pratilipi%%' AND m.timestamp IS NOT NULL
    GROUP BY 1 ORDER BY 1
""")

monthly_cost_df = query_df(conn, """
    SELECT DATE_TRUNC('month', m.timestamp)::date as month, m.model,
           SUM(t.input_tokens) as input_tok, SUM(t.output_tokens) as output_tok,
           SUM(t.cached_tokens) as cached_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.org ILIKE 'pratilipi%%' AND m.timestamp IS NOT NULL
    GROUP BY 1, 2 ORDER BY 1
""")

def calc_cost(row):
    return _calc_cost(row.get("model") or "", float(row["input_tok"] or 0),
                      float(row["output_tok"] or 0), float(row["cached_tok"] or 0))

monthly_cost_df["cost"] = monthly_cost_df.apply(calc_cost, axis=1)
cost_by_month = monthly_cost_df.groupby("month")["cost"].sum().reset_index()

zero_edit_monthly = query_df(conn, """
    WITH session_months AS (
        SELECT s.id, DATE_TRUNC('month', MIN(m.timestamp))::date as month
        FROM sessions s JOIN messages m ON m.session_id = s.id
        WHERE s.org ILIKE 'pratilipi%%' AND m.timestamp IS NOT NULL
        GROUP BY s.id
    ),
    session_edits AS (
        SELECT m.session_id, COUNT(*) as edits
        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%%'
        AND tc.tool_name IN ('Edit', 'Write', 'MultiEdit', 'MultiEditTool')
        GROUP BY m.session_id
    )
    SELECT sm.month, COUNT(*) as total,
           COUNT(CASE WHEN COALESCE(se.edits,0) = 0 THEN 1 END) as zero_edit
    FROM session_months sm LEFT JOIN session_edits se ON se.session_id = sm.id
    WHERE sm.month IS NOT NULL
    GROUP BY sm.month ORDER BY sm.month
""")

monthly_trends = {"headline": f"AI usage over {span_months} months ({first_date} to {last_date})", "months": []}
month_labels = {
    "2025-09-01": "Sep 25", "2025-10-01": "Oct 25", "2025-11-01": "Nov 25",
    "2025-12-01": "Dec 25", "2026-01-01": "Jan 26", "2026-02-01": "Feb 26",
}
for _, row in monthly.iterrows():
    m_str = str(row["month"])
    label = month_labels.get(m_str, m_str[:7])
    cost_row = cost_by_month[cost_by_month["month"] == row["month"]]
    cost_val = round(float(cost_row["cost"].iloc[0]), 2) if len(cost_row) > 0 else 0
    ze_row = zero_edit_monthly[zero_edit_monthly["month"] == row["month"]]
    ze_pct = round(float(ze_row["zero_edit"].iloc[0]) / float(ze_row["total"].iloc[0]) * 100, 1) if len(ze_row) > 0 and float(ze_row["total"].iloc[0]) > 0 else 0
    monthly_trends["months"].append({
        "month": m_str, "label": label,
        "sessions": int(row["sessions"]), "messages": int(row["messages"]),
        "user_msgs": int(row["user_msgs"]),
        "cost_usd": cost_val, "zero_edit_pct": ze_pct,
    })

# Per-developer monthly adoption
dev_monthly = query_df(conn, """
    SELECT s.user_email, DATE_TRUNC('month', m.timestamp)::date as month,
           COUNT(DISTINCT s.id) as sessions, s.source
    FROM messages m JOIN sessions s ON s.id = m.session_id
    WHERE s.org ILIKE 'pratilipi%%' AND m.timestamp IS NOT NULL
    GROUP BY s.user_email, DATE_TRUNC('month', m.timestamp)::date, s.source
    ORDER BY s.user_email, month
""")
dev_timeline = {}
for _, row in dev_monthly.iterrows():
    email = row["user_email"]
    name = email.split("@")[0]
    if name not in dev_timeline:
        dev_timeline[name] = {"email": email, "months": []}
    dev_timeline[name]["months"].append({
        "month": str(row["month"]),
        "label": month_labels.get(str(row["month"]), str(row["month"])[:7]),
        "sessions": int(row["sessions"]), "source": row["source"],
    })
monthly_trends["developer_adoption"] = list(dev_timeline.values())

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


# ===========================================================================
# 3. PROMPT EFFECTIVENESS (count and length vs success)
# ===========================================================================
print("\n3. Computing prompt effectiveness...")

prompt_count_vs_edits = query_df(conn, """
WITH session_stats AS (
    SELECT m.session_id,
           COUNT(CASE WHEN m.msg_type = 'user' THEN 1 END) as user_msgs,
           COUNT(CASE WHEN tc.tool_name IN ('Edit','Write') THEN 1 END) as edits
    FROM messages m JOIN sessions s ON s.id = m.session_id
    LEFT JOIN tool_calls tc ON tc.message_id = m.id
    WHERE s.org LIKE 'pratilipi%%'
    GROUP BY m.session_id
)
SELECT CASE WHEN user_msgs = 1 THEN '1 prompt'
            WHEN user_msgs BETWEEN 2 AND 3 THEN '2-3 prompts'
            WHEN user_msgs BETWEEN 4 AND 7 THEN '4-7 prompts'
            WHEN user_msgs BETWEEN 8 AND 15 THEN '8-15 prompts'
            ELSE '16+ prompts' END as bucket,
       MIN(user_msgs) as sort_key, COUNT(*) as sessions,
       ROUND(AVG(edits)::numeric, 1) as avg_edits,
       ROUND(100.0 * COUNT(CASE WHEN edits > 0 THEN 1 END) / COUNT(*), 1) as pct_with_edits
FROM session_stats GROUP BY 1 ORDER BY MIN(user_msgs)
""")

prompt_length_vs_edits = query_df(conn, """
WITH first_prompts AS (
    SELECT m.session_id, MIN(m.id) as first_msg_id
    FROM messages m JOIN sessions s ON s.id = m.session_id
    WHERE s.org LIKE 'pratilipi%%' AND m.msg_type = 'user' GROUP BY m.session_id
),
session_data AS (
    SELECT fp.session_id, LENGTH(m.content) as prompt_len,
           (SELECT COUNT(*) FROM tool_calls tc2 JOIN messages m2 ON m2.id = tc2.message_id
            WHERE m2.session_id = fp.session_id AND tc2.tool_name IN ('Edit','Write')) as edits
    FROM first_prompts fp JOIN messages m ON m.id = fp.first_msg_id
)
SELECT CASE WHEN prompt_len < 50 THEN '<50 chars'
            WHEN prompt_len < 200 THEN '50-200 chars'
            WHEN prompt_len < 500 THEN '200-500 chars'
            WHEN prompt_len < 1000 THEN '500-1K chars'
            ELSE '1K+ chars' END as bucket,
       MIN(prompt_len) as sort_key, COUNT(*) as sessions,
       ROUND(AVG(edits)::numeric, 1) as avg_edits,
       ROUND(100.0 * COUNT(CASE WHEN edits > 0 THEN 1 END) / COUNT(*), 1) as pct_with_edits
FROM session_data WHERE prompt_len IS NOT NULL GROUP BY 1 ORDER BY MIN(prompt_len)
""")

chart_data = load("chart_data.json")
chart_data["prompt_count_effectiveness"] = {
    "data": prompt_count_vs_edits.to_dict("records"),
    "methodology": "Edit success = session has at least one Edit/Write tool call.",
}
chart_data["prompt_length_effectiveness"] = {
    "data": prompt_length_vs_edits.to_dict("records"),
    "methodology": "First user message length vs session edit count.",
}
save("chart_data.json", chart_data)


# ===========================================================================
# 4. STRUGGLE FILES (5+ edits to same file per session)
# ===========================================================================
print("\n4. Computing struggle files...")

struggle_files = query_df(conn, """
SELECT tc.tool_input->>'file_path' as file_path,
       m.session_id, COUNT(*) as edit_count, 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 s.org LIKE 'pratilipi%%'
  AND tc.tool_name IN ('Edit','Write','MultiEdit','MultiEditTool')
  AND tc.tool_input->>'file_path' IS NOT NULL
GROUP BY tc.tool_input->>'file_path', m.session_id, s.user_email
HAVING COUNT(*) >= 5
ORDER BY COUNT(*) DESC LIMIT 20
""")

struggle_list = []
for _, row in struggle_files.iterrows():
    fp = row["file_path"]
    if fp:
        fp = re.sub(r'^/[^/]+/[^/]+/[^/]+/', '', fp)
    struggle_list.append({
        "file": fp or row["file_path"],
        "session_id": row["session_id"],
        "edit_count": int(row["edit_count"]),
        "developer": row["user_email"].split("@")[0] if "@" in str(row["user_email"]) else str(row["user_email"]),
    })

chart_data = load("chart_data.json")
chart_data["struggle_files"] = {
    "data": struggle_list,
    "methodology": "Files edited 5+ times in a single session — indicates iterative fixing.",
}
save("chart_data.json", chart_data)


# ===========================================================================
# 5. DAY-OF-WEEK PATTERNS
# ===========================================================================
print("\n5. Computing day-of-week patterns...")

dow_data = query_df(conn, """
WITH session_days AS (
    SELECT m.session_id,
           TO_CHAR(MIN(m.timestamp) AT TIME ZONE 'Asia/Kolkata', 'Dy') as day_name,
           EXTRACT(DOW FROM MIN(m.timestamp) AT TIME ZONE 'Asia/Kolkata') as dow,
           COUNT(CASE WHEN tc.tool_name IN ('Edit','Write') THEN 1 END) as edits,
           COUNT(CASE WHEN m.msg_type = 'user' THEN 1 END) as user_msgs
    FROM messages m JOIN sessions s ON s.id = m.session_id
    LEFT JOIN tool_calls tc ON tc.message_id = m.id
    WHERE s.org LIKE 'pratilipi%%'
    GROUP BY m.session_id
)
SELECT day_name, dow::int as dow, COUNT(*) as sessions,
    ROUND(100.0 * COUNT(CASE WHEN edits = 0 THEN 1 END) / COUNT(*), 1) as zero_edit_pct,
    ROUND(AVG(edits)::numeric, 1) as avg_edits,
    ROUND(AVG(user_msgs)::numeric, 1) as avg_prompts
FROM session_days GROUP BY day_name, dow ORDER BY dow
""")

chart_data = load("chart_data.json")
chart_data["day_of_week"] = {
    "data": dow_data.to_dict("records"),
    "methodology": "Session day determined by first message timestamp in IST.",
}
save("chart_data.json", chart_data)


# ===========================================================================
# 6. DEVELOPER LEARNING CURVES (monthly zero-edit trends)
# ===========================================================================
print("\n6. Computing developer learning curves...")

learning_curves = query_df(conn, """
WITH session_edits AS (
    SELECT m.session_id,
           TO_CHAR(MIN(m.timestamp) AT TIME ZONE 'Asia/Kolkata', 'YYYY-MM') as month,
           s.user_email,
           COUNT(CASE WHEN tc.tool_name IN ('Edit','Write','MultiEdit','MultiEditTool') THEN 1 END) as edit_count
    FROM messages m JOIN sessions s ON s.id = m.session_id
    LEFT JOIN tool_calls tc ON tc.message_id = m.id
    WHERE s.org LIKE 'pratilipi%%'
    GROUP BY m.session_id, s.user_email
)
SELECT month, user_email, COUNT(*) as sessions,
       COUNT(CASE WHEN edit_count = 0 THEN 1 END) as zero_edit,
       ROUND(100.0 * COUNT(CASE WHEN edit_count = 0 THEN 1 END) / COUNT(*), 1) as zero_edit_pct,
       ROUND(AVG(edit_count)::numeric, 1) as avg_edits
FROM session_edits GROUP BY month, user_email ORDER BY month, user_email
""")

lc_by_dev = {}
for _, row in learning_curves.iterrows():
    dev = row["user_email"].split("@")[0] if "@" in str(row["user_email"]) else str(row["user_email"])
    if dev not in lc_by_dev:
        lc_by_dev[dev] = []
    lc_by_dev[dev].append({
        "month": row["month"], "sessions": int(row["sessions"]),
        "zero_edit_pct": float(row["zero_edit_pct"]), "avg_edits": float(row["avg_edits"]),
    })

chart_data = load("chart_data.json")
chart_data["developer_learning_curves"] = {
    "data": lc_by_dev,
    "methodology": "Monthly zero-edit rate per developer. Lower is better.",
}
save("chart_data.json", chart_data)


# ===========================================================================
# 7. INTER-SESSION CADENCE
# ===========================================================================
print("\n7. Computing inter-session cadence...")

cadence_data = query_df(conn, """
WITH session_times AS (
    SELECT s.user_email, MIN(m.timestamp) as session_start,
           LAG(MIN(m.timestamp)) OVER (PARTITION BY s.user_email ORDER BY MIN(m.timestamp)) as prev_start
    FROM messages m JOIN sessions s ON s.id = m.session_id
    WHERE s.org LIKE 'pratilipi%%'
    GROUP BY s.id, s.user_email
)
SELECT user_email,
       ROUND(AVG(EXTRACT(EPOCH FROM session_start - prev_start) / 3600)::numeric, 1) as avg_gap_hours,
       ROUND(PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY EXTRACT(EPOCH FROM session_start - prev_start) / 3600)::numeric, 1) as median_gap_hours,
       COUNT(*) as session_transitions
FROM session_times WHERE prev_start IS NOT NULL GROUP BY user_email
""")

cadence_list = []
for _, row in cadence_data.iterrows():
    dev = row["user_email"].split("@")[0] if "@" in str(row["user_email"]) else str(row["user_email"])
    cadence_list.append({
        "developer": dev,
        "avg_gap_hours": float(row["avg_gap_hours"]),
        "median_gap_hours": float(row["median_gap_hours"]),
        "session_transitions": int(row["session_transitions"]),
    })

chart_data = load("chart_data.json")
chart_data["session_cadence"] = {"data": cadence_list}
save("chart_data.json", chart_data)


# ===========================================================================
# 8. DEVELOPER EFFICIENCY (edits per dollar)
# ===========================================================================
print("\n8. Computing developer efficiency...")

dev_efficiency = query_df(conn, """
WITH dev_edits AS (
    SELECT s.user_email, COUNT(DISTINCT tc.tool_id) as edit_count
    FROM sessions s JOIN messages m ON m.session_id = s.id
    JOIN tool_calls tc ON tc.message_id = m.id
    WHERE s.org LIKE 'pratilipi%%'
    AND tc.tool_name IN ('Edit', 'Write', 'MultiEdit', 'MultiEditTool')
    GROUP BY s.user_email
),
dev_tokens AS (
    SELECT s.user_email, m.model,
           SUM(tu.input_tokens) as total_input, SUM(tu.output_tokens) as total_output
    FROM sessions s JOIN messages m ON m.session_id = s.id
    JOIN token_usage tu ON tu.message_id = m.id
    WHERE s.org LIKE 'pratilipi%%'
    GROUP BY s.user_email, m.model
)
SELECT de.user_email, de.edit_count, dt.model, dt.total_input, dt.total_output
FROM dev_edits de JOIN dev_tokens dt ON dt.user_email = de.user_email
""")

# Aggregate cost per developer using per-model pricing
dev_eff_map = {}
for _, row in dev_efficiency.iterrows():
    email = row["user_email"]
    if email not in dev_eff_map:
        dev_eff_map[email] = {"edit_count": int(row["edit_count"]), "cost": 0}
    dev_eff_map[email]["cost"] += _calc_cost(row["model"] or "", row["total_input"] or 0, row["total_output"] or 0)

dev_eff_list = []
for email, d in dev_eff_map.items():
    edits_per_dollar = round(d["edit_count"] / max(d["cost"], 0.01), 1)
    dev_eff_list.append({
        "developer": email,
        "edit_count": d["edit_count"],
        "estimated_cost": round(d["cost"], 2),
        "edits_per_dollar": edits_per_dollar,
    })
dev_eff_list.sort(key=lambda x: x["edits_per_dollar"], reverse=True)

overview = load("overview.json")
overview["developer_efficiency"] = dev_eff_list
save("overview.json", overview)


# ===========================================================================
# 9. ZERO-EDIT RATE BY HOUR
# ===========================================================================
print("\n9. Computing zero-edit rate by hour...")

hourly_df = query_df(conn, """
WITH session_hours AS (
    SELECT s.id as session_id,
           EXTRACT(HOUR FROM MIN(m.timestamp) AT TIME ZONE 'Asia/Kolkata') as hour_ist
    FROM sessions s JOIN messages m ON m.session_id = s.id
    WHERE s.org LIKE 'pratilipi%%' GROUP BY s.id
),
session_edits AS (
    SELECT DISTINCT m.session_id FROM messages m
    JOIN sessions s ON s.id = m.session_id
    JOIN tool_calls tc ON tc.message_id = m.id
    WHERE s.org LIKE 'pratilipi%%'
    AND tc.tool_name IN ('Edit', 'Write', 'MultiEdit', 'MultiEditTool')
)
SELECT sh.hour_ist::int as hour, COUNT(*) as total_sessions,
       COUNT(*) - COUNT(se.session_id) as zero_edit_sessions,
       ROUND(100.0 * (COUNT(*) - COUNT(se.session_id)) / COUNT(*), 1) as zero_edit_pct
FROM session_hours sh LEFT JOIN session_edits se ON se.session_id = sh.session_id
GROUP BY sh.hour_ist ORDER BY sh.hour_ist
""")

zero_edit_by_hour = []
for _, row in hourly_df.iterrows():
    zero_edit_by_hour.append({
        "hour": int(row["hour"]),
        "hour_label": f"{int(row['hour']):02d}:00 IST",
        "total_sessions": int(row["total_sessions"]),
        "zero_edit_sessions": int(row["zero_edit_sessions"]),
        "zero_edit_pct": float(row["zero_edit_pct"]),
    })

chart_data = load("chart_data.json")
chart_data["zero_edit_rate_by_hour"] = {
    "data": zero_edit_by_hour,
    "methodology": "Percentage of sessions started in each hour (IST) that produced zero edits.",
}
save("chart_data.json", chart_data)


# ===========================================================================
# 10. COST PER PRODUCTIVE SESSION
# ===========================================================================
print("\n10. Computing cost per productive session...")

sessions_with_edits_df = query_df(conn, """
    SELECT COUNT(DISTINCT m.session_id) as cnt
    FROM messages m JOIN sessions s ON s.id = m.session_id
    JOIN tool_calls tc ON tc.message_id = m.id
    WHERE s.org LIKE 'pratilipi%%'
    AND tc.tool_name IN ('Edit', 'Write', 'MultiEdit', 'MultiEditTool')
""")
sessions_with_edits = int(sessions_with_edits_df["cnt"].iloc[0])

overview = load("overview.json")
total_cost = overview.get("estimated_cost_usd", 0)
total_sessions = overview.get("total_sessions", 1)
cost_per_session = round(total_cost / total_sessions, 2) if total_sessions else 0
cost_per_productive = round(total_cost / sessions_with_edits, 2) if sessions_with_edits else 0

overview["sessions_with_edits"] = sessions_with_edits
overview["cost_per_session"] = cost_per_session
overview["cost_per_productive_session"] = cost_per_productive
# Cap savings_pct at 100%
if overview.get("savings_pct", 0) > 100:
    overview["savings_pct"] = 100.0
save("overview.json", overview)


# ===========================================================================
# 11. QUALITATIVE ANALYSIS (LLM-powered)
# ===========================================================================
print("\n11. Running qualitative analysis...")

msgs = query_df(conn, f"""
    SELECT m.id, m.session_id, m.msg_type, m.timestamp, m.content,
           s.user_email, s.source
    FROM messages m JOIN sessions s ON s.id = m.session_id
    WHERE {SESSION_FILTER_SQL}
    ORDER BY m.session_id, m.timestamp
""")

sessions_raw = query_df(conn, f"""
    SELECT s.* FROM sessions s WHERE {SESSION_FILTER_SQL}
    ORDER BY s.first_seen
""")

# Compute active times
active_times = {}
for sid, grp in msgs.groupby("session_id"):
    ts = grp["timestamp"].sort_values()
    if len(ts) < 2:
        active_times[sid] = 0.0
        continue
    gaps = ts.diff().dropna().dt.total_seconds()
    active_sec = gaps[gaps < IDLE_THRESHOLD_SEC].sum()
    active_times[sid] = round(active_sec / 60.0, 2)
sessions_raw["active_min"] = sessions_raw["id"].map(active_times).fillna(0)

user_msgs = msgs[msgs["msg_type"] == "user"].copy()
user_msgs["clean"] = user_msgs["content"].apply(strip_system_prompts)
user_msgs = user_msgs[user_msgs["clean"].str.len() > 5]

tc = query_df(conn, f"""
    SELECT tc.tool_name, tc.tool_input, m.session_id, 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 {SESSION_FILTER_SQL}
    ORDER BY m.timestamp
""")

# --- 11a: Session narratives for top 20 sessions by activity ---
print("  11a: Session narratives...")

session_msg_counts = user_msgs.groupby("session_id").size().sort_values(ascending=False)
top_sessions = session_msg_counts.head(20).index.tolist()

SESSION_ARC_SYSTEM = """You are analyzing a developer-AI coding session. Given the user's prompts, determine:
1. What was the developer trying to accomplish?
2. Did they pivot or change direction?
3. Was the session successful, partially successful, or abandoned?
4. One-sentence narrative summary

Respond in JSON:
{"goal": "brief goal", "pivots": ["direction changes"], "outcome": "successful|partial|abandoned|exploration_only",
 "narrative": "One sentence story", "difficulty": "low|medium|high"}"""

session_narratives = []
for sid in top_sessions:
    s_msgs = user_msgs[user_msgs["session_id"] == sid].sort_values("timestamp")
    prompts = s_msgs["clean"].tolist()[:15]
    if not prompts:
        continue
    s_info = sessions_raw[sessions_raw["id"] == sid].iloc[0] if sid in sessions_raw["id"].values else None
    s_tc = tc[tc["session_id"] == sid]
    n_edits = (s_tc["tool_name"].isin(EDIT_TOOLS)).sum()
    n_explores = (s_tc["tool_name"].isin(EXPLORE_TOOLS)).sum()

    prompt_text = "\n---\n".join(p[:300] for p in prompts)
    context = f"\nTool calls: {len(s_tc)} ({n_edits} edits, {n_explores} explores)\nTotal user prompts: {len(prompts)}"

    try:
        arc = llm_call_json(f"Session prompts:\n{prompt_text}\n\nContext:{context}",
                           system=SESSION_ARC_SYSTEM, max_tokens=400)
    except Exception as e:
        print(f"    LLM failed for {sid[:12]}: {e}")
        arc = {"goal": "unknown", "narrative": f"Session with {len(prompts)} prompts", "outcome": "unknown", "difficulty": "medium"}

    # Filter out empty/warmup opening prompts
    opening = prompts[0][:200] if prompts else ""
    if is_empty_prompt(opening) and len(prompts) > 1:
        for p in prompts[1:]:
            if not is_empty_prompt(p) and len(p) > 5:
                opening = p[:200]
                break

    arc["session_id"] = sid
    arc["n_prompts"] = len(prompts)
    arc["n_tool_calls"] = len(s_tc)
    arc["n_edits"] = int(n_edits)
    arc["active_min"] = round(active_times.get(sid, 0), 1)
    arc["developer"] = s_info["user_email"].split("@")[0] if s_info is not None else "unknown"
    arc["source"] = s_info["source"] if s_info is not None else "unknown"
    arc["opening_prompt"] = opening
    session_narratives.append(arc)

print(f"    {len(session_narratives)} session narratives generated")

# --- 11b: Developer behavioral analysis ---
print("  11b: Developer behavioral analysis...")

DEV_BEHAVIOR_SYSTEM = """You are a developer behavior analyst. Analyze how developers use AI coding assistants.

CRITICAL RULES:
- Do NOT restate metrics. Focus on behavioral patterns from reading the actual prompts.
- Every pattern must cite specific evidence.

Given sample prompts, identify:
1. Working style (delegator/collaborator/explorer/builder/firefighter)
2. 3-4 unique behavioral patterns with evidence
3. Blind spots
4. Recommendations

Respond in JSON:
{"working_style": "style name",
 "behavioral_patterns": [{"pattern": "observation", "evidence": "quote", "implication": "impact"}],
 "blind_spots": ["specific things"],
 "recommendations": [{"action": "step", "expected_impact": "benefit"}],
 "assessment": "2-3 sentence behavioral profile"}"""

developer_summaries = []
for email in sessions_raw["user_email"].unique():
    dev_msgs = user_msgs[user_msgs["user_email"] == email]
    dev_sessions = sessions_raw[sessions_raw["user_email"] == email]
    dev_tc = tc[tc["user_email"] == email]

    if len(dev_sessions) < 3:
        continue

    first_prompts = dev_msgs.sort_values("timestamp").groupby("session_id").first()
    prompt_samples = first_prompts["clean"].tolist()[:15]

    mid_prompts = []
    for sid, grp in dev_msgs.groupby("session_id"):
        sorted_msgs = grp.sort_values("timestamp")
        if len(sorted_msgs) >= 4:
            mid = sorted_msgs.iloc[2:5]["clean"].tolist()
            mid_prompts.extend(m[:200] for m in mid if len(m) > 10)
    mid_prompts = mid_prompts[:10]

    context = f"Developer: {email.split('@')[0]}\nCLI tools: {dev_sessions['source'].value_counts().to_dict()}\n\n"
    context += "OPENING PROMPTS:\n" + "\n---\n".join(p[:250] for p in prompt_samples[:12])
    if mid_prompts:
        context += "\n\nMID-SESSION PROMPTS:\n" + "\n---\n".join(mid_prompts[:8])

    try:
        result = llm_call_json(context, system=DEV_BEHAVIOR_SYSTEM, max_tokens=1200)
    except Exception as e:
        print(f"    LLM failed for {email}: {e}")
        result = {"working_style": "unknown", "behavioral_patterns": [], "blind_spots": [],
                  "recommendations": [], "assessment": "Analysis could not be completed."}

    result["email"] = email
    result["developer"] = email.split("@")[0]
    result["sessions"] = int(len(dev_sessions))
    result["active_hours"] = round(dev_sessions["active_min"].sum() / 60, 2)
    developer_summaries.append(result)

print(f"    {len(developer_summaries)} developer summaries generated")

# --- 11c: Tacit knowledge extraction ---
print("  11c: Tacit knowledge extraction...")

TACIT_SYSTEM = """Extract implicit codebase knowledge from developer-AI conversations.
Categories: architecture, conventions, gotchas, domain, tooling

Respond in JSON:
{"facts": [{"category": "category", "fact": "the knowledge", "evidence": "brief quote"}],
 "claude_md_suggestions": ["suggested lines for CLAUDE.md"]}"""

all_clean_msgs = user_msgs["clean"].tolist()
tacit_facts = []
claude_md_suggestions = []
batch_size = 40
for i in range(0, min(len(all_clean_msgs), 200), batch_size):
    batch = all_clean_msgs[i:i+batch_size]
    combined = "\n---\n".join(m[:300] for m in batch)
    try:
        result = llm_call_json(f"Developer messages:\n{combined}", system=TACIT_SYSTEM, max_tokens=1000)
        tacit_facts.extend(result.get("facts", []))
        claude_md_suggestions.extend(result.get("claude_md_suggestions", []))
    except Exception as e:
        print(f"    Tacit batch {i} failed: {e}")

# Deduplicate
deduped_facts = []
for fact in tacit_facts:
    words = set(fact.get("fact", "").lower().split())
    is_dup = False
    for existing in deduped_facts:
        ex_words = set(existing.get("fact", "").lower().split())
        if existing.get("category") == fact.get("category"):
            overlap = len(words & ex_words) / max(len(words | ex_words), 1)
            if overlap > 0.5:
                existing["times_seen"] = existing.get("times_seen", 1) + 1
                is_dup = True
                break
    if not is_dup:
        fact["times_seen"] = 1
        deduped_facts.append(fact)
deduped_facts.sort(key=lambda f: -f.get("times_seen", 1))
claude_md_suggestions = list(dict.fromkeys(claude_md_suggestions))

# --- 11d: Abandonment analysis ---
print("  11d: Abandonment analysis...")

sessions_with_edits_set = set(tc[tc["tool_name"].isin(EDIT_TOOLS)]["session_id"].unique())
abandoned_candidates = []
for sid, grp in user_msgs.groupby("session_id"):
    if sid in sessions_with_edits_set or len(grp) < 3:
        continue
    prompts = grp["clean"].tolist()
    abandoned_candidates.append({
        "session_id": sid,
        "n_prompts": len(prompts),
        "n_tool_calls": len(tc[tc["session_id"] == sid]),
        "active_min": active_times.get(sid, 0),
        "opening_prompt": prompts[0][:200] if prompts else "",
        "developer": grp["user_email"].iloc[0].split("@")[0],
    })
abandoned_candidates.sort(key=lambda x: -x["n_prompts"])

ABANDON_SYSTEM = """A developer started this AI coding session but never produced any code edits.
Possible reasons: wrong_tool, stuck, exploration_only, context_switching, insufficient_context
Respond in JSON: {"reason": "category", "explanation": "1-2 sentences"}"""

abandonment_analysis = []
for cand in abandoned_candidates[:10]:
    context = f"Opening prompt: {cand['opening_prompt']}\nTool calls: {cand['n_tool_calls']}\nUser prompts: {cand['n_prompts']}\nActive time: {cand['active_min']:.1f} min"
    try:
        result = llm_call_json(context, system=ABANDON_SYSTEM, max_tokens=200)
        result.update(cand)
        abandonment_analysis.append(result)
    except Exception:
        cand["reason"] = "unknown"
        cand["explanation"] = "LLM analysis failed"
        abandonment_analysis.append(cand)

# --- Build qualitative_insights.json ---
print("  Building qualitative_insights.json...")

abandon_reasons = Counter(a.get("reason", "unknown") for a in abandonment_analysis)

qualitative = {
    "methodology": (
        f"Qualitative analysis powered by gpt-4o-mini with disk caching. "
        f"System tags stripped. "
        f"Analyzed {len(session_narratives)} session narratives, {len(developer_summaries)} developer profiles, "
        f"{len(deduped_facts)} tacit knowledge facts, {len(abandonment_analysis)} abandonment analyses."
    ),
    "session_narratives": session_narratives,
    "developer_summaries": developer_summaries,
    "tacit_knowledge": {
        "facts": deduped_facts[:30],
        "patterns": [f.get("fact", "") for f in deduped_facts if f.get("fact")][:30],
        "claude_md_suggestions": claude_md_suggestions[:15],
        "headline": f"{len(deduped_facts)} pieces of implicit codebase knowledge",
    },
    "abandonment_analysis": {
        "total_abandoned": len(abandoned_candidates),
        "analyzed": len(abandonment_analysis),
        "reasons": dict(abandon_reasons.most_common()),
        "details": abandonment_analysis[:10],
    },
}

# Key findings
findings = []
if abandonment_analysis:
    top_reason = abandon_reasons.most_common(1)[0] if abandon_reasons else ("unknown", 0)
    findings.append({
        "title": f"Most common reason for abandoned sessions: {top_reason[0].replace('_', ' ')}",
        "detail": f"{top_reason[1]} of {len(abandonment_analysis)} analyzed.",
        "category": "behavioral",
    })
for ds in developer_summaries:
    patterns = ds.get("behavioral_patterns", [])
    for p in patterns[:2]:
        findings.append({
            "title": f"{ds['developer']}: {p.get('pattern', '')[:80]}",
            "detail": p.get("implication", p.get("evidence", "")),
            "category": "developer_behavior",
        })
qualitative["key_qualitative_findings"] = findings

save("qualitative_insights.json", qualitative)

# Update developer profiles with behavioral assessments
dp = load("developer_profiles.json")
dev_assessments = {ds.get("developer", ""): ds.get("assessment", "") for ds in developer_summaries}
for dev in dp.get("developers", []):
    email = dev.get("email", "")
    name = email.split("@")[0]
    if name in dev_assessments:
        dev["llm_summary"] = dev_assessments[name]
save("developer_profiles.json", dp)


# ===========================================================================
# 12. FINAL SYNC — ensure all files copied to dashboard
# ===========================================================================
print("\n12. Syncing all JSON files to dashboard...")
for name in OUTPUT_DIR.glob("*.json"):
    dst = DASHBOARD_DIR / name.name
    if not dst.exists() or name.read_text() != dst.read_text():
        dst.parent.mkdir(parents=True, exist_ok=True)
        dst.write_text(name.read_text())
        print(f"  synced {name.name}")

conn.close()
print("\nEnrichment complete.")
