#!/usr/bin/env python3
"""Run a test suite N times and record every attempt to a history file.

Detects the result format on its own:
  - Playwright : JSON reporter output  (--reporter=json)
  - Selenium / JUnit / TestNG : JUnit XML (surefire-reports, target, test-results)

Usage:
  python collect.py --command "npx playwright test --reporter=json" --runs 5
  python collect.py --command "mvn -q test" --runs 5 --format junit
  python collect.py --command "npx playwright test --reporter=json" --runs 5 --fresh

Writes .flaky/history.jsonl, one JSON object per test attempt. Then run analyze.py.
Stdlib only. No install step.
"""
import argparse, json, os, subprocess, sys, glob, time, datetime
import xml.etree.ElementTree as ET


# ---------- playwright ----------
def parse_playwright(stdout, run):
    """Walk the Playwright JSON reporter tree. Every result (incl. retries) is an attempt."""
    start = stdout.find("{")
    if start == -1:
        return []
    try:
        data = json.loads(stdout[start:])
    except json.JSONDecodeError:
        return []
    rows = []

    def walk(suite, path):
        title = suite.get("title", "")
        here = path + [title] if title else path
        for spec in suite.get("specs", []) or []:
            file = spec.get("file") or suite.get("file") or ""
            # Playwright repeats the file as the top suite title. Drop it so the id
            # reads "tests/checkout.spec.ts > describe > test", not the name twice.
            base = os.path.basename(file)
            crumbs = [p for p in here if p and p != file and p != base]
            name = " > ".join(crumbs + [spec.get("title", "")])
            tid = f"{file} > {name}" if file else name
            for test in spec.get("tests", []) or []:
                for res in test.get("results", []) or []:
                    st = res.get("status")
                    st = {"passed": "passed", "failed": "failed", "timedOut": "failed",
                          "interrupted": "failed", "skipped": "skipped"}.get(st, st)
                    msg = ""
                    for err_key in ("error", "errors"):
                        err = res.get(err_key)
                        if isinstance(err, dict):
                            msg = (err.get("message") or "").strip()
                        elif isinstance(err, list) and err:
                            msg = (err[0].get("message") or "").strip()
                        if msg:
                            break
                    rows.append({"run": run, "framework": "playwright", "id": tid, "file": file,
                                 "name": spec.get("title", ""), "status": st,
                                 "duration_ms": res.get("duration"),
                                 "message": msg.splitlines()[0][:300] if msg else ""})
        for child in suite.get("suites", []) or []:
            walk(child, here)

    for s in data.get("suites", []) or []:
        walk(s, [])
    return rows


# ---------- junit / testng (selenium, maven surefire) ----------
JUNIT_GLOBS = ["**/surefire-reports/*.xml", "**/failsafe-reports/*.xml",
               "**/target/**/TEST-*.xml", "**/test-results/**/*.xml",
               "**/build/test-results/**/*.xml", "**/reports/**/TEST-*.xml"]


def find_junit(since):
    seen, files = set(), []
    for g in JUNIT_GLOBS:
        for f in glob.glob(g, recursive=True):
            if f in seen or not os.path.isfile(f):
                continue
            # only files this run touched, so old reports do not leak in
            if since and os.path.getmtime(f) < since - 1:
                continue
            seen.add(f)
            files.append(f)
    return files


def parse_junit(files, run):
    rows = []
    for f in files:
        try:
            root = ET.parse(f).getroot()
        except ET.ParseError:
            continue
        suites = [root] if root.tag == "testsuite" else root.iter("testsuite")
        for suite in suites:
            for tc in suite.iter("testcase"):
                cls = tc.get("classname") or suite.get("name") or ""
                name = tc.get("name") or ""
                tid = f"{cls}.{name}" if cls else name
                status, msg = "passed", ""
                if tc.find("skipped") is not None:
                    status = "skipped"
                for tag in ("failure", "error"):
                    node = tc.find(tag)
                    if node is not None:
                        status = "failed"
                        msg = (node.get("message") or (node.text or "")).strip()
                        break
                dur = tc.get("time")
                rows.append({"run": run, "framework": "selenium/junit", "id": tid, "file": cls,
                             "name": name, "status": status,
                             "duration_ms": int(float(dur) * 1000) if dur else None,
                             "message": msg.splitlines()[0][:300] if msg else ""})
    return rows


def main():
    ap = argparse.ArgumentParser(description="Run a suite N times and record every attempt.")
    ap.add_argument("--command", required=True, help='e.g. "npx playwright test --reporter=json"')
    ap.add_argument("--runs", type=int, default=5)
    ap.add_argument("--format", choices=["auto", "playwright", "junit"], default="auto")
    ap.add_argument("--history", default=".flaky/history.jsonl")
    ap.add_argument("--fresh", action="store_true", help="wipe history first (new baseline)")
    a = ap.parse_args()

    os.makedirs(os.path.dirname(a.history) or ".", exist_ok=True)
    if a.fresh and os.path.exists(a.history):
        os.remove(a.history)

    start_run = 1
    if os.path.exists(a.history):
        with open(a.history) as fh:
            runs = [json.loads(l).get("run", 0) for l in fh if l.strip()]
        start_run = (max(runs) + 1) if runs else 1

    total = 0
    for i in range(a.runs):
        run = start_run + i
        t0 = time.time()
        print(f"[{run}] {a.command}", flush=True)
        proc = subprocess.run(a.command, shell=True, capture_output=True, text=True)
        # A non-zero exit is EXPECTED here: failing tests are the signal, not an error.

        rows = []
        if a.format in ("auto", "playwright"):
            rows = parse_playwright(proc.stdout, run)
        if not rows and a.format in ("auto", "junit"):
            rows = parse_junit(find_junit(t0), run)

        if not rows:
            print(f"    no results parsed. Playwright needs --reporter=json; "
                  f"JUnit/TestNG needs XML reports on disk.", file=sys.stderr)
        else:
            with open(a.history, "a") as fh:
                for r in rows:
                    r["ts"] = datetime.datetime.now().isoformat(timespec="seconds")
                    fh.write(json.dumps(r) + "\n")
            p = sum(1 for r in rows if r["status"] == "passed")
            f = sum(1 for r in rows if r["status"] == "failed")
            total += len(rows)
            print(f"    {len(rows)} attempts  {p} passed  {f} failed  ({time.time()-t0:.1f}s)")

    print(f"\n{total} attempts appended to {a.history}")
    print("Next: python analyze.py --html flaky-report.html")


if __name__ == "__main__":
    main()