#!/usr/bin/env python3
"""
AtomCode Datalog Analyzer — Post-hoc evaluation of agent performance.

Scans datalog/*.md files, detects antipatterns, and generates an evaluation
report with issues ranked by severity.

Usage:
    python3 scripts/analyze_datalogs.py /path/to/datalog/
    python3 scripts/analyze_datalogs.py /path/to/datalog/2026-03-25_10-42-58.md
    python3 scripts/analyze_datalogs.py --deep /path/to/datalog/   # Claude Code deep analysis on FAIL turns
"""

import json
import re
import sys
import os
import subprocess
from dataclasses import dataclass, field
from pathlib import Path


@dataclass
class Issue:
    severity: str  # HIGH, MEDIUM, LOW, BUG
    rule: str
    message: str
    steps: list = field(default_factory=list)


@dataclass
class TurnAnalysis:
    file: str
    user_prompt: str
    total_steps: int
    duration_secs: float
    issues: list = field(default_factory=list)
    files_read: dict = field(default_factory=dict)   # filename -> count
    files_edited: dict = field(default_factory=dict)  # filename -> count
    tools_used: dict = field(default_factory=dict)    # tool -> count
    bash_commands: list = field(default_factory=list)
    round_ms: list = field(default_factory=list)      # per-round LLM latency (ms)
    tool_ms: list = field(default_factory=list)       # (tool, dur_ms) per tool call


def _arg_path(args: str) -> str:
    """Best-effort target path from a tool-call's JSON args (current datalogs record
    args as JSON like `{"path":"src/x.rs"}` / `{"file_path":"…"}`). Falls back to a
    trailing-segment heuristic for non-JSON/legacy args. Returns the basename."""
    raw = args
    try:
        obj = json.loads(args)
        if isinstance(obj, dict):
            for key in ("path", "file_path", "file", "target_file", "filename"):
                if isinstance(obj.get(key), str):
                    raw = obj[key]
                    break
    except (ValueError, TypeError):
        pass
    raw = raw.rsplit("/", 1)[-1] if "/" in raw else raw
    return re.sub(r'\s+L\d+-\d+$', '', raw)


def parse_datalog(path: str) -> dict:
    """Parse a datalog .md file (current grammar emitted by datalog.rs) into
    structured data.

    Grammar (per turn file):
      - `## User\\n```\\n<prompt>\\n``` `
      - per round: `### Turn N` + `  _[request: …]_`, then tool-call bullets
        `- <tool> `<args>`` and a per-round line `  _[tokens: … dur=Nms]_`
        (or `  _[dur=Nms]_` when usage is absent).
      - per tool: ``**Tool result:** `<tool>` (`<id>`, <status>[, dur=Nms])`` then
        a fenced output block.
      - `**Stats:** N turns, M tool calls, X.Xs, …` at the end.
    A "step" here is one executed tool call (tool + args + output), reconstructed by
    pairing tool-call bullets with tool-result blocks in execution order.
    """
    text = open(path).read()

    # Extract user prompt.
    user_match = re.search(r'## User\n```\n(.*?)\n```', text, re.DOTALL)
    user_prompt = user_match.group(1).strip() if user_match else ""

    # Stats: "N turns, M tool calls, X.Xs" (only `duration` is load-bearing here;
    # `total_steps` is derived from the reconstructed steps below).
    stats_match = re.search(
        r'\*\*Stats:\*\*\s+\d+\s+turns?,\s+\d+\s+tool calls?,\s+([\d.]+)s', text
    )
    duration = float(stats_match.group(1)) if stats_match else 0.0

    # Tool-call bullets: `- <tool> `<args>`` (backticks inside args are escaped as \`).
    call_pat = re.compile(r'^- (\w[\w./-]*) `(.*)`\s*$', re.MULTILINE)
    calls = [(m.group(1), m.group(2).replace('\\`', '`')) for m in call_pat.finditer(text)]

    # Tool-result blocks (authoritative per executed tool): name, status, optional
    # duration, fenced output.
    result_pat = re.compile(
        r'\*\*Tool result:\*\* `([^`]+)` \(`[^`]*`, (\w+)(?:, dur=(\d+)ms)?\)\n```\n(.*?)\n```',
        re.DOTALL,
    )
    results = [
        {
            "tool": m.group(1),
            "status": m.group(2),
            "dur_ms": int(m.group(3)) if m.group(3) else None,
            "output": m.group(4).strip(),
        }
        for m in result_pat.finditer(text)
    ]

    # Reconstruct steps by pairing each result with its call's args in execution order.
    steps = []
    for i, r in enumerate(results):
        if i < len(calls) and calls[i][0] == r["tool"]:
            args = calls[i][1]
        else:  # order drifted (rare); fall back to the first same-named call.
            args = next((a for (n, a) in calls if n == r["tool"]), "")
        steps.append({
            "num": i + 1,
            "tool": r["tool"],
            "args": args,
            "output": r["output"],
            "status": r["status"],
            "dur_ms": r["dur_ms"],
        })

    # Per-round LLM latency: `_[tokens: … dur=Nms]_` or standalone `_[dur=Nms]_`.
    round_ms = [int(x) for x in re.findall(r'_\[(?:tokens:[^\]]*?)?dur=(\d+)ms\]_', text)]
    # Per-tool duration from the result blocks.
    tool_ms = [(s["tool"], s["dur_ms"]) for s in steps if s["dur_ms"] is not None]

    resp_match = re.search(r'\*\*Response:\*\*\n(.*?)(?=\n---|\Z)', text, re.DOTALL)
    response = resp_match.group(1).strip() if resp_match else ""

    return {
        "file": os.path.basename(path),
        "user_prompt": user_prompt,
        "total_steps": len(steps),
        "duration": duration,
        "steps": steps,
        "round_ms": round_ms,
        "tool_ms": tool_ms,
        "response": response,
    }


def analyze(data: dict) -> TurnAnalysis:
    """Run all antipattern rules against a parsed datalog."""
    t = TurnAnalysis(
        file=data["file"],
        user_prompt=data["user_prompt"],
        total_steps=data["total_steps"],
        duration_secs=data["duration"],
        round_ms=data.get("round_ms", []),
        tool_ms=data.get("tool_ms", []),
    )

    steps = data["steps"]
    sleep_count = 0
    bash_cmds = {}  # normalized cmd -> [step_nums]

    for s in steps:
        tool = s["tool"]
        args = s["args"]
        output = s["output"]
        num = s["num"]

        # Track tool usage
        t.tools_used[tool] = t.tools_used.get(tool, 0) + 1

        # Track file reads
        if tool == "read_file":
            fname = _arg_path(args)
            t.files_read[fname] = t.files_read.get(fname, 0) + 1

        # Track file edits
        if tool in ("edit_file", "search_replace"):
            fname = _arg_path(args)
            t.files_edited[fname] = t.files_edited.get(fname, 0) + 1

        # Track file writes (overwrites)
        if tool == "write_file":
            fname = _arg_path(args)
            if "Overwrote" in output or "Overwrote" in args:
                t.issues.append(Issue(
                    "HIGH", "write-existing-file",
                    f"Step {num}: write_file overwrote existing file {fname}",
                    [num],
                ))

        # Track bash commands
        if tool == "bash":
            cmd = args
            t.bash_commands.append((num, cmd))

            # Sleep detection
            if cmd.startswith("sleep ") or "&& sleep " in cmd or "; sleep " in cmd:
                sleep_count += 1

            # Normalize for repeat detection
            norm = re.sub(r'2>&1|2>/dev/null|> /\S+', '', cmd).strip()
            norm = re.sub(r'\s+', ' ', norm)
            if norm not in bash_cmds:
                bash_cmds[norm] = []
            bash_cmds[norm].append(num)

            # Bash reading files
            first_word = cmd.split()[0] if cmd else ""
            if first_word in ("cat", "head", "tail", "sed", "awk", "wc"):
                t.issues.append(Issue(
                    "LOW", "bash-read-file",
                    f"Step {num}: used bash '{first_word}' instead of read_file/grep tool",
                    [num],
                ))

            # Check for foreground server start
            if ("run dev" in cmd or "spring-boot:run" in cmd or "flask run" in cmd) \
                    and "nohup" not in cmd and "&" not in cmd:
                t.issues.append(Issue(
                    "MEDIUM", "foreground-server",
                    f"Step {num}: started server in foreground (should use nohup/&)",
                    [num],
                ))

        # Edit failures
        if tool in ("edit_file", "search_replace") and (
            "old_string not found" in output or output.startswith("x ") or output.startswith("✗")
        ):
            t.issues.append(Issue(
                "MEDIUM", "edit-failed",
                f"Step {num}: edit_file failed (old_string not found)",
                [num],
            ))

        # Rename failures
        if "Failed to rename" in output:
            t.issues.append(Issue(
                "BUG", "rename-failed",
                f"Step {num}: atomic rename failed (dev server file lock?)",
                [num],
            ))

        # Loop/block detection
        if "BLOCKED" in output:
            t.issues.append(Issue(
                "WARN", "blocked",
                f"Step {num}: tool call blocked by loop detection",
                [num],
            ))
        if "Loop in cleanup" in output or "force-terminated" in output.lower():
            t.issues.append(Issue(
                "HIGH", "force-terminated",
                f"Step {num}: turn force-terminated due to loop",
                [num],
            ))

    # --- Aggregate rules ---

    # Re-read same file 3+ times
    for fname, count in t.files_read.items():
        if count >= 3:
            t.issues.append(Issue(
                "HIGH", "re-read",
                f"'{fname}' read {count} times (expected: 1-2)",
            ))
        elif count >= 2:
            t.issues.append(Issue(
                "LOW", "re-read",
                f"'{fname}' read {count} times",
            ))

    # Same file edited 3+ times
    for fname, count in t.files_edited.items():
        if count >= 3:
            t.issues.append(Issue(
                "HIGH", "multi-edit",
                f"'{fname}' edited {count} times (should be 1-2 comprehensive edits)",
            ))

    # Sleep polling
    if sleep_count >= 2:
        t.issues.append(Issue(
            "MEDIUM", "sleep-loop",
            f"{sleep_count} sleep commands this turn (polling antipattern)",
        ))

    # Repeated bash commands
    for cmd, step_nums in bash_cmds.items():
        if len(step_nums) >= 3:
            t.issues.append(Issue(
                "HIGH", "repeated-command",
                f"Same command executed {len(step_nums)} times at steps {step_nums}: '{cmd[:60]}'",
                step_nums,
            ))
        elif len(step_nums) >= 2:
            t.issues.append(Issue(
                "LOW", "repeated-command",
                f"Same command executed {len(step_nums)} times at steps {step_nums}: '{cmd[:60]}'",
                step_nums,
            ))

    # No final verification
    if t.total_steps > 3:
        last_tools = [s["tool"] for s in steps[-3:]] if len(steps) >= 3 else []
        has_verify = any(t == "bash" for t in last_tools)
        if not has_verify and t.total_steps > 5:
            t.issues.append(Issue(
                "MEDIUM", "no-final-verify",
                "No build/curl verification in the last 3 steps",
            ))

    # Read but not edited (wasted reads)
    if t.files_edited:
        read_not_edited = set(t.files_read.keys()) - set(t.files_edited.keys())
        # Filter out common reference files
        ignore = {"index.ts", "api.ts", "router", "App.vue", "style.css"}
        wasted = [f for f in read_not_edited if f not in ignore and t.files_read.get(f, 0) >= 2]
        if wasted:
            t.issues.append(Issue(
                "LOW", "read-not-edited",
                f"Files read but never edited: {', '.join(wasted)}",
            ))

    # Step count check
    prompt_lower = t.user_prompt.lower()
    if any(kw in prompt_lower for kw in ["启动", "start", "install", "安装"]):
        expected = 4
    elif any(kw in prompt_lower for kw in ["修复", "fix", "bug", "改", "调整"]):
        expected = 6
    elif any(kw in prompt_lower for kw in ["实现", "功能", "feature", "创建"]):
        expected = 20
    else:
        expected = 10

    if t.total_steps > expected * 2:
        t.issues.append(Issue(
            "HIGH", "step-overrun",
            f"{t.total_steps} steps (expected ~{expected} for this task type, 2x over)",
        ))
    elif t.total_steps > expected * 1.5:
        t.issues.append(Issue(
            "MEDIUM", "step-overrun",
            f"{t.total_steps} steps (expected ~{expected}, 1.5x over)",
        ))

    # Sort issues: HIGH > MEDIUM > LOW > BUG > WARN
    severity_order = {"HIGH": 0, "BUG": 1, "MEDIUM": 2, "WARN": 3, "LOW": 4}
    t.issues.sort(key=lambda i: severity_order.get(i.severity, 5))

    return t


def format_report(analyses: list) -> str:
    """Generate a markdown evaluation report."""
    lines = []
    lines.append("# AtomCode Evaluation Report\n")

    # Summary
    total = len(analyses)
    high_issues = sum(1 for a in analyses if any(i.severity == "HIGH" for i in a.issues))
    avg_steps = sum(a.total_steps for a in analyses) / total if total else 0
    avg_duration = sum(a.duration_secs for a in analyses) / total if total else 0
    total_issues = sum(len(a.issues) for a in analyses)

    lines.append("## Summary\n")
    lines.append(f"| Metric | Value |")
    lines.append(f"|--------|-------|")
    lines.append(f"| Turns analyzed | {total} |")
    lines.append(f"| Turns with HIGH issues | {high_issues}/{total} |")
    lines.append(f"| Total issues found | {total_issues} |")
    lines.append(f"| Avg steps/turn | {avg_steps:.1f} |")
    lines.append(f"| Avg duration/turn | {avg_duration:.0f}s |")
    lines.append("")

    # Issue frequency
    from collections import Counter
    rule_counts = Counter()
    for a in analyses:
        for i in a.issues:
            rule_counts[f"[{i.severity}] {i.rule}"] += 1
    if rule_counts:
        lines.append("## Issue Frequency\n")
        lines.append("| Issue | Count |")
        lines.append("|-------|-------|")
        for rule, count in rule_counts.most_common():
            lines.append(f"| {rule} | {count} |")
        lines.append("")

    # Per-turn details
    lines.append("## Per-Turn Details\n")
    for a in analyses:
        status = "PASS" if not any(i.severity == "HIGH" for i in a.issues) else "FAIL"
        lines.append(f"### {a.file} — {status}")
        lines.append(f"- **Prompt:** \"{a.user_prompt[:80]}\"")
        lines.append(f"- **Steps:** {a.total_steps} | **Duration:** {a.duration_secs:.0f}s")

        # Request-level timing (from the per-round / per-tool `dur=` markers).
        timing = []
        if a.round_ms:
            timing.append(
                f"{len(a.round_ms)} rounds, slowest {max(a.round_ms)}ms, total {sum(a.round_ms)}ms"
            )
        timed_tools = [(name, ms) for (name, ms) in a.tool_ms if ms is not None]
        if timed_tools:
            slow_tool, slow_ms = max(timed_tools, key=lambda x: x[1])
            timing.append(f"slowest tool `{slow_tool}` {slow_ms}ms")
        if timing:
            lines.append(f"- **Timing:** {'; '.join(timing)}")

        if a.files_read:
            reads = ", ".join(f"{f}({c}x)" for f, c in sorted(a.files_read.items(), key=lambda x: -x[1]) if c > 1)
            if reads:
                lines.append(f"- **Re-reads:** {reads}")

        if a.issues:
            lines.append(f"- **Issues ({len(a.issues)}):**")
            for i in a.issues:
                lines.append(f"  - [{i.severity}] {i.rule}: {i.message}")
        else:
            lines.append(f"- No issues detected")
        lines.append("")

    return "\n".join(lines)


DEEP_ANALYSIS_PROMPT = """You are evaluating an AI coding agent called "atomcode".
Below is a log of one turn: the user's request and every tool call the agent made.
Also included are rule-based issues already detected.

Your job:
1. Compare with how Claude Code would handle the same task (fewer steps, better decisions).
2. Identify the ROOT CAUSE of each failure (model reasoning? framework limitation? prompt issue?).
3. Suggest ONE specific, actionable improvement to atomcode's framework or system prompt.
4. Rate overall quality: A (Claude Code level), B (acceptable), C (needs work), D (failure).

Be concise. Output format:

## Grade: X
## Claude Code would: (2-3 sentences)
## Root causes: (bullet list)
## Suggested fix: (one specific change)

---

### Rule-based issues already found:
{issues}

### Full datalog:
{datalog}
"""


def deep_analyze(analysis: TurnAnalysis, datalog_path: str) -> str:
    """Call claude -p for deep analysis of a FAIL turn."""
    datalog_content = open(datalog_path).read()
    issues_text = "\n".join(f"- [{i.severity}] {i.rule}: {i.message}" for i in analysis.issues)

    prompt = DEEP_ANALYSIS_PROMPT.format(
        issues=issues_text or "(none)",
        datalog=datalog_content,
    )

    try:
        result = subprocess.run(
            ["claude", "-p", prompt],
            capture_output=True,
            text=True,
            timeout=120,
        )
        if result.returncode == 0 and result.stdout.strip():
            return result.stdout.strip()
        else:
            return f"(claude -p failed: {result.stderr.strip()[:200]})"
    except FileNotFoundError:
        return "(claude CLI not found — install Claude Code to enable deep analysis)"
    except subprocess.TimeoutExpired:
        return "(claude -p timed out after 120s)"


def main():
    deep_mode = "--deep" in sys.argv
    args = [a for a in sys.argv[1:] if not a.startswith("--")]

    if not args:
        print("Usage: python3 analyze_datalogs.py [--deep] <path-to-datalog-dir-or-file>")
        sys.exit(1)

    target = args[0]
    files = []

    if os.path.isdir(target):
        files = sorted(Path(target).glob("*.md"))
    elif os.path.isfile(target):
        files = [Path(target)]
    else:
        print(f"Error: {target} not found")
        sys.exit(1)

    analyses = []
    file_paths = {}  # analysis -> original file path
    for f in files:
        try:
            data = parse_datalog(str(f))
            if data["total_steps"] == 0:
                continue
            analysis = analyze(data)
            analyses.append(analysis)
            file_paths[id(analysis)] = str(f)
        except Exception as e:
            print(f"Warning: failed to parse {f}: {e}", file=sys.stderr)

    if not analyses:
        print("No valid datalogs found.")
        sys.exit(0)

    report = format_report(analyses)
    print(report)

    # Deep analysis on FAIL turns
    if deep_mode:
        fail_turns = [a for a in analyses if any(i.severity == "HIGH" for i in a.issues)]
        if not fail_turns:
            print("\n## Deep Analysis\nNo FAIL turns to analyze.")
            return

        print(f"\n## Deep Analysis (Claude Code)\n")
        print(f"Analyzing {len(fail_turns)} FAIL turns...\n")

        for a in fail_turns:
            path = file_paths.get(id(a), "")
            print(f"### {a.file}")
            print(f"**Prompt:** \"{a.user_prompt[:80]}\" | **Steps:** {a.total_steps}\n")
            result = deep_analyze(a, path)
            print(result)
            print()


if __name__ == "__main__":
    main()