#!/usr/bin/env python3
"""Agrega tempo de stage (log do pipeline) + custo VLM (vlm_usage.jsonl) num JSON de run.

Uso:
  ./venv/bin/python scripts/pipeline_run_metrics.py \
      --log path/to/recut-v02-full.log \
      --workdir path/to/cut-v0.2-dev \
      --label v0.2 \
      --out path/to/run_metrics.json
"""
from __future__ import annotations

import argparse, json, os, re, sys
from datetime import datetime, timezone


STAGE_OK = re.compile(r"^\s*✓\s*OK\s*\((\d+)s\)\s*$")
STAGE_FAIL = re.compile(r"^\s*✗\s*(TIMEOUT|FALHOU).*?\((\d+)s")
STAGE_BANNER = re.compile(r"^\s{2}(\d|3|4|5|6|7|8|[a-zA-Z]).*")  # desc line after ====


def parse_log(path: str) -> dict:
    stages = []
    order = []
    pending_desc = None
    start_line = None
    if not os.path.exists(path):
        return {"stages": [], "wall_s_from_log": 0, "log": path}
    with open(path, errors="replace") as f:
        lines = f.readlines()
    for i, line in enumerate(lines):
        if line.startswith("START "):
            start_line = line.strip()
        # description is the line after ==== that is indented with 2 spaces and not cmd list
        if line.startswith("===="):
            # next non-empty non-cmd line is desc
            for j in range(i + 1, min(i + 5, len(lines))):
                s = lines[j].rstrip("\n")
                if s.startswith("  [") or s.startswith("  './") or s.startswith("  '/"):
                    continue
                if s.startswith("  ") and not s.startswith("  ✓") and not s.startswith("  ✗"):
                    pending_desc = s.strip()
                    break
            continue
        m = STAGE_OK.match(line.rstrip("\n"))
        if m and pending_desc:
            dt = int(m.group(1))
            stages.append({"desc": pending_desc, "dt_s": dt, "ok": True})
            order.append(pending_desc)
            pending_desc = None
            continue
        m = STAGE_FAIL.search(line.rstrip("\n"))
        if m and pending_desc:
            dt = int(m.group(2)) if m.lastindex and m.lastindex >= 2 else 0
            # try group
            mm = re.search(r"\((\d+)s", line)
            dt = int(mm.group(1)) if mm else 0
            stages.append({"desc": pending_desc, "dt_s": dt, "ok": False, "raw": line.strip()})
            pending_desc = None
    wall = sum(s["dt_s"] for s in stages)
    return {
        "log": path,
        "start_banner": start_line,
        "stages": stages,
        "wall_s_from_log": wall,
        "wall_min_from_log": round(wall / 60.0, 2),
    }


def parse_vlm_ledger(path: str) -> dict:
    if not os.path.exists(path):
        return {"path": path, "calls": 0, "exists": False}
    rows = []
    with open(path, errors="replace") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except Exception:
                pass
    tin = sum(int(r.get("prompt_tokens") or 0) for r in rows)
    tout = sum(int(r.get("completion_tokens") or 0) for r in rows)
    cost = sum(float(r.get("cost_usd_est") or 0) for r in rows)
    el = sum(float(r.get("elapsed_s") or 0) for r in rows)
    by_kind = {}
    for r in rows:
        k = r.get("kind") or "?"
        by_kind.setdefault(k, {"calls": 0, "prompt_tokens": 0, "completion_tokens": 0, "cost_usd_est": 0.0})
        by_kind[k]["calls"] += 1
        by_kind[k]["prompt_tokens"] += int(r.get("prompt_tokens") or 0)
        by_kind[k]["completion_tokens"] += int(r.get("completion_tokens") or 0)
        by_kind[k]["cost_usd_est"] += float(r.get("cost_usd_est") or 0)
    return {
        "path": path,
        "exists": True,
        "calls": len(rows),
        "prompt_tokens": tin,
        "completion_tokens": tout,
        "cost_usd_est": round(cost, 4),
        "vlm_elapsed_s_sum": round(el, 1),
        "by_kind": by_kind,
    }


def estimate_if_no_ledger(log_info: dict, workdir: str) -> dict:
    """Estimativa bruta se o ledger VLM não existir (processo já em voo sem hook)."""
    # v0.1 storyboard: ~226 janelas × 1 call; v0.2: × 3
    # use ~1050 prompt tokens / image from earlier probe
    story_dt = 0
    for s in log_info.get("stages") or []:
        if "STORYBOARD" in s.get("desc", "").upper() or "storyboard" in s.get("desc", "").lower():
            story_dt = s["dt_s"]
    # detect viterbi from env note in workdir
    v02 = os.path.exists(os.path.join(workdir, "V0.2_RUN.txt"))
    n_win = 226
    calls = n_win * (3 if v02 else 1)
    # judges: rough from log later
    pin, pout = 0.00015, 0.0025
    tin = calls * 1050
    tout = calls * 25
    cost = (tin / 1000) * pin + (tout / 1000) * pout
    return {
        "estimated": True,
        "note": "Sem vlm_usage.jsonl completo — estimativa por contagem de janelas (probe ~1050 in / 25 out tokens)",
        "assumed_windows": n_win,
        "assumed_frames_per_window": 3 if v02 else 1,
        "assumed_calls": calls,
        "prompt_tokens_est": tin,
        "completion_tokens_est": tout,
        "cost_usd_est": round(cost, 4),
        "storyboard_stage_s": story_dt,
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--log", required=True)
    ap.add_argument("--workdir", required=True)
    ap.add_argument("--label", default="run")
    ap.add_argument("--out", default="")
    args = ap.parse_args()

    log_info = parse_log(args.log)
    ledger_path = os.path.join(args.workdir, "vlm_usage.jsonl")
    vlm = parse_vlm_ledger(ledger_path)
    out = {
        "label": args.label,
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "workdir": args.workdir,
        "pipeline": log_info,
        "vlm": vlm,
    }
    if not vlm.get("calls"):
        out["vlm_estimate"] = estimate_if_no_ledger(log_info, args.workdir)

    text = json.dumps(out, indent=2, ensure_ascii=False)
    print(text)
    out_path = args.out or os.path.join(args.workdir, "run_metrics.json")
    with open(out_path, "w") as f:
        f.write(text)
    print(f"\nwrote {out_path}", file=sys.stderr)


if __name__ == "__main__":
    main()
