#!/usr/bin/env python3
"""CC0-1.0. Bounded offline QCOW2 observations; 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
CASES = ("intact", "base-replaced", "base-missing", "overlay-only-relocated",
         "chain-relocated", "payload-byte-flip", "safe-rebase", "unsafe-rebase", "base-truncated")
BASELINE = bytes([34]) * CLUSTER + bytes(CLUSTER) + bytes([17]) * (2 * CLUSTER)
CHANGED = bytes([34]) * CLUSTER + bytes(CLUSTER) + bytes([51]) * (2 * CLUSTER)


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


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


def digest(path):
    return sha(path.read_bytes()) if path.exists() else None


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

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

    def parsed(self, *args):
        return json.loads(self.call("qemu-img", *args)["stdout"])


def scenario(img, io, case, repetition, parent=None):
    require(case in CASES, "Unknown case")
    with tempfile.TemporaryDirectory(prefix="hesela-qcow2-", dir=parent) as temp:
        root = Path(temp)
        r = Runner(img, io, root)
        (root / "base.raw").write_bytes(bytes([17]) * SIZE)
        (root / "expected.raw").write_bytes(BASELINE)
        r.call("qemu-img", "create", "-f", "qcow2", "-o", "cluster_size=65536",
               "-F", "raw", "-b", "base.raw", "overlay.qcow2")
        r.call("qemu-io", "-f", "qcow2", "-c", "write -P 34 0 65536", "overlay.qcow2")
        r.call("qemu-io", "-f", "qcow2", "-c", "write -z 65536 65536", "overlay.qcow2")
        r.call("qemu-img", "convert", "-f", "qcow2", "-O", "raw",
               "overlay.qcow2", "before.raw")
        require((root / "before.raw").read_bytes() == BASELINE, "Bad initial payload")
        top = root / "overlay.qcow2"
        before_top, before_base = digest(top), digest(root / "base.raw")
        mapping = r.parsed("map", "--output=json", "-f", "qcow2", "overlay.qcow2")
        expected = BASELINE
        image = "overlay.qcow2"
        if case == "base-replaced":
            (root / "base.raw").write_bytes(bytes([51]) * SIZE)
            expected = CHANGED
        elif case == "base-truncated":
            (root / "base.raw").write_bytes(bytes([17]) * (2 * CLUSTER))
            expected = bytes([34]) * CLUSTER + bytes(3 * CLUSTER)
        elif case == "base-missing":
            (root / "base.raw").unlink()
        elif case in ("overlay-only-relocated", "chain-relocated"):
            (root / "restore").mkdir()
            shutil.move(top, root / "restore/overlay.qcow2")
            if case == "chain-relocated":
                shutil.move(root / "base.raw", root / "restore/base.raw")
            else:
                (root / "base.raw").unlink()
            image = "restore/overlay.qcow2"
            top = root / image
        elif case == "payload-byte-flip":
            extent = next(item for item in mapping
                          if item["start"] == 0 and item["data"] and item["depth"] == 0)
            with top.open("r+b") as output:
                output.seek(extent["offset"])
                output.write(bytes([35]))
            expected = bytes([35]) + BASELINE[1:]
        elif case in ("safe-rebase", "unsafe-rebase"):
            (root / "newbase.raw").write_bytes(bytes([51]) * SIZE)
            flags = ["-u"] if case == "unsafe-rebase" else []
            r.call("qemu-img", "rebase", *flags, "-f", "qcow2", "-F", "raw",
                   "-b", "newbase.raw", image)
            if case == "unsafe-rebase":
                expected = CHANGED
        missing = case in ("base-missing", "overlay-only-relocated")
        check = r.call("qemu-img", "check", "--output=json", "-f", "qcow2", image,
                       allowed=(1,) if missing else (0,))
        parsed_check = json.loads(check["stdout"]) if check["stdout"].strip() else None
        compare = r.call("qemu-img", "compare", "-f", "qcow2", "-F", "raw",
                         image, "expected.raw", allowed=(0, 1, 2))
        conversion = r.call("qemu-img", "convert", "-f", "qcow2", "-O", "raw",
                            image, "after.raw", allowed=(0, 1))
        actual = None
        if missing:
            require(compare["exitCode"] == 2 and conversion["exitCode"] == 1,
                    "Missing dependency did not fail as expected")
        else:
            require(conversion["exitCode"] == 0, "Conversion failed")
            actual = (root / "after.raw").read_bytes()
            require(actual == expected, "Full payload mismatch")
            require(compare["exitCode"] == int(expected != BASELINE), "Compare mismatch")
        after_top = digest(top)
        unchanged = case in ("intact", "base-replaced", "base-missing",
                             "overlay-only-relocated", "chain-relocated", "base-truncated")
        require((before_top == after_top) == unchanged, "Top-image hash expectation failed")
        selected_base = (root / "newbase.raw" if "rebase" in case else
                         root / ("restore/base.raw" if "relocated" in case else "base.raw"))
        return {"case": case, "repetition": repetition, "checkExit": check["exitCode"],
                "compareExit": compare["exitCode"], "exportExit": conversion["exitCode"],
                "topHashUnchanged": before_top == after_top,
                "guestMatchesBaseline": actual == BASELINE if actual is not None else None,
                "guestSha256": sha(actual) if actual is not None else None,
                "expectedGuestSha256": sha(expected) if not missing else None,
                "baselineGuestSha256": sha(BASELINE),
                "changedGuestBytes": sum(a != b for a, b in zip(actual, BASELINE))
                    if actual is not None else None,
                "topBeforeSha256": before_top, "topAfterSha256": after_top,
                "baseBeforeSha256": before_base, "selectedBaseSha256": digest(selected_base),
                "initialMap": mapping, "checkResult": parsed_check, "commands": r.commands}


def tail_case(img, io, repetition, parent=None):
    with tempfile.TemporaryDirectory(prefix="hesela-qcow2-tail-", dir=parent) as temp:
        root = Path(temp)
        r = Runner(img, io, root)
        small = bytes([17]) * CLUSTER
        (root / "small.raw").write_bytes(small)
        (root / "large.raw").write_bytes(small + bytes(CLUSTER))
        normal = r.call("qemu-img", "compare", "-f", "raw", "-F", "raw",
                        "small.raw", "large.raw")
        strict = r.call("qemu-img", "compare", "-s", "-f", "raw", "-F", "raw",
                        "small.raw", "large.raw", allowed=(1,))
        return {"case": "zero-tail-size", "repetition": repetition,
                "compareExit": normal["exitCode"], "strictCompareExit": strict["exitCode"],
                "smallSizeBytes": len(small), "largeSizeBytes": 2 * len(small),
                "smallSha256": digest(root / "small.raw"),
                "largeSha256": digest(root / "large.raw"), "commands": r.commands}


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--qemu-img", default="qemu-img")
    p.add_argument("--qemu-io", default="qemu-io")
    p.add_argument("--repetitions", type=int, default=3)
    p.add_argument("--temp-dir")
    args = p.parse_args()
    p.error("repetitions must be 1..10") if not 1 <= args.repetitions <= 10 else None
    img, io = shutil.which(args.qemu_img), shutil.which(args.qemu_io)
    require(img is not None and io is not None, "qemu-img and qemu-io are required")
    versions = {}
    for name, binary in (("qemuImg", img), ("qemuIo", io)):
        v = subprocess.run([binary, "--version"], capture_output=True, text=True,
                           check=True, timeout=20)
        versions[name] = {"version": v.stdout.splitlines()[0], "binarySha256": digest(Path(binary))}
    observations = []
    for repetition in range(1, args.repetitions + 1):
        for case in CASES:
            observations.append(scenario(img, io, case, repetition, args.temp_dir))
        observations.append(tail_case(img, io, repetition, args.temp_dir))
    result = {"schemaVersion": 1, "generatedAt": datetime.now(timezone.utc).isoformat(),
              "kind": "measured-tool-behavior-with-synthetic-inputs", "license": "CC0-1.0",
              "environment": {"python": platform.python_version(), "kernel": platform.release(),
                              **versions}, "clusterBytes": CLUSTER, "virtualSizeBytes": SIZE,
              "repetitions": args.repetitions, "harnessSha256": digest(Path(__file__)),
              "observations": observations}
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()