#!/usr/bin/env python3
import collections
import glob
import os
import re
import subprocess
import sys
import time
from datetime import datetime
from pathlib import Path


def now_iso():
    return datetime.now().astimezone().isoformat(timespec="seconds")


def wc_l(path):
    if not os.path.exists(path):
        return 0
    with open(path, "rb") as f:
        return sum(1 for _ in f)


def read_lines(path):
    if not os.path.exists(path):
        return []
    with open(path) as f:
        return [l.rstrip("\n") for l in f if l.strip()]


def write_lines(path, lines):
    with open(path, "w") as f:
        for l in lines:
            f.write(l + "\n")


def begin_test_case_names(qpa_path):
    names = set()
    if not os.path.exists(qpa_path):
        return names
    txt = open(qpa_path, errors="replace").read()
    for line in txt.splitlines():
        if line.startswith("#beginTestCaseResult "):
            names.add(line[len("#beginTestCaseResult "):].strip())
    return names


def main():
    lab_root = Path(os.environ.get("LAB_ROOT", Path(__file__).resolve().parent.parent))
    drivers_root = Path(os.environ.get("DRIVERS_ROOT", lab_root.parent))
    mesa_src = Path(os.environ.get("MESA_SRC", drivers_root / "intel-hard-graphics"))
    mesa_build = Path(os.environ.get("MESA_BUILD", mesa_src / "build"))
    cts_dir = Path(os.environ.get("CTS_DIR", mesa_src / "VK-GL-CTS/build-cts/external/vulkancts/modules/vulkan"))
    cts_mustpass = Path(os.environ.get("CTS_MUSTPASS", mesa_src / "VK-GL-CTS/external/vulkancts/mustpass/main/vk-default"))
    icd = Path(os.environ.get("VK_DRIVER_FILES", mesa_build / "src/intel/vulkan_hasvk/intel_hasvk13_devenv_icd.x86_64.json"))
    exclude = Path(os.environ.get("EXCLUDE", lab_root / "cts/cts-exclude-gen7.txt"))

    scope = sys.argv[1] if len(sys.argv) > 1 else "mustpass"
    chunk_size = int(sys.argv[2]) if len(sys.argv) > 2 else 100000
    pause = int(sys.argv[3]) if len(sys.argv) > 3 else 30
    outdir = Path(os.environ.get("OUTDIR", Path.cwd() / "cts-chunked"))

    stop_on_hang = os.environ.get("STOP_ON_HANG", "0") == "1"
    max_hangs = int(os.environ.get("MAX_HANGS", "3"))
    max_chunks = int(os.environ.get("MAX_CHUNKS", "3"))
    max_subruns = int(os.environ.get("MAX_SUBRUNS", "40"))
    hangs = 0

    deqp = cts_dir / "deqp-vk"
    for p in (deqp, icd):
        if not p.exists():
            print(f"missing: {p}", file=sys.stderr)
            sys.exit(1)

    outdir.mkdir(parents=True, exist_ok=True)
    os.chdir(cts_dir)

    src_txt = outdir / "src.txt"
    cases_txt = outdir / "cases.txt"
    if not (src_txt.exists() and src_txt.stat().st_size > 0):
        if scope == "all":
            out = subprocess.run([str(deqp), "--deqp-runmode=stdout-caselist"],
                                  capture_output=True, text=True).stdout
            lines = [l[len("TEST: "):] for l in out.splitlines() if l.startswith("TEST: ")]
            write_lines(src_txt, lines)
        else:
            lines = []
            for f in sorted(glob.glob(str(cts_mustpass / "*.txt"))):
                lines.extend(l for l in read_lines(f) if l.startswith("dEQP-VK"))
            write_lines(src_txt, lines)

    if exclude.exists() and exclude.stat().st_size > 0:
        excl = set(read_lines(exclude))
        write_lines(cases_txt, [l for l in read_lines(src_txt) if l not in excl])
    else:
        write_lines(cases_txt, read_lines(src_txt))

    todo_txt = outdir / "todo.txt"
    prev = sorted(glob.glob(str(outdir / "chunk-*.qpa")))
    if prev:
        done = set()
        for f in prev:
            done |= begin_test_case_names(f)
        done_txt = outdir / "done.txt"
        write_lines(done_txt, sorted(done))
        cases = read_lines(cases_txt)
        todo = [l for l in cases if l not in done]
        write_lines(todo_txt, todo)
        print(f"resume: {len(done)} recorded, {len(todo)} left")
    else:
        cases = read_lines(cases_txt)
        write_lines(todo_txt, cases)
        print(f"{len(cases)} cases to run")

    if not (todo_txt.exists() and todo_txt.stat().st_size > 0):
        print("nothing left to run — summarising")

    todo = read_lines(todo_txt)
    total = len(todo)
    n = 0
    start = 1
    ran = 0
    rc = 0
    sub = 0

    while todo and start <= total:
        n += 1
        while glob.glob(str(outdir / f"chunk-{n}.*.qpa")):
            n += 1

        list_path = outdir / f"chunk-{n}.txt"
        chunk_lines = todo[start - 1: start - 1 + chunk_size]
        write_lines(list_path, chunk_lines)
        if not chunk_lines:
            break

        print(f"chunk {n}: {len(chunk_lines)} cases (from line {start} of {total})")
        want = len(chunk_lines)
        work_path = outdir / f"chunk-{n}.work"
        write_lines(work_path, chunk_lines)
        sub = 0

        while work_path.exists() and work_path.stat().st_size > 0:
            sub += 1
            if sub > max_subruns:
                print(f"  gave up after {max_subruns} subruns")
                break

            qpa = outdir / f"chunk-{n}.{sub}.qpa"
            log = outdir / f"chunk-{n}.{sub}.log"
            env = dict(os.environ, VK_DRIVER_FILES=str(icd))
            with open(log, "w") as lf:
                proc = subprocess.run(
                    [str(deqp), f"--deqp-caselist-file={work_path}",
                     f"--deqp-log-filename={qpa}",
                     "--deqp-log-images=disable", "--deqp-log-shader-sources=disable"],
                    stdout=lf, stderr=subprocess.STDOUT, env=env)
            rc = proc.returncode

            work_lines = read_lines(work_path)
            left = len(work_lines)
            got = sum(1 for l in open(qpa, errors="replace").read().splitlines()
                      if l.startswith("#beginTestCaseResult")) if qpa.exists() else 0
            if got >= left:
                break

            done_now = begin_test_case_names(qpa)
            next_lines = [l for l in work_lines if l not in done_now]
            if len(next_lines) == left:
                culprit = work_lines[0] if work_lines else ""
                with open(outdir / "crashers.txt", "a") as cf:
                    cf.write(f"{now_iso()} crashed before recording: {culprit}\n")
                next_lines = work_lines[1:]
            write_lines(work_path, next_lines)
            print(f"  subrun {sub}: rc={rc}, {got} done, {len(next_lines)} left — restarting")
            time.sleep(pause)

        for suffix in (".work", ".work.done", ".work.next"):
            p = outdir / f"chunk-{n}{suffix}"
            if p.exists():
                p.unlink()

        done_n = 0
        bad_log = False
        for f in sorted(glob.glob(str(outdir / f"chunk-{n}.*.qpa"))):
            done_n += sum(1 for l in open(f, errors="replace").read().splitlines()
                          if l.startswith("#beginTestCaseResult"))
        for f in sorted(glob.glob(str(outdir / f"chunk-{n}.*.log"))):
            if "DeviceLost" in open(f, errors="replace").read():
                bad_log = True
        print(f"  rc={rc}, {done_n}/{want} results ({sub} subrun)")

        bad = done_n < want or bad_log
        wait_s = pause
        if bad:
            hangs += 1
            with open(outdir / "hangs.txt", "a") as hf:
                hf.write(f"{now_iso()} chunk {n} ended badly (rc={rc})\n")
            if stop_on_hang:
                print("  chunk ended badly — stopping (STOP_ON_HANG=1)")
                print("  rerun the same command to continue from here")
                break
            wait_s = pause * 4
            print(f"  chunk ended badly (#{hangs}) — cooling down {wait_s}s")
            if max_hangs > 0 and hangs >= max_hangs:
                print(f"  {hangs} bad chunks — stopping (MAX_HANGS={max_hangs})")
                break

        start += chunk_size
        ran += 1
        if max_chunks > 0 and ran >= max_chunks and start <= total:
            print()
            print(f"{ran} chunks done this boot (MAX_CHUNKS={max_chunks}) — stopping here.")
            print(f"REBOOT, then rerun the same command to continue from case {start}.")
            break
        if start <= total:
            time.sleep(wait_s)

    res = {}
    for f in sorted(glob.glob(str(outdir / "chunk-*.qpa"))):
        txt = open(f, errors="replace").read()
        for block in re.split(r"#beginTestCaseResult ", txt)[1:]:
            name = block.split("\n", 1)[0].strip()
            m = re.search(r'<Result StatusCode="(\w+)">', block)
            res[name] = m.group(1) if m else "Crash"

    counts = collections.Counter(res.values())
    print(f"\n{sum(counts.values())} results")
    for k, v in counts.most_common():
        print(f"  {k:16} {v}")

    with open(outdir / "failures.txt", "w") as fh:
        for name, status in sorted(res.items()):
            if status not in ("Pass", "NotSupported"):
                fh.write(f"{status}\t{name}\n")
    print(f"\nfailures -> {outdir}/failures.txt")


if __name__ == "__main__":
    main()