#!/usr/bin/env python3
"""
Runs musrview --png on a given msr-file, then compares the generated PNGs
against reference images using pixel-level comparison (Pillow).

Usage:
  musrview_check.py <musrview> <msr-file> <ref-dir> <test-name>
                    [--tol T] [--generate] [musrview-opts...]

Modes:
  default     Compare generated PNGs against references in <ref-dir>/<test-name>/
  --generate  Generate reference PNGs into <ref-dir>/<test-name>/ (no comparison)

Tolerance metric: mean absolute pixel difference normalised to [0, 1].
  0.0 = identical, 1.0 = maximally different.  Default tolerance: 0.01 (~1%).
"""

import argparse
import glob
import os
import shutil
import subprocess
import sys
import tempfile


def pixel_diff(img_a_path, img_b_path):
    """Return the mean absolute pixel difference normalised to [0, 1]."""
    from PIL import Image
    import numpy as np

    a = np.asarray(Image.open(img_a_path).convert("RGBA"), dtype=np.float64)
    b = np.asarray(Image.open(img_b_path).convert("RGBA"), dtype=np.float64)

    if a.shape != b.shape:
        return 1.0  # completely different dimensions

    return np.mean(np.abs(a - b)) / 255.0


def main():
    # ---- argument parsing ----------------------------------------------------
    parser = argparse.ArgumentParser(description="musrview PNG integration test")
    parser.add_argument("musrview", help="path to musrview executable")
    parser.add_argument("msr_file", help="path to msr input file")
    parser.add_argument("ref_dir", help="root reference directory")
    parser.add_argument("test_name", help="test name (subdirectory in ref_dir)")
    parser.add_argument("--tol", type=float, default=0.01,
                        help="tolerance for pixel comparison (default 0.01)")
    parser.add_argument("--generate", action="store_true",
                        help="generate reference PNGs instead of comparing")

    # everything after the known args is forwarded to musrview
    args, musrview_opts = parser.parse_known_args()

    msr_basename = os.path.splitext(os.path.basename(args.msr_file))[0]
    ref_subdir = os.path.join(args.ref_dir, args.test_name)

    # ---- run musrview in a temporary directory -------------------------------
    work_dir = os.path.dirname(os.path.abspath(args.msr_file))
    with tempfile.TemporaryDirectory() as tmp_dir:
        # snapshot existing PNGs so we only pick up newly created ones
        pre_existing = set(glob.glob(os.path.join(work_dir, f"{msr_basename}_*.png")))

        cmd = [args.musrview, args.msr_file, "--png"] + musrview_opts
        print(f"running: {' '.join(cmd)}")
        result = subprocess.run(cmd, capture_output=True, text=True,
                                cwd=work_dir,
                                env={**os.environ, "MUSRVIEW_PNG_DIR": tmp_dir})
        if result.returncode != 0:
            print(f"**ERROR** musrview returned exit code {result.returncode}")
            print(result.stdout + result.stderr)
            return 1

        # collect only newly created PNGs (exclude pre-existing ones)
        generated = sorted(
            p for p in glob.glob(os.path.join(work_dir, f"{msr_basename}_*.png"))
            if p not in pre_existing
        )
        if not generated:
            generated = sorted(glob.glob(os.path.join(tmp_dir, f"{msr_basename}_*.png")))
        if not generated:
            print(f"**ERROR** no PNGs matching '{msr_basename}_*.png' found")
            print(f"  checked: {work_dir}")
            if result.stdout:
                print(result.stdout)
            return 1

        # ---- generate mode ---------------------------------------------------
        if args.generate:
            os.makedirs(ref_subdir, exist_ok=True)
            for png in generated:
                dst = os.path.join(ref_subdir, os.path.basename(png))
                shutil.copy2(png, dst)
                print(f"  saved reference: {dst}")
            # clean up generated PNGs from work_dir
            for png in generated:
                if os.path.dirname(os.path.abspath(png)) == os.path.abspath(work_dir):
                    os.remove(png)
            print(f"GENERATE: {len(generated)} reference PNG(s) written to {ref_subdir}")
            return 0

        # ---- compare mode ----------------------------------------------------
        if not os.path.isdir(ref_subdir):
            print(f"**ERROR** reference directory not found: {ref_subdir}")
            return 1

        failures = 0
        compared = 0
        for png_path in generated:
            name = os.path.basename(png_path)
            ref_path = os.path.join(ref_subdir, name)
            if not os.path.isfile(ref_path):
                print(f"FAIL: no reference PNG for {name}")
                failures += 1
                continue

            diff = pixel_diff(png_path, ref_path)
            compared += 1
            if diff > args.tol:
                print(f"FAIL: {name}  diff={diff:.6f} > tol={args.tol:.6f}")
                failures += 1
            else:
                print(f"PASS: {name}  diff={diff:.6f} <= tol={args.tol:.6f}")

        # clean up generated PNGs from work_dir
        for png in generated:
            if os.path.dirname(os.path.abspath(png)) == os.path.abspath(work_dir):
                os.remove(png)

        if compared == 0:
            print("**ERROR** no PNGs were compared")
            return 1

        if failures:
            print(f"\n{failures} of {len(generated)} PNG(s) FAILED")
            return 1

        print(f"\nAll {compared} PNG(s) PASSED (tol={args.tol})")
        return 0


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