#!/usr/bin/env bash
# trace-analysis.sh — Parse opencode.log into structured agent trace JSONL
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
PROJECT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"

# --- Source platform for PYTHON_CMD ---
# shellcheck source=platform.sh
. "$SCRIPT_DIR/platform.sh"

# --- Paths ---
LOG_FILE="$HOME/.local/share/opencode/log/opencode.log"
TRACE_DIR="$PROJECT_DIR/.opencode-kit/logs"
TRACE_FILE="$TRACE_DIR/trace.jsonl"
OFFSET_FILE="$TRACE_DIR/.parse_offset"

# --- ANSI colors (use $'...' for actual ESC bytes) ---
RED=$'\033[0;31m'
GREEN=$'\033[0;32m'
YELLOW=$'\033[1;33m'
CYAN=$'\033[0;36m'
BOLD=$'\033[1m'
RESET=$'\033[0m'

# --- Helpers ---
ensure_dirs() {
  mkdir -p "$TRACE_DIR"
}

die() {
  echo "Error: $*" >&2
  exit 1
}

# ============================================================
# cmd_parse — Read new lines from opencode.log, append to trace.jsonl
# ============================================================
cmd_parse() {
  ensure_dirs

  if [ ! -f "$LOG_FILE" ]; then
    echo "opencode.log not found at: $LOG_FILE"
    exit 0
  fi

  # Determine offset (last processed line number)
  local offset=0
  if [ -f "$OFFSET_FILE" ]; then
    offset=$(cat "$OFFSET_FILE")
  fi

  # Count total lines using Python
  local total_lines
  total_lines=$("$PYTHON_CMD" -c "import sys; print(sum(1 for _ in open('$LOG_FILE')))")

  if [ "$total_lines" -le "$offset" ]; then
    echo "No new entries (offset=$offset, total=$total_lines)"
    exit 0
  fi

  echo "Parsing lines $((offset + 1))-$total_lines from opencode.log..."

  # Export paths for Python subprocess
  export TRACE_FILE PARSE_OFFSET="$offset"

  # Run Python parser via heredoc (single-quoted delimiter = no bash expansion)
  "$PYTHON_CMD" <<'PYEOF' || die "Python parse failed"
import json, re, sys, os

log_file = os.path.expanduser("~/.local/share/opencode/log/opencode.log")
trace_file = os.environ.get("TRACE_FILE", "")
offset = int(os.environ.get("PARSE_OFFSET", "0"))

# Read all lines
with open(log_file) as f:
    all_lines = f.readlines()

# Build run_state from ALL lines for correct agent context
run_state = {}
for line in all_lines:
    m = re.search(r'\bmessage=stream\s', line)
    if m:
        run_m = re.search(r'\brun=(\S+)', line)
        agent_m = re.search(r'\bagent=(\S+)', line)
        mode_m = re.search(r'\bmode=(\S+)', line)
        if run_m and agent_m and mode_m:
            run_state[run_m.group(1)] = {
                'agent': agent_m.group(1),
                'mode': mode_m.group(1),
            }

# Process only new lines
new_lines = all_lines[offset:]
entries = []

for line in new_lines:
    line = line.rstrip('\n')
    if not line:
        continue

    ts_m = re.search(r'^timestamp=(\S+)', line)
    run_m = re.search(r'\brun=(\S+)', line)
    if not ts_m or not run_m:
        continue
    ts = ts_m.group(1)
    run = run_m.group(1)

    agent = run_state.get(run, {}).get('agent', 'unknown')
    mode = run_state.get(run, {}).get('mode', 'unknown')

    # --- message=evaluated permission=... ---
    if re.search(r'\bmessage=evaluated\b', line):
        perm_m = re.search(r'\bpermission=(\S+)', line)
        action_m = re.search(r'\baction\.action=(\S+)', line)
        pattern_m = re.search(r'\bpattern="([^"]+)"', line)
        if perm_m and action_m:
            entry = {
                'ts': ts,
                'tool': perm_m.group(1),
                'action': action_m.group(1),
                'agent': agent,
                'mode': mode,
                'run': run,
            }
            if pattern_m:
                entry['pattern'] = pattern_m.group(1)
            entries.append(entry)

    # --- message=stream ... ---
    elif re.search(r'\bmessage=stream\b', line):
        agent_m = re.search(r'\bagent=(\S+)', line)
        mode_m = re.search(r'\bmode=(\S+)', line)
        model_m = re.search(r'\bmodelID=(\S+)', line)
        ses_m = re.search(r'\bsession\.id=(\S+)', line)
        if agent_m:
            entry = {
                'ts': ts,
                'event': 'stream',
                'agent': agent_m.group(1),
                'run': run,
            }
            if mode_m:
                entry['mode'] = mode_m.group(1)
            if model_m:
                entry['model'] = model_m.group(1)
            if ses_m:
                entry['session_id'] = ses_m.group(1)
            entries.append(entry)

    # --- message=created id=... ---
    elif re.search(r'\bmessage=created\b', line):
        id_m = re.search(r'\bid=(\S+)', line)
        agent_m = re.search(r'\bagent=(\S+)', line)
        if id_m:
            entry = {
                'ts': ts,
                'event': 'session_created',
                'session_id': id_m.group(1),
                'run': run,
            }
            if agent_m:
                entry['agent'] = agent_m.group(1)
            entries.append(entry)

# Append entries to trace.jsonl
if entries:
    with open(trace_file, 'a') as f:
        for e in entries:
            f.write(json.dumps(e, ensure_ascii=False) + '\n')
    print(f'Appended {len(entries)} entries to {trace_file}')
else:
    print('No parseable entries found')
PYEOF

  # Save new offset
  echo "$total_lines" >"$OFFSET_FILE"
  echo "Offset updated to $total_lines"
}

# ============================================================
# cmd_report — Generate compliance report from trace.jsonl
# ============================================================
cmd_report() {
  if [ ! -f "$TRACE_FILE" ]; then
    echo "No trace data found. Run 'parse' first."
    exit 0
  fi

  export TRACE_FILE

  "$PYTHON_CMD" <<'PYEOF' || die "Python report failed"
import json, sys, os
from collections import defaultdict

trace_file = os.environ.get("TRACE_FILE", "")
if not trace_file or not os.path.exists(trace_file):
    print("trace.jsonl not found at: " + trace_file)
    sys.exit(1)

entries = []
with open(trace_file) as f:
    for line in f:
        line = line.strip()
        if line:
            try:
                entries.append(json.loads(line))
            except json.JSONDecodeError:
                pass

if not entries:
    print("No entries in trace file.")
    sys.exit(0)

# ANSI colors (defined in Python since heredoc avoids bash expansion)
GREEN = "\033[0;32m"
YELLOW = "\033[1;33m"
RED = "\033[0;31m"
CYAN = "\033[0;36m"
BOLD = "\033[1m"
RESET = "\033[0m"

# Aggregate by run and agent
runs = defaultdict(list)
for e in entries:
    key = (e.get("run", "unknown"), e.get("agent", "unknown"))
    runs[key].append(e)

total = len(entries)

# Tool usage: count by tool + action
tool_counter = defaultdict(lambda: defaultdict(int))
for e in entries:
    if "tool" in e and "action" in e:
        tool_counter[e["tool"]][e["action"]] += 1

# Permission compliance
allow_count = sum(1 for e in entries if e.get("action") == "allow")
ask_count = sum(1 for e in entries if e.get("action") == "ask")
deny_count = sum(1 for e in entries if e.get("action") == "deny")
total_decisions = allow_count + ask_count + deny_count

# Agent dispatch
agent_entries = defaultdict(lambda: {"count": 0, "modes": defaultdict(int)})
for e in entries:
    a = e.get("agent", "unknown")
    agent_entries[a]["count"] += 1
    agent_entries[a]["modes"][e.get("mode", "unknown")] += 1

# State transitions per (run, agent)
state_chains = {}
for (run, agent), evts in sorted(runs.items()):
    evts.sort(key=lambda x: x.get("ts", ""))
    seen = set()
    states = []
    for e in evts:
        if "tool" in e:
            tool = e["tool"]
            if tool.startswith("lean-ctx_ctx_read"):
                s = "READ"
            elif tool.startswith("lean-ctx_ctx_shell") or tool == "bash":
                s = "EXEC"
            elif tool == "task":
                s = "DISPATCH"
            elif tool.startswith("lean-ctx_ctx_search"):
                s = "SEARCH"
            elif tool.startswith("lean-ctx_ctx_edit") or tool == "write" or tool == "edit":
                s = "WRITE"
            elif tool in ("read",):
                s = "READ"
            else:
                s = tool.upper()[:10]
        elif e.get("event") == "session_created":
            s = "INIT"
        elif e.get("event") == "stream":
            s = "STREAM"
        else:
            s = "?"
        if s not in seen or s in ("STREAM",):
            states.append(s)
            if s not in ("?", "STREAM"):
                seen.add(s)
    if states:
        state_chains[(run, agent)] = states

# ── Print Report ──
print(f"\n{BOLD}📊 Execution Trace Report{RESET}")
print(f"{'═' * 45}")

for (run, agent), evts in sorted(runs.items()):
    n = len(evts)
    print(f"\n{CYAN}Run:{RESET} {run}  {CYAN}Agent:{RESET} {agent}  {CYAN}Entries:{RESET} {n}")

# Tool usage header
print(f"\n{BOLD}Tool Usage:{RESET}")
sorted_tools = sorted(tool_counter.items(), key=lambda x: -sum(x[1].values()))
for tool, actions in sorted_tools:
    parts = [f"{a}: {c}" for a, c in sorted(actions.items())]
    total_tool = sum(actions.values())
    print(f"  {tool:<30s} {total_tool:>4}x ({', '.join(parts)})")

# Compliance
if total_decisions > 0:
    decisions_pct = round(allow_count / total_decisions * 100) if total_decisions else 0
    ask_pct = round(ask_count / total_decisions * 100) if total_decisions else 0
    deny_pct = round(deny_count / total_decisions * 100) if total_decisions else 0
    print(f"\n{BOLD}Permission Compliance:{RESET}")
    print(f"  {GREEN}✅ allow: {allow_count:>4} ({decisions_pct:>2}%){RESET}")
    print(f"  {YELLOW}⚠️  ask:   {ask_count:>4} ({ask_pct:>2}%){RESET}")
    print(f"  {RED}❌ deny:  {deny_count:>4} ({deny_pct:>2}%){RESET}")

# Agent dispatch
print(f"\n{BOLD}Agent Dispatch:{RESET}")
for agent, info in sorted(agent_entries.items(), key=lambda x: -x[1]["count"]):
    modes_str = ", ".join(f"{m}" for m in info["modes"])
    print(f"  {agent:<20s} {info['count']:>4} entries ({modes_str})")

# State transitions
print(f"\n{BOLD}State Transitions:{RESET}")
for (run, agent), states in sorted(state_chains.items()):
    chain = " → ".join(states[:10])
    if len(states) > 10:
        chain += " → ..."
    print(f"  {agent:<15s} [{run:<8s}] {chain}")

# Summary
print(f"\n{BOLD}Summary:{RESET}")
print(f"  Total entries: {total}")
all_runs = set(e.get("run", "") for e in entries)
all_agents = set(e.get("agent", "") for e in entries)
print(f"  Unique runs:   {len(all_runs)}")
print(f"  Unique agents: {len(all_agents)}")
events = set(e.get("event", "") for e in entries if "event" in e)
events.discard("")
if events:
    print(f"  Event types:   {', '.join(sorted(events))}")
print()
PYEOF
}

# ============================================================
# cmd_clear — Clear trace data
# ============================================================
cmd_clear() {
  if [ -f "$TRACE_FILE" ]; then
    : >"$TRACE_FILE"
    echo "Cleared $TRACE_FILE"
  fi
  if [ -f "$OFFSET_FILE" ]; then
    rm "$OFFSET_FILE"
    echo "Removed $OFFSET_FILE"
  fi
  echo "Trace data cleared."
}

# ============================================================
# Main dispatch
# ============================================================
main() {
  local cmd="${1:-help}"

  case "$cmd" in
    parse)
      cmd_parse
      ;;
    report)
      cmd_report
      ;;
    clear)
      cmd_clear
      ;;
    help|--help|-h)
      echo "Usage: $0 <command>"
      echo ""
      echo "Commands:"
      echo "  parse   Parse opencode.log and append new entries to trace.jsonl"
      echo "  report  Generate compliance report from trace.jsonl"
      echo "  clear   Clear trace.jsonl"
      ;;
    *)
      die "Unknown command: $cmd. Use: parse | report | clear"
      ;;
  esac
}

main "$@"