#!/usr/bin/env python3
"""PARSER agent: parse 34 Codex CLI JSONL sessions and compute heuristic metrics."""

import json
import os
import re
from collections import Counter
from datetime import datetime
from pathlib import Path
from statistics import mean

REPO_ROOT = Path("/home/sagar/trace")
INPUT_DIRS = [REPO_ROOT / "docs/internal/2026/01" / d for d in ("28", "29", "30", "31")]
OUTPUT_DIR = REPO_ROOT / "docs/run1-09022026/analysis-outputs"
PARSED_DIR = OUTPUT_DIR / "parsed"
METRICS_FILE = OUTPUT_DIR / "metrics.json"
STATUS_FILE = OUTPUT_DIR / "coordination/parser.json"


def write_status(status, current_task, completed_tasks, sessions_parsed, sessions_total, errors):
    STATUS_FILE.write_text(json.dumps({
        "status": status,
        "current_task": current_task,
        "completed_tasks": completed_tasks,
        "sessions_parsed": sessions_parsed,
        "sessions_total": sessions_total,
        "errors": errors,
    }, indent=2))


def truncate(text, max_len=200):
    if not text:
        return ""
    s = str(text)
    return s[:max_len] + "..." if len(s) > max_len else s


def extract_text_from_content(content):
    """Extract text from a content array."""
    if not content:
        return ""
    if isinstance(content, str):
        return content
    parts = []
    for item in content:
        if isinstance(item, str):
            parts.append(item)
        elif isinstance(item, dict):
            if item.get("type") in ("input_text", "text"):
                parts.append(item.get("text", ""))
            elif item.get("type") == "output_text":
                parts.append(item.get("text", ""))
    return "\n".join(parts)


def parse_session(filepath: Path) -> dict:
    """Parse a single JSONL session file into a ParsedSession."""
    filename = filepath.name
    # Extract date and session_id from filename
    # Format: rollout-2026-01-28T15-51-17-019c054d-499a-7f32-bee4-5a701f1ce557.jsonl
    m = re.match(r"rollout-(\d{4}-\d{2}-\d{2})T[\d-]+-([0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})\.jsonl", filename)
    date = m.group(1) if m else "unknown"
    session_id = m.group(2) if m else "unknown"

    events = []
    parse_errors = []
    with open(filepath, "r") as f:
        for lineno, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                events.append(json.loads(line))
            except json.JSONDecodeError as e:
                parse_errors.append(f"Line {lineno}: {e}")

    raw_event_count = len(events)
    timestamps = []
    model = None
    cli_version = None
    approval_policy = None
    effort = None
    summary_from_context = None
    user_messages = []
    assistant_messages = []
    tool_calls = []
    files_touched = set()
    user_msg_idx = 0
    asst_msg_idx = 0

    for evt in events:
        ts = evt.get("timestamp", "")
        if ts:
            timestamps.append(ts)

        evt_type = evt.get("type", "")
        payload = evt.get("payload", {})
        if not isinstance(payload, dict):
            continue

        if evt_type == "session_meta":
            cli_version = payload.get("cli_version", cli_version)

        elif evt_type == "turn_context":
            model = payload.get("model", model)
            approval_policy = payload.get("approval_policy", approval_policy)
            effort = payload.get("effort", effort)
            s = payload.get("summary")
            if s:
                summary_from_context = s

        elif evt_type == "response_item":
            role = payload.get("role", "")
            p_type = payload.get("type", "")
            content = payload.get("content", [])

            if p_type == "function_call":
                name = payload.get("name", "unknown")
                args = payload.get("arguments", "")
                tool_calls.append({
                    "name": name,
                    "input_summary": truncate(args),
                    "output_summary": ""
                })
                # Track files from tool calls
                if isinstance(args, str):
                    try:
                        args_obj = json.loads(args)
                    except (json.JSONDecodeError, TypeError):
                        args_obj = {}
                elif isinstance(args, dict):
                    args_obj = args
                else:
                    args_obj = {}
                for key in ("file_path", "path", "filename", "file"):
                    v = args_obj.get(key)
                    if v and isinstance(v, str):
                        files_touched.add(v)
                # Also check command for file references
                cmd = args_obj.get("command", "")
                if cmd and isinstance(cmd, str):
                    # Extract paths that look like file references
                    pass

            elif p_type == "function_call_output":
                output = payload.get("output", "")
                # Attach to last tool call
                if tool_calls:
                    tool_calls[-1]["output_summary"] = truncate(output)

            elif role in ("developer", "user"):
                text = extract_text_from_content(content)
                if text.strip():
                    user_messages.append({
                        "index": user_msg_idx,
                        "text": text,
                        "timestamp": ts
                    })
                    user_msg_idx += 1

            elif role == "assistant":
                text = extract_text_from_content(content)
                if text.strip():
                    assistant_messages.append({
                        "index": asst_msg_idx,
                        "text": text,
                        "timestamp": ts
                    })
                    asst_msg_idx += 1

    # Compute timestamps
    first_timestamp = timestamps[0] if timestamps else ""
    last_timestamp = timestamps[-1] if timestamps else ""
    duration_seconds = 0.0
    if first_timestamp and last_timestamp:
        try:
            fmt_patterns = ["%Y-%m-%dT%H:%M:%S.%fZ", "%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%dT%H:%M:%S.%f", "%Y-%m-%dT%H:%M:%S"]
            t1 = t2 = None
            for fmt in fmt_patterns:
                try:
                    t1 = datetime.strptime(first_timestamp, fmt)
                    break
                except ValueError:
                    continue
            for fmt in fmt_patterns:
                try:
                    t2 = datetime.strptime(last_timestamp, fmt)
                    break
                except ValueError:
                    continue
            if t1 and t2:
                duration_seconds = round((t2 - t1).total_seconds(), 1)
        except Exception:
            pass

    turn_count = len(user_messages)

    return {
        "filename": filename,
        "date": date,
        "session_id": session_id,
        "first_timestamp": first_timestamp,
        "last_timestamp": last_timestamp,
        "duration_seconds": duration_seconds,
        "model": model,
        "cli_version": cli_version,
        "approval_policy": approval_policy,
        "effort": effort,
        "raw_event_count": raw_event_count,
        "user_messages": user_messages,
        "assistant_messages": assistant_messages,
        "tool_calls": tool_calls,
        "turn_count": turn_count,
        "files_touched": sorted(files_touched),
        "summary_from_context": summary_from_context,
    }, parse_errors


def classify_outcome(session: dict) -> str:
    raw = session["raw_event_count"]
    asst = session["assistant_messages"]
    user = session["user_messages"]
    turn_count = session["turn_count"]

    # context-only
    if raw <= 5 and (not asst or (len(asst) == 1 and len(asst[0].get("text", "")) < 100)):
        return "context-only"

    # abandoned
    if turn_count > 3:
        all_msgs = []
        for m in user:
            all_msgs.append(("user", m.get("timestamp", "")))
        for m in asst:
            all_msgs.append(("assistant", m.get("timestamp", "")))
        all_msgs.sort(key=lambda x: x[1])
        if all_msgs and all_msgs[-1][0] == "user":
            return "abandoned"

    # resolved
    if asst:
        all_msgs = []
        for m in user:
            all_msgs.append(("user", m.get("timestamp", ""), m.get("text", "")))
        for m in asst:
            all_msgs.append(("assistant", m.get("timestamp", ""), m.get("text", "")))
        all_msgs.sort(key=lambda x: x[1])
        if all_msgs and all_msgs[-1][0] == "assistant":
            # Check for resolved indicators
            all_user_text = " ".join(m.get("text", "") for m in user).lower()
            resolved_words = ["thanks", "done", "works", "perfect", "great"]
            has_tool_calls = len(session.get("tool_calls", [])) > 0
            if has_tool_calls or any(w in all_user_text for w in resolved_words):
                return "resolved"

    return "partial"


def classify_session_type(session: dict) -> str:
    all_user_text = " ".join(m.get("text", "") for m in session["user_messages"]).lower()

    bug_keywords = ["fix", "error", "crash", "fail", "bug", "broken", "issue", "traceback", "exception"]
    feature_keywords = ["add", "implement", "create", "build", "new feature", "write"]
    exploration_keywords = ["how", "why", "explain", "what", "?", "help me understand"]
    setup_keywords = ["install", "configure", "setup", "init", "dependency", "pip", "npm", "config"]
    refactor_keywords = ["refactor", "clean", "optimize", "improve", "rename", "restructure"]

    # Check in priority order
    for kw in bug_keywords:
        if kw in all_user_text:
            return "bug-fix"
    for kw in feature_keywords:
        if kw in all_user_text:
            return "feature"
    for kw in setup_keywords:
        if kw in all_user_text:
            return "setup"
    for kw in refactor_keywords:
        if kw in all_user_text:
            return "refactor"
    for kw in exploration_keywords:
        if kw in all_user_text:
            return "exploration"
    return "unknown"


def compute_prompt_specificity(session: dict) -> int:
    all_user_text = " ".join(m.get("text", "") for m in session["user_messages"])
    text_lower = all_user_text.lower()
    total_len = len(all_user_text)
    score = 1

    has_error = any(w in text_lower for w in ["error", "traceback", "exception", "stack trace"])
    has_file_ref = bool(re.search(r'[\w/]+\.\w{1,5}', all_user_text))
    has_line_num = bool(re.search(r'line \d+|:\d+', text_lower))
    has_tried = any(w in text_lower for w in ["i tried", "i've tried", "attempted", "already"])
    has_proposed_fix = any(w in text_lower for w in ["should be", "change to", "replace with", "instead of"])

    if total_len > 50:
        score = 2
    if has_error or has_file_ref:
        score = 3
    if (has_error or has_file_ref) and (has_tried or total_len > 200):
        score = 4
    if has_file_ref and has_line_num and has_proposed_fix:
        score = 5

    return score


def compute_struggle_score(session: dict) -> int:
    score = 0
    turn_count = session["turn_count"]
    if turn_count > 10:
        score += 2
    if turn_count > 20:
        score += 2

    # Repeated similar phrases
    user_texts = [m.get("text", "").lower().strip() for m in session["user_messages"]]
    if len(user_texts) >= 3:
        # Check for near-duplicate messages
        for i in range(len(user_texts)):
            for j in range(i + 1, len(user_texts)):
                if user_texts[i] and user_texts[j] and (
                    user_texts[i] == user_texts[j] or
                    (len(user_texts[i]) > 20 and user_texts[i] in user_texts[j]) or
                    (len(user_texts[j]) > 20 and user_texts[j] in user_texts[i])
                ):
                    score += 2
                    break
            else:
                continue
            break

    # Error keywords in multiple user messages
    error_msg_count = sum(1 for t in user_texts if any(w in t for w in ["error", "fail"]))
    if error_msg_count >= 2:
        score += 2

    outcome = session.get("metrics", {}).get("outcome", classify_outcome(session))
    if outcome == "abandoned":
        score += 2

    if session["duration_seconds"] > 600 and outcome != "resolved":
        score += 2

    return min(score, 10)


def compute_time_of_day(first_timestamp: str) -> str:
    if not first_timestamp:
        return "unknown"
    for fmt in ["%Y-%m-%dT%H:%M:%S.%fZ", "%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%dT%H:%M:%S.%f", "%Y-%m-%dT%H:%M:%S"]:
        try:
            dt = datetime.strptime(first_timestamp, fmt)
            hour = dt.hour
            if 6 <= hour < 12:
                return "morning"
            elif 12 <= hour < 18:
                return "afternoon"
            elif 18 <= hour < 24:
                return "evening"
            else:
                return "night"
        except ValueError:
            continue
    return "unknown"


def main():
    # Collect all JSONL files
    jsonl_files = []
    for d in INPUT_DIRS:
        if d.exists():
            jsonl_files.extend(sorted(d.glob("*.jsonl")))

    total = len(jsonl_files)
    print(f"Found {total} JSONL files to parse")

    all_errors = []
    completed_tasks = ["directories_created"]
    parsed_sessions = []

    for i, filepath in enumerate(jsonl_files):
        print(f"  [{i+1}/{total}] Parsing {filepath.name}")
        try:
            session, errors = parse_session(filepath)
            if errors:
                all_errors.extend([f"{filepath.name}: {e}" for e in errors])

            # Write parsed session
            out_path = PARSED_DIR / f"{filepath.name}.json"
            out_path.write_text(json.dumps(session, indent=2, ensure_ascii=False))
            parsed_sessions.append(session)

        except Exception as e:
            err_msg = f"{filepath.name}: {e}"
            all_errors.append(err_msg)
            print(f"    ERROR: {e}")

        # Update status every 5 sessions
        if (i + 1) % 5 == 0 or (i + 1) == total:
            write_status(
                "in_progress",
                f"parsing session {filepath.stem}",
                completed_tasks,
                i + 1,
                total,
                all_errors[:20]  # Keep errors manageable
            )

    completed_tasks.append(f"{total}_sessions_parsed")
    print(f"Parsed {len(parsed_sessions)} sessions with {len(all_errors)} errors")

    # Step 3: Compute metrics for each session
    print("Computing heuristic metrics...")
    for session in parsed_sessions:
        outcome = classify_outcome(session)
        session_type = classify_session_type(session)
        iteration_count = session["turn_count"]
        specificity = compute_prompt_specificity(session)
        time_of_day = compute_time_of_day(session["first_timestamp"])

        session["metrics"] = {
            "outcome": outcome,
            "session_type": session_type,
            "iteration_count": iteration_count,
            "struggle_score": 0,  # placeholder, computed after outcome is set
            "prompt_specificity": specificity,
            "time_of_day": time_of_day,
        }
        # Now compute struggle with outcome available
        session["metrics"]["struggle_score"] = compute_struggle_score(session)

        # Re-write parsed file with metrics
        out_path = PARSED_DIR / f"{session['filename']}.json"
        out_path.write_text(json.dumps(session, indent=2, ensure_ascii=False))

    completed_tasks.append("metrics_computed")

    # Step 4: Aggregated metrics.json
    print("Writing aggregated metrics.json...")
    outcome_counts = Counter(s["metrics"]["outcome"] for s in parsed_sessions)
    type_counts = Counter(s["metrics"]["session_type"] for s in parsed_sessions)
    tod_counts = Counter(s["metrics"]["time_of_day"] for s in parsed_sessions)
    date_counts = Counter(s["date"] for s in parsed_sessions)

    avg_turns = round(mean(s["turn_count"] for s in parsed_sessions), 1) if parsed_sessions else 0
    avg_duration = round(mean(s["duration_seconds"] for s in parsed_sessions), 1) if parsed_sessions else 0
    avg_spec = round(mean(s["metrics"]["prompt_specificity"] for s in parsed_sessions), 1) if parsed_sessions else 0
    avg_struggle = round(mean(s["metrics"]["struggle_score"] for s in parsed_sessions), 1) if parsed_sessions else 0

    metrics_agg = {
        "total_sessions": len(parsed_sessions),
        "by_outcome": {k: outcome_counts.get(k, 0) for k in ["resolved", "abandoned", "partial", "context-only"]},
        "by_type": {k: type_counts.get(k, 0) for k in ["bug-fix", "feature", "setup", "exploration", "refactor", "unknown"]},
        "avg_turns": avg_turns,
        "avg_duration_seconds": avg_duration,
        "avg_specificity": avg_spec,
        "avg_struggle_score": avg_struggle,
        "by_date": dict(sorted(date_counts.items())),
        "by_time_of_day": {k: tod_counts.get(k, 0) for k in ["morning", "afternoon", "evening", "night"]},
        "sessions": [
            {
                "filename": s["filename"],
                "outcome": s["metrics"]["outcome"],
                "type": s["metrics"]["session_type"],
                "turns": s["turn_count"],
                "duration": s["duration_seconds"],
                "specificity": s["metrics"]["prompt_specificity"],
                "struggle": s["metrics"]["struggle_score"],
            }
            for s in parsed_sessions
        ],
    }

    METRICS_FILE.write_text(json.dumps(metrics_agg, indent=2))
    completed_tasks.append("metrics_json_written")

    # Step 5: Final status
    write_status("completed", "done", completed_tasks, len(parsed_sessions), total, all_errors[:20])
    print(f"Done! Wrote metrics.json and {len(parsed_sessions)} parsed session files.")


if __name__ == "__main__":
    main()
