#!/usr/bin/env python
"""Framedump rendering test runner for PPSSPPHeadless.

Runs GE frame dumps (".ppdmp", possibly zipped) through PPSSPPHeadless with a
set of rendering variants (command line option sets), generating reference
images when missing and comparing against existing ones when present. Produces
a self-contained HTML report and returns a nonzero exit code on failure.

This script is test-set agnostic: it reads its configuration from a JSON file
(similar to frametests/frametests.json) that points at the test data tree,
so it can be used against different test sets (the CI one, or bigger private
ones on custom machines).

Usage:
    frametests.py [OPTIONS] [CONFIG.json]

Example:
    python3 frametests.py frametests/frametests.json
"""

import argparse
import base64
import glob
import html
import json
import os
import re
import shlex
import shutil
import subprocess
import platform
import sys
import time
from pathlib import Path

# test.py-style candidate paths for the headless binary, relative to the
# current working directory, in preference order.
# The machine's own architecture comes first, the same way test.py picks: an x64 build runs on
# Windows-on-ARM too, under emulation, so looking for it first quietly tests the emulated build.
HEADLESS_CANDIDATES = ([
    "Windows/ARM64/Debug/PPSSPPHeadless.exe",
    "Windows/ARM64/Release/PPSSPPHeadless.exe",
] if platform.machine().lower() in ("arm64", "aarch64") else []) + [
    "Windows/x64/Debug/PPSSPPHeadless.exe",
    "Windows/Debug/PPSSPPHeadless.exe",
    "Windows/x64/Release/PPSSPPHeadless.exe",
    "Windows/Release/PPSSPPHeadless.exe",
    "build/PPSSPPHeadless",
    "build-headless/PPSSPPHeadless",
    "build*/PPSSPPHeadless",
    "PPSSPPHeadless",
    "ppsspp/PPSSPPHeadless",
]

MSE_RE = re.compile(r"Screenshot MSE: ([0-9.eE+-]+)")
LOG_EMBED_LIMIT = 64 * 1024

STATUS_PASS = "PASS"
STATUS_FAIL = "FAIL"
STATUS_NEW = "NEW"
STATUS_ERROR = "ERROR"


def find_headless(config_dir, config_path):
    """Locate the PPSSPPHeadless binary. The config headlessPath wins, then
    the PPSSPP_HEADLESS env var, then candidate paths (newest by mtime)."""
    candidates = []
    if config_path:
        p = Path(config_path)
        candidates.append(p if p.is_absolute() else config_dir / p)
    env_path = os.environ.get("PPSSPP_HEADLESS")
    if env_path:
        candidates.append(Path(env_path))
    for pattern in HEADLESS_CANDIDATES:
        candidates.extend(Path(m) for m in glob.glob(pattern))
    found = [c.resolve() for c in candidates if c.is_file()]
    if not found:
        return None
    return max(found, key=lambda p: p.stat().st_mtime)


def strip_dump_extensions(name):
    """Derive the base name for a dump, stripping .zip and .ppdmp extensions."""
    base = name
    lower = base.lower()
    if lower.endswith(".zip"):
        base = base[:-4]
        lower = base.lower()
    if lower.endswith(".ppdmp"):
        base = base[:-6]
    return base


def collect_dumps(test_root):
    """Recursively collect frame dumps (.ppdmp files and .zip wrappers)."""
    dumps = []
    for dirpath, dirnames, filenames in os.walk(test_root):
        dirnames[:] = [d for d in sorted(dirnames) if not d.startswith(".")]
        for filename in sorted(filenames):
            lower = filename.lower()
            if lower.endswith(".ppdmp") or lower.endswith(".zip"):
                dumps.append(Path(dirpath) / filename)
    return sorted(dumps)


def run_test(headless, dump_path, variant_args, ref_path, actual_path, diff_path, max_mse, timeout, output_dir):
    """Run one dump with one variant's args. Returns (status, mse, output, timed_out)."""
    generate = not ref_path.exists()
    args = [str(headless)] + list(variant_args)
    if generate:
        args.append("--screenshot-save=" + str(ref_path))
    else:
        args.extend([
            "--screenshot-save=" + str(actual_path),
            "--screenshot=" + str(ref_path),
            "--max-mse=" + str(max_mse),
            "--screenshot-diff=" + str(diff_path),
        ])
    args.append(str(dump_path))

    timed_out = False
    try:
        proc = subprocess.Popen(
            args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
            cwd=str(output_dir), text=True, encoding="utf-8", errors="replace")
        output, _ = proc.communicate(timeout=timeout)
    except subprocess.TimeoutExpired:
        proc.kill()
        output, _ = proc.communicate()
        timed_out = True
    except OSError as e:
        # e.g. the binary couldn't be launched at all.
        return STATUS_ERROR, None, "Failed to launch headless binary: %s\n" % e, False
    returncode = proc.returncode

    # Headless may drop these in the cwd on failure (outside GITHUB_ACTIONS).
    for stray in ("__testfailure.bmp", "__testcompare.png"):
        try:
            (output_dir / stray).unlink()
        except OSError:
            pass

    match = MSE_RE.search(output)
    mse = float(match.group(1)) if match else None

    if timed_out:
        return STATUS_FAIL, mse, output, True
    if generate:
        if returncode != 0 or not ref_path.exists():
            return STATUS_ERROR, mse, output, False
        return STATUS_NEW, mse, output, False
    if returncode != 0:
        return STATUS_FAIL, mse, output, False
    return STATUS_PASS, mse, output, False


def image_data_uri(path):
    """Read an image file as a PNG data URI, or None if missing."""
    try:
        data = path.read_bytes()
    except OSError:
        return None
    return "data:image/png;base64," + base64.b64encode(data).decode("ascii")


def write_report(report_path, results, variants):
    rows = []
    for res in results:
        cls = "failed" if res["status"] in (STATUS_FAIL, STATUS_ERROR) else ("new" if res["status"] == STATUS_NEW else "")
        mse = ('<span class="mse">MSE: %.6f</span>' % res["mse"]) if res["mse"] is not None else ""
        detail = '<span class="detail">%s</span>' % html.escape(res["detail"]) if res["detail"] else ""
        ref_img = '<img class="img" src="%s" alt="reference" title="reference">' % res["ref_uri"] if res["ref_uri"] else ""
        actual_img = '<img class="img" src="%s" alt="actual" title="actual">' % res["actual_uri"] if res["actual_uri"] else ""
        diff_img = '<img class="img diff" src="%s" alt="diff" title="diff">' % res["diff_uri"] if res["diff_uri"] else ""
        log = ('<details class="log"><summary>log</summary><pre>%s</pre></details>' % html.escape(res["log"])) if res["log"] else ""
        rows.append("""        <div class="test {cls}">
          <div class="head">
            <span class="status">{status}</span>
            <span class="variant">{variant}</span>
            <span class="name">{name}</span>
            {mse}{detail}
          </div>
          <div class="images">
            {ref_img}{actual_img}{diff_img}
          </div>
          {log}
        </div>""".format(
            cls=cls,
            status=html.escape(res["status"]),
            variant=html.escape(res["variant"]),
            name=html.escape(res["dump_name"]),
            mse=mse,
            detail=detail,
            ref_img=ref_img,
            actual_img=actual_img,
            diff_img=diff_img,
            log=log,
        ))

    counts = {}
    for res in results:
        counts[res["status"]] = counts.get(res["status"], 0) + 1
    summary = ", ".join("%s: %d" % (k, v) for k, v in sorted(counts.items()))

    variant_rows = ""
    for variant in variants:
        var_counts = {}
        for res in results:
            if res["variant"] == variant:
                var_counts[res["status"]] = var_counts.get(res["status"], 0) + 1
        variant_rows += "<tr><td>%s</td><td>%s</td></tr>" % (
            html.escape(variant), html.escape(", ".join("%s: %d" % (k, v) for k, v in sorted(var_counts.items()))))

    report = """<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<title>PPSSPP frametests report</title>
<style>
  body {{ font-family: sans-serif; margin: 1em; }}
  .test {{ border: 1px solid #ccc; border-radius: 4px; margin: 0.5em 0; padding: 0.5em; }}
  .test.failed {{ border-color: #c33; background: #fee; }}
  .test.new {{ border-color: #cc3; background: #ffe; }}
  .head {{ font-weight: bold; }}
  .status {{ color: #393; }}
  .failed .status {{ color: #c33; }}
  .new .status {{ color: #c93; }}
  .variant {{ color: #666; margin: 0 0.5em; }}
  .mse {{ color: #c33; margin-left: 0.5em; }}
  .detail {{ color: #c33; margin-left: 0.5em; }}
  .images {{ margin-top: 0.5em; }}
  .img {{ width: 256px; border: 1px solid #999; margin-right: 0.5em; vertical-align: top; }}
  .img.diff {{ width: 512px; }}
  .log pre {{ white-space: pre-wrap; overflow-x: auto; max-height: 20em; }}
  .failed .img {{ border-color: #c33; }}
  table {{ border-collapse: collapse; }}
  td, th {{ border: 1px solid #ccc; padding: 0.2em 0.6em; }}
</style>
</head>
<body>
<h1>PPSSPP frametests report</h1>
<p>{summary}</p>
<h2>Variants</h2>
<table>
<tr><th>variant</th><th>results</th></tr>
{variant_rows}
</table>
<h2>Results</h2>
{rows}
</body>
</html>""".format(summary=html.escape(summary), variant_rows=variant_rows, rows="\n".join(rows))

    report_path.write_text(report, encoding="utf-8")
    return report_path


def main():
    parser = argparse.ArgumentParser(description="Run PPSSPP framedump rendering tests.")
    parser.add_argument("config", nargs="?", default="frametests/frametests.json",
                        help="Path to the JSON configuration file (default: frametests/frametests.json)")
    parser.add_argument("--filter", default="", help="Only run dumps whose path contains this substring (case-insensitive)")
    parser.add_argument("--strict", action="store_true",
                        help="Treat missing reference images as failures (configuration error)")
    parser.add_argument("--out-mode", choices=["all", "failures"], default="all",
                        help="What to keep in the output directory: everything, or only failures (default: all)")
    args = parser.parse_args()

    strict = args.strict or bool(os.environ.get("GITHUB_ACTIONS"))

    config_path = Path(args.config)
    if not config_path.is_file():
        print("ERROR: config file not found: %s" % config_path, file=sys.stderr)
        return 2
    config_dir = config_path.parent
    with open(config_path, "r", encoding="utf-8") as f:
        config = json.load(f)

    test_root = (config_dir / config["testRoot"]).resolve()
    ref_root = (config_dir / config["refRoot"]).resolve()
    output_root = (config_dir / config.get("outputRoot", "out")).resolve()
    timeout = float(config.get("timeout", 60))
    max_mse = float(config.get("maxMse", 0.0))
    raw_variants = config.get("variants", {})
    if not raw_variants:
        print("ERROR: no variants defined in %s" % config_path, file=sys.stderr)
        return 2

    # Each variant has a suffix (used for its output images) and a compare-suffix
    # (the reference image it compares against).  Reference images are only
    # generated for variants whose suffix matches their compare-suffix, so that
    # other variants can share a reference (e.g. gl comparing against the soft
    # reference).  A plain string entry means args only, with both suffixes
    # defaulting to the variant name.
    variants = {}
    for key, entry in raw_variants.items():
        if isinstance(entry, str):
            args_str = entry
            suffix = key
        else:
            args_str = entry.get("args", "")
            suffix = entry.get("suffix", key)
        compare_suffix = entry.get("compare-suffix", suffix) if not isinstance(entry, str) else suffix
        variants[key] = {"args": shlex.split(args_str), "suffix": suffix, "compare_suffix": compare_suffix}

    headless = find_headless(config_dir, config.get("headlessPath", ""))
    if headless is None:
        print("ERROR: PPSSPPHeadless binary not found. Set 'headlessPath' in %s or PPSSPP_HEADLESS, or run from the repo root." % config_path, file=sys.stderr)
        return 2

    if not test_root.is_dir():
        print("ERROR: test root not found: %s (check 'testRoot' in %s)" % (test_root, config_path), file=sys.stderr)
        return 2

    dumps = collect_dumps(test_root)
    if not dumps:
        print("ERROR: no frame dumps found under %s" % test_root, file=sys.stderr)
        return 2

    if args.filter:
        filter_lower = args.filter.lower()
        dumps = [d for d in dumps if filter_lower in str(d.relative_to(test_root)).lower()]

    log_dir = output_root / "logs"
    diff_dir = output_root / "diffs"
    actual_dir = output_root / "actuals"
    generated_dir = output_root / "generated"
    for d in (log_dir, diff_dir, actual_dir, generated_dir):
        d.mkdir(parents=True, exist_ok=True)

    print("Headless: %s" % headless)
    print("Test root: %s" % test_root)
    variant_desc = []
    for key, v in variants.items():
        if v["compare_suffix"] == v["suffix"]:
            variant_desc.append(key)
        else:
            variant_desc.append("%s (ref: %s)" % (key, v["compare_suffix"]))
    print("Variants: %s" % ", ".join(variant_desc))
    print("Running %d dumps..." % len(dumps))

    results = []
    failures = 0
    errors = 0
    new_refs = 0
    start_time = time.time()

    for dump in dumps:
        rel = dump.relative_to(test_root)
        base = strip_dump_extensions(dump.name)
        rel_dir = rel.parent
        for variant, vinfo in variants.items():
            variant_args = vinfo["args"]
            generate_refs = vinfo["suffix"] == vinfo["compare_suffix"]
            ref_path = ref_root / rel_dir / ("%s-%s.png" % (base, vinfo["compare_suffix"]))
            actual_path = actual_dir / rel_dir / ("%s-%s.png" % (base, variant))
            diff_path = diff_dir / rel_dir / ("%s-%s.png" % (base, variant))
            log_path = log_dir / rel_dir / ("%s-%s.log" % (base, variant))
            for p in (ref_path, actual_path, diff_path, log_path):
                p.parent.mkdir(parents=True, exist_ok=True)

            if not ref_path.exists() and not generate_refs:
                status = STATUS_ERROR
                mse = None
                output = ""
                timed_out = False
            else:
                status, mse, output, timed_out = run_test(
                    headless, dump, variant_args, ref_path, actual_path, diff_path,
                    max_mse, timeout, output_root)

            detail = ""
            if status == STATUS_ERROR and not ref_path.exists() and not generate_refs:
                detail = "no reference image '%s' available for this variant (references are only generated for variants whose suffix matches their compare-suffix)" % vinfo["compare_suffix"]
            elif status == STATUS_FAIL:
                if timed_out:
                    detail = "timed out"
                elif mse is None:
                    if "Unable to read screenshot" in output:
                        status = STATUS_ERROR
                        detail = "reference image could not be loaded (corrupt or unreadable)"
                    else:
                        detail = "no screenshot MSE reported"
                else:
                    detail = "MSE %.6f exceeds maximum %.6f" % (mse, max_mse)
            if status == STATUS_ERROR and not detail:
                detail = "headless failed or produced no reference image"
            if status == STATUS_NEW:
                detail = "reference image generated"

            # Track counters and emit GitHub Actions annotations.
            if status == STATUS_FAIL:
                failures += 1
                if os.environ.get("GITHUB_ACTIONS"):
                    print("::error file=%s::%s failed (%s)" % (rel, variant, detail))
            elif status == STATUS_ERROR:
                errors += 1
                if os.environ.get("GITHUB_ACTIONS"):
                    print("::error file=%s::%s errored (%s)" % (rel, variant, detail))
            elif status == STATUS_NEW:
                new_refs += 1
                gen_path = generated_dir / rel_dir / ("%s-%s.png" % (base, vinfo["compare_suffix"]))
                gen_path.parent.mkdir(parents=True, exist_ok=True)
                shutil.copyfile(ref_path, gen_path)
                if strict:
                    failures += 1
                    if os.environ.get("GITHUB_ACTIONS"):
                        print("::error file=%s::%s missing reference image (config error; generated one saved to %s)" % (rel, variant, gen_path))

            # Keep artifacts (log, actual, diff) for everything in "all" mode,
            # or only for failures/errors otherwise.
            keep_artifacts = args.out_mode == "all" or status in (STATUS_FAIL, STATUS_ERROR)
            if keep_artifacts:
                log_path.write_text(output, encoding="utf-8", errors="replace")
            else:
                for p in (log_path, actual_path, diff_path):
                    try:
                        p.unlink()
                    except OSError:
                        pass

            if status != STATUS_PASS:
                print("[%s] %s (%s)%s" % (status, rel, variant, " - " + detail if detail else ""))

            # In "failures" mode, only embed images for non-passing tests to
            # keep the report (and thus the CI artifact) small.
            ref_uri = None
            if args.out_mode == "all" or status != STATUS_PASS:
                if "could not be loaded" not in detail:
                    ref_uri = image_data_uri(ref_path)
            results.append({
                "dump_name": str(rel),
                "variant": variant,
                "status": status,
                "mse": mse,
                "detail": detail,
                "log": output if keep_artifacts and len(output) < LOG_EMBED_LIMIT else "",
                "ref_uri": ref_uri,
                "actual_uri": image_data_uri(actual_path) if actual_path.exists() else None,
                "diff_uri": image_data_uri(diff_path) if diff_path.exists() else None,
            })

    elapsed = time.time() - start_time
    print("Done in %.1fs: %d tests, %d failures, %d errors, %d new references." % (
        elapsed, len(results), failures, errors, new_refs))

    report_path = output_root / "report.html"
    write_report(report_path, results, list(variants.keys()))
    print("Report written to: %s" % report_path)

    if failures or errors:
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())