#!/usr/bin/env python3
import argparse
import json
import os
import struct
import subprocess
import sys
import tempfile


PALETTE = {
    "background": [16, 16, 24, 255],
    "foreground": [240, 48, 96, 255],
    "compute_to_draw": [48, 224, 128, 255],
    "offscreen_sampled": [64, 128, 240, 255],
    "mrt": [240, 64, 64, 255],
    "indirect_instance": [64, 240, 64, 255],
    "offscreen": [64, 64, 240, 255],
}


def _read_token(f):
    token = bytearray()
    while True:
        c = f.read(1)
        if not c:
            return None
        if c == b"#":
            f.readline()
            continue
        if not c.isspace():
            token.extend(c)
            break
    while True:
        c = f.read(1)
        if not c or c.isspace():
            break
        token.extend(c)
    return token.decode("ascii")


def load_ppm(path):
    with open(path, "rb") as f:
        magic = _read_token(f)
        if magic != "P6":
            raise ValueError("unsupported PPM magic")
        width = int(_read_token(f))
        height = int(_read_token(f))
        max_value = int(_read_token(f))
        if max_value != 255:
            raise ValueError("unsupported PPM max value")
        data = f.read(width * height * 3)
    if len(data) != width * height * 3:
        raise ValueError("truncated PPM image")
    rgba = bytearray(width * height * 4)
    for i in range(width * height):
        rgba[i * 4:i * 4 + 3] = data[i * 3:i * 3 + 3]
        rgba[i * 4 + 3] = 255
    return width, height, bytes(rgba)


def load_with_pillow(path):
    from PIL import Image
    img = Image.open(path).convert("RGBA")
    return img.width, img.height, img.tobytes()


def load_with_magick(path):
    result = subprocess.run(
        ["magick", path, "rgba:-"],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    if result.returncode != 0:
        raise ValueError(result.stderr.decode("utf-8", "replace").strip() or "magick failed")
    size = subprocess.run(
        ["magick", "identify", "-format", "%w %h", path],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
    )
    if size.returncode != 0:
        raise ValueError(size.stderr.strip() or "magick identify failed")
    width, height = [int(v) for v in size.stdout.split()]
    return width, height, result.stdout


def load_image(path):
    lower = path.lower()
    if lower.endswith(".ppm"):
        return load_ppm(path)
    try:
        return load_with_pillow(path)
    except Exception:
        return load_with_magick(path)


def color_for(value):
    if isinstance(value, str):
        if value not in PALETTE:
            raise ValueError("unknown palette color: " + value)
        return PALETTE[value]
    if isinstance(value, list) and len(value) in (3, 4):
        out = [int(v) for v in value]
        if len(out) == 3:
            out.append(255)
        return out
    raise ValueError("invalid expected color")


def pixel_at(width, data, x, y):
    idx = (y * width + x) * 4
    return list(data[idx:idx + 4])


def within(actual, expected, tolerance):
    return all(abs(int(a) - int(e)) <= tolerance for a, e in zip(actual, expected))


def evaluate_assertion(width, height, data, assertion):
    kind = assertion.get("kind")
    expected = color_for(assertion.get("expected"))
    tolerance = int(assertion.get("tolerance", 0))
    result = {
        "kind": kind,
        "expected": expected,
        "tolerance": tolerance,
        "pass": False,
    }
    if kind == "point":
        x = int(assertion["x"])
        y = int(assertion["y"])
        result["x"] = x
        result["y"] = y
        if x < 0 or y < 0 or x >= width or y >= height:
            result["reason"] = "point outside image"
            return result
        actual = pixel_at(width, data, x, y)
        result["actual"] = actual
        result["pass"] = within(actual, expected, tolerance)
        if not result["pass"]:
            result["reason"] = "point color mismatch"
        return result
    if kind == "rect":
        x = int(assertion["x"])
        y = int(assertion["y"])
        w = int(assertion["w"])
        h = int(assertion["h"])
        result.update({"x": x, "y": y, "w": w, "h": h})
        if x < 0 or y < 0 or w <= 0 or h <= 0 or x + w > width or y + h > height:
            result["reason"] = "rect outside image"
            return result
        mismatches = []
        for yy in range(y, y + h):
            for xx in range(x, x + w):
                actual = pixel_at(width, data, xx, yy)
                if not within(actual, expected, tolerance):
                    mismatches.append({"x": xx, "y": yy, "actual": actual})
                    if len(mismatches) >= 8:
                        break
            if len(mismatches) >= 8:
                break
        result["mismatches"] = mismatches
        result["pass"] = len(mismatches) == 0
        if not result["pass"]:
            result["reason"] = "rect color mismatch"
        return result
    result["reason"] = "unknown assertion kind"
    return result


def evaluate(path, spec):
    if spec.get("source_kind") == "screenshot":
        return {
            "ok": False,
            "path": path,
            "source_kind": "screenshot",
            "error": "screenshots are smoke artifacts, not proof artifacts",
        }
    width, height, data = load_image(path)
    expected_width = spec.get("width")
    expected_height = spec.get("height")
    results = []
    ok = True
    if expected_width is not None and int(expected_width) != width:
        ok = False
        results.append({"kind": "dimensions", "pass": False, "actual": [width, height], "expected": [expected_width, expected_height]})
    if expected_height is not None and int(expected_height) != height:
        ok = False
        if not results or results[-1].get("kind") != "dimensions":
            results.append({"kind": "dimensions", "pass": False, "actual": [width, height], "expected": [expected_width, expected_height]})
    for assertion in spec.get("assertions", []):
        item = evaluate_assertion(width, height, data, assertion)
        results.append(item)
        ok = ok and item["pass"]
    return {
        "ok": ok,
        "path": path,
        "source_kind": spec.get("source_kind", "readback"),
        "origin": spec.get("origin", "top-left"),
        "format": "RGBA8",
        "width": width,
        "height": height,
        "assertions": results,
    }


def write_ppm(path, width, height, fill):
    rgb = bytes(fill[:3])
    with open(path, "wb") as f:
        f.write(("P6\n%d %d\n255\n" % (width, height)).encode("ascii"))
        f.write(rgb * width * height)


def patch_rect_ppm(path, width, x, y, w, h, color):
    with open(path, "r+b") as f:
        data = f.read()
        header_end = data.find(b"\n255\n") + 5
        pixels = bytearray(data[header_end:])
        for yy in range(y, y + h):
            for xx in range(x, x + w):
                idx = (yy * width + xx) * 3
                pixels[idx:idx + 3] = bytes(color[:3])
        f.seek(header_end)
        f.write(pixels)


def selftest():
    with tempfile.TemporaryDirectory() as tmp:
        on = os.path.join(tmp, "feature_on.ppm")
        off = os.path.join(tmp, "feature_off.ppm")
        write_ppm(on, 4, 4, PALETTE["background"])
        write_ppm(off, 4, 4, PALETTE["background"])
        patch_rect_ppm(on, 4, 1, 1, 2, 2, PALETTE["foreground"])
        spec = {
            "source_kind": "fixture",
            "width": 4,
            "height": 4,
            "origin": "top-left",
            "assertions": [
                {"kind": "point", "x": 1, "y": 1, "expected": "foreground"},
                {"kind": "rect", "x": 0, "y": 0, "w": 1, "h": 4, "expected": "background"},
            ],
        }
        passed = evaluate(on, spec)
        failed = evaluate(off, spec)
        screenshot = evaluate(on, {"source_kind": "screenshot", "assertions": []})
        assert passed["ok"], passed
        assert not failed["ok"], failed
        assert not screenshot["ok"], screenshot
    print("visual_assert selftest ok")
    return 0


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("image", nargs="?")
    parser.add_argument("--spec", help="JSON assertion spec path")
    parser.add_argument("--json", dest="json_path")
    parser.add_argument("--selftest", action="store_true")
    args = parser.parse_args()
    if args.selftest:
        return selftest()
    if not args.image or not args.spec:
        parser.error("image and --spec are required unless --selftest is used")
    with open(args.spec, "r", encoding="utf-8") as f:
        spec = json.load(f)
    result = evaluate(args.image, spec)
    encoded = json.dumps(result, sort_keys=True)
    if args.json_path:
        with open(args.json_path, "w", encoding="utf-8") as f:
            f.write(encoded + "\n")
    else:
        print(encoded)
    return 0 if result.get("ok") else 1


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