#!/usr/bin/env python3
"""CC0-1.0. Offline synthetic QCOW2 fixtures; never accepts an existing image."""
import argparse
import hashlib
import json
import platform
import shutil
import subprocess
import tempfile
from datetime import datetime, timezone
from pathlib import Path

CLUSTER = 65536
SIZE = 4 * CLUSTER
BASELINE = bytes([17]) * CLUSTER + bytes(CLUSTER) + bytes([51]) * CLUSTER + bytes([68]) * CLUSTER
CASES = ("intact", "data-byte-flip", "data-replaced", "data-truncated",
         "data-missing", "metadata-only-moved", "pair-moved", "pair-moved-cwd",
         "pair-moved-explicit", "runtime-override")


def require(ok, message):
    if not ok:
        raise RuntimeError(message)


def sha(data):
    return hashlib.sha256(data).hexdigest()


def snapshot(path):
    if not path.exists():
        return None
    data = path.read_bytes()
    return {"size": len(data), "sha256": sha(data)}


class Runner:
    def __init__(self, img, io, root):
        self.img, self.io, self.root = img, io, root
        self.commands = []
        self.cwd = root

    def call(self, tool, *args, checked=True):
        binary = self.img if tool == "qemu-img" else self.io
        proc = subprocess.run([binary, *map(str, args)], cwd=self.cwd,
                              capture_output=True, text=True, timeout=20)
        record = {"cwd": str(self.cwd.relative_to(self.root)), "command": [tool, *map(str, args)], "exitCode": proc.returncode,
                  "stdout": proc.stdout.replace(str(self.root), "<fixture>"),
                  "stderr": proc.stderr.replace(str(self.root), "<fixture>")}
        self.commands.append(record)
        if checked:
            require(proc.returncode == 0, json.dumps(record))
        return record


def scenario(img, io, case, raw, repetition):
    require(case in CASES and type(raw) is bool, "Invalid fixture")
    with tempfile.TemporaryDirectory(prefix="hesela-qcow2-external-") as temp:
        root = Path(temp)
        r = Runner(img, io, root)
        image = "disk.qcow2"
        top, data = root / image, root / "payload.bin"
        (root / "expected.raw").write_bytes(BASELINE)
        r.call("qemu-img", "create", "-f", "qcow2", "-o",
               "cluster_size=65536,data_file=payload.bin,data_file_raw=" + ("on" if raw else "off"),
               image, str(SIZE))
        for index, pattern in enumerate((17, 34, 51, 68)):
            r.call("qemu-io", "-f", "qcow2", "-c",
                   f"write -P {pattern} {index * CLUSTER} {CLUSTER}", image)
        r.call("qemu-io", "-f", "qcow2", "-c", f"write -z {CLUSTER} {CLUSTER}", image)
        r.call("qemu-img", "convert", "-f", "qcow2", "-O", "raw", image, "before.raw")
        require((root / "before.raw").read_bytes() == BASELINE, "Invalid baseline")
        before_top, before_data = snapshot(top), snapshot(data)
        raw_before = data.read_bytes()
        initial_raw_difference = sum(a != b for a, b in zip(raw_before, BASELINE))
        require(len(raw_before) == SIZE, "Unexpected data-file size")
        r.call("qemu-img", "map", "--output=json", "-f", "qcow2", image)
        initial_info = r.call("qemu-img", "info", "--backing-chain", "--output=json", "-f", "qcow2", image)

        if case == "data-byte-flip":
            with data.open("r+b") as stream:
                stream.write(bytes([16]))
        elif case == "data-replaced":
            data.write_bytes(bytes([85]) * SIZE)
        elif case == "data-truncated":
            with data.open("r+b") as stream:
                stream.truncate(2 * CLUSTER)
        elif case == "data-missing":
            data.unlink()
        elif case in ("metadata-only-moved", "pair-moved", "pair-moved-cwd", "pair-moved-explicit"):
            restore = root / "restore"
            restore.mkdir()
            top.rename(restore / image)
            if case != "metadata-only-moved":
                data.rename(restore / "payload.bin")
            else:
                data.unlink()
            image = "restore/" + image
            top, data = root / image, restore / "payload.bin"
            if case == "pair-moved-cwd":
                r.cwd = restore
                image = "disk.qcow2"
                (restore / "expected.raw").write_bytes(BASELINE)
        elif case == "runtime-override":
            (root / "replacement.bin").write_bytes(bytes([85]) * SIZE)

        inputs = ["-f", "qcow2", image]
        if case == "pair-moved-explicit":
            inputs = ["--image-opts",
                      "driver=qcow2,file.filename=restore/disk.qcow2,data-file.driver=file,data-file.filename=restore/payload.bin"]
        if case == "runtime-override":
            inputs = ["--image-opts",
                      "driver=qcow2,file.filename=disk.qcow2,data-file.driver=file,data-file.filename=replacement.bin"]
        check = r.call("qemu-img", "check", "--output=json", *inputs, checked=False)
        info = r.call("qemu-img", "info", "--backing-chain", "--output=json", *inputs, checked=False)
        # compare --image-opts applies to both operands.
        if "--image-opts" in inputs:
            compare = r.call("qemu-img", "compare", *inputs,
                             "driver=raw,file.filename=expected.raw", checked=False)
        else:
            compare = r.call("qemu-img", "compare", "-f", "qcow2", "-F", "raw",
                             image, "expected.raw", checked=False)
        export = r.call("qemu-img", "convert", "-O", "raw", *inputs, "after.raw", checked=False)
        payload = (r.cwd / "after.raw").read_bytes() if export["exitCode"] == 0 else None
        if payload is not None:
            require(len(payload) == SIZE, "Unexpected guest capacity")
        return {"case": case, "rawMode": raw, "repetition": repetition,
                "initialDirectRawChangedBytes": initial_raw_difference,
                "before": {"metadata": before_top, "data": before_data},
                "after": {"metadata": snapshot(top), "data": snapshot(data),
                          "override": snapshot(root / "replacement.bin")},
                "metadataUnchanged": snapshot(top) == before_top,
                "checkExit": check["exitCode"], "infoExit": info["exitCode"],
                "compareExit": compare["exitCode"], "exportExit": export["exitCode"],
                "changedGuestBytes": None if payload is None else sum(a != b for a, b in zip(payload, BASELINE)),
                "guestSha256": None if payload is None else sha(payload),
                "guestClusterSha256": None if payload is None else [
                    sha(payload[n:n + CLUSTER]) for n in range(0, SIZE, CLUSTER)],
                "initialInfo": json.loads(initial_info["stdout"]),
                "commands": r.commands}


def collect(img, io, repetitions):
    require(type(repetitions) is int and 1 <= repetitions <= 5, "Repetitions must be 1..5")
    binaries = {}
    for name, path in (("qemuImg", img), ("qemuIo", io)):
        binaries[name] = {
            "version": subprocess.run([path, "--version"], check=True, capture_output=True,
                                      text=True, timeout=10).stdout.splitlines()[0],
            "binarySha256": sha(Path(path).read_bytes())}
    return {"schemaVersion": 1, "generatedAt": datetime.now(timezone.utc).isoformat(),
            "evidenceType": "native-tool-observations-on-synthetic-inputs",
            "license": "CC0-1.0", "harnessSha256": sha(Path(__file__).read_bytes()),
            "environment": {"python": platform.python_version(), "kernel": platform.release(), **binaries},
            "fixture": {"clusterBytes": CLUSTER, "virtualBytes": SIZE,
                        "baselineSha256": sha(BASELINE), "baselineClusterBytes": [17, 0, 51, 68],
                        "repetitions": repetitions, "cases": list(CASES), "rawModes": [False, True]},
            "rows": [scenario(img, io, case, raw, rep)
                     for rep in range(1, repetitions + 1) for raw in (False, True) for case in CASES]}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--qemu-img", default="qemu-img")
    parser.add_argument("--qemu-io", default="qemu-io")
    parser.add_argument("--repetitions", type=int, default=3)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    require(not args.output.exists(), "Refusing to overwrite output")
    img, io = shutil.which(args.qemu_img), shutil.which(args.qemu_io)
    require(img and io, "qemu-img and qemu-io are required")
    result = collect(img, io, args.repetitions)
    with args.output.open("x") as stream:
        json.dump(result, stream, indent=2)
        stream.write("\n")


if __name__ == "__main__":
    main()