#!/usr/bin/env python3
"""Independent multi-block directory and xattr oracle for B19a."""

from __future__ import annotations

import argparse
from concurrent.futures import ThreadPoolExecutor
import hashlib
import json
import os
from pathlib import Path
import statistics
import struct
import threading
import time


SUPER = 1024
MAGIC = 0xE0F5E1E2
CRC32C_POLY = 0x82F63B78


class Reject(RuntimeError):
    def __init__(self, errno_name: str, point: str):
        super().__init__(f"{errno_name} at {point}")
        self.errno_name = errno_name
        self.point = point


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def crc32c(data: bytes, seed: int = 0xFFFFFFFF) -> int:
    value = seed
    for byte in data:
        value ^= byte
        for _ in range(8):
            value = (value >> 1) ^ (CRC32C_POLY if value & 1 else 0)
    return value & 0xFFFFFFFF


def load_spec(path: Path) -> dict:
    spec = json.loads(path.read_text(encoding="ascii"))
    if (
        spec.get("schema") != 1
        or spec.get("batch") != "B19a"
        or spec.get("candidate") != "P15-022"
    ):
        raise SystemExit("invalid B19a oracle spec identity")
    return spec


class ReadStats:
    def __init__(self) -> None:
        self.lock = threading.Lock()
        self.calls = 0
        self.bytes = 0
        self.block_reads = 0

    def add(self, offset: int, length: int, block_size: int) -> None:
        with self.lock:
            self.calls += 1
            self.bytes += length
            if length:
                self.block_reads += (
                    (offset + length - 1) // block_size - offset // block_size + 1
                )

    def reset(self) -> None:
        with self.lock:
            self.calls = 0
            self.bytes = 0
            self.block_reads = 0

    def snapshot(self) -> dict[str, int]:
        with self.lock:
            return {
                "bytes": self.bytes,
                "calls": self.calls,
                "provider_block_reads": self.block_reads,
            }


class Reader:
    def __init__(self, path: Path):
        self.path = path
        self.fd = os.open(path, os.O_RDONLY)
        self.size = os.fstat(self.fd).st_size
        self.stats = ReadStats()
        header = self.read(SUPER, 144, "super")
        if struct.unpack_from("<I", header)[0] != MAGIC:
            raise Reject("EINTEGRITY", "super.magic")
        self.feature_compat = struct.unpack_from("<I", header, 8)[0]
        self.block_bits = header[12]
        if self.block_bits < 9 or self.block_bits > 16:
            raise Reject("EINTEGRITY", "super.block-size")
        self.block_size = 1 << self.block_bits
        self.root_nid = struct.unpack_from("<H", header, 14)[0]
        self.blocks = struct.unpack_from("<I", header, 36)[0]
        self.limit = self.blocks << self.block_bits
        self.meta_blkaddr = struct.unpack_from("<I", header, 40)[0]
        self.xattr_blkaddr = struct.unpack_from("<I", header, 44)[0]
        if self.limit > self.size:
            raise Reject("EINTEGRITY", "super.bounds")
        if self.feature_compat & 1:
            expected = struct.unpack_from("<I", header, 4)[0]
            raw = bytearray(self.read(SUPER, self.block_size - SUPER, "super.checksum"))
            raw[4:8] = bytes(4)
            if crc32c(bytes(raw)) != expected:
                raise Reject("EINTEGRITY", "super.checksum")
        self.stats.reset()

    def close(self) -> None:
        os.close(self.fd)

    def __enter__(self) -> "Reader":
        return self

    def __exit__(self, *_: object) -> None:
        self.close()

    def read(self, offset: int, length: int, point: str) -> bytes:
        limit = self.limit if hasattr(self, "limit") else self.size
        if offset < 0 or length < 0 or offset > limit or length > limit - offset:
            raise Reject("EINTEGRITY", f"{point}.bounds")
        data = os.pread(self.fd, length, offset)
        if len(data) != length:
            raise Reject("EINTEGRITY", f"{point}.short")
        block_size = self.block_size if hasattr(self, "block_size") else 4096
        self.stats.add(offset, length, block_size)
        return data

    def inode(self, nid: int) -> dict[str, int]:
        offset = (self.meta_blkaddr << self.block_bits) + (nid << 5)
        raw = self.read(offset, 64, "inode")
        inode_format, xattr_count = struct.unpack_from("<HH", raw)
        inode_size = 64 if inode_format & 1 else 32
        xattr_size = 0 if xattr_count == 0 else 12 + 4 * (xattr_count - 1)
        size = struct.unpack_from("<Q" if inode_size == 64 else "<I", raw, 8)[0]
        return {
            "nid": nid,
            "offset": offset,
            "inode_size": inode_size,
            "xattr_size": xattr_size,
            "layout": (inode_format >> 1) & 7,
            "size": size,
            "start_block": struct.unpack_from("<I", raw, 16)[0],
        }

    def inode_data(self, inode: dict[str, int], logical: int, length: int) -> bytes:
        if logical > inode["size"] or length > inode["size"] - logical:
            raise Reject("EINTEGRITY", "inode.data-range")
        if inode["layout"] == 0:
            physical = (inode["start_block"] << self.block_bits) + logical
        elif inode["layout"] == 2:
            tail_start = ((inode["size"] + self.block_size - 1) // self.block_size - 1) * self.block_size
            if logical < tail_start:
                if logical + length > tail_start:
                    raise Reject("EINTEGRITY", "inode.inline-crossing")
                physical = (inode["start_block"] << self.block_bits) + logical
            else:
                physical = (
                    inode["offset"]
                    + inode["inode_size"]
                    + inode["xattr_size"]
                    + logical
                    - tail_start
                )
        else:
            raise Reject("EOPNOTSUPP", "inode.layout")
        return self.read(physical, length, "inode.data")

    @staticmethod
    def parse_dirblock(data: bytes, block_index: int) -> list[dict[str, object]]:
        if len(data) < 12:
            raise Reject("EINTEGRITY", "directory.header")
        first_name = struct.unpack_from("<H", data, 8)[0]
        if first_name < 12 or first_name >= len(data) or first_name % 12:
            raise Reject("EINTEGRITY", "directory.name-offset")
        count = first_name // 12
        entries = []
        previous_offset = 0
        previous_name: bytes | None = None
        for index in range(count):
            offset = index * 12
            nid, start, file_type = struct.unpack_from("<QHB", data, offset)
            end = (
                struct.unpack_from("<H", data, offset + 20)[0]
                if index + 1 < count
                else len(data)
            )
            if (
                start < first_name
                or start >= len(data)
                or (index == 0 and start != first_name)
                or (index != 0 and start <= previous_offset)
                or end <= start
                or end > len(data)
            ):
                raise Reject("EINTEGRITY", "directory.name-bounds")
            span = data[start:end]
            if index + 1 < count:
                if b"\0" in span:
                    raise Reject("EINTEGRITY", "directory.name-nul")
                name = span
            else:
                name = span.split(b"\0", 1)[0]
            if not name or len(name) > 255 or b"/" in name:
                raise Reject("EINTEGRITY", "directory.name")
            if previous_name is not None and previous_name >= name:
                raise Reject("EINTEGRITY", "directory.order")
            entries.append(
                {
                    "block": block_index,
                    "entry": index,
                    "file_type": file_type,
                    "name": name,
                    "nid": nid,
                }
            )
            previous_offset = start
            previous_name = name
        return entries

    @staticmethod
    def validate_boundary(
        left: list[dict[str, object]], right: list[dict[str, object]]
    ) -> None:
        if left[-1]["name"] >= right[0]["name"]:
            raise Reject("EINTEGRITY", "directory.boundary-order")

    @staticmethod
    def resolve_entries(entries: list[dict[str, object]], name: bytes) -> int:
        matches = [int(entry["nid"]) for entry in entries if entry["name"] == name]
        if not matches:
            raise Reject("ENOENT", "path.missing")
        if len(matches) != 1:
            raise Reject("EINTEGRITY", "path.duplicate")
        return matches[0]

    def directory_entries(self, inode: dict[str, int]) -> list[dict[str, object]]:
        entries = []
        previous_name: bytes | None = None
        logical = 0
        block_index = 0
        while logical < inode["size"]:
            length = min(self.block_size, inode["size"] - logical)
            block_entries = self.parse_dirblock(
                self.inode_data(inode, logical, length), block_index
            )
            if entries:
                self.validate_boundary(entries, block_entries)
            entries.extend(block_entries)
            previous_name = block_entries[-1]["name"]
            logical += length
            block_index += 1
        return entries

    def resolve(self, path: str) -> dict[str, int]:
        inode = self.inode(self.root_nid)
        for component in path.strip("/").encode("ascii").split(b"/"):
            if not component:
                continue
            inode = self.inode(
                self.resolve_entries(self.directory_entries(inode), component)
            )
        return inode

    def body_offset(self, inode: dict[str, int]) -> int:
        return inode["offset"] + inode["inode_size"]

    @staticmethod
    def parse_entry(raw: bytes, offset: int, limit: int, point: str) -> tuple[dict, int]:
        if offset > limit or 4 > limit - offset:
            raise Reject("EINTEGRITY", f"{point}.header")
        name_length, name_index, value_length = struct.unpack_from("<BBH", raw, offset)
        total = (4 + name_length + value_length + 3) & ~3
        if total > limit - offset:
            raise Reject("EINTEGRITY", f"{point}.bounds")
        name_start = offset + 4
        name_end = name_start + name_length
        value_end = name_end + value_length
        name = raw[name_start:name_end]
        if b"\0" in name:
            raise Reject("EINTEGRITY", f"{point}.name-nul")
        return (
            {
                "index": name_index,
                "name": name,
                "value": raw[name_end:value_end],
            },
            total,
        )

    def shared_entry(self, shared_id: int) -> dict:
        offset = (self.xattr_blkaddr << self.block_bits) + shared_id * 4
        header = self.read(offset, 4, "shared.header")
        name_length, _, value_length = struct.unpack("<BBH", header)
        total = (4 + name_length + value_length + 3) & ~3
        raw = self.read(offset, total, "shared.entry")
        entry, parsed = self.parse_entry(raw, 0, len(raw), "shared")
        if parsed != total:
            raise Reject("EINTEGRITY", "shared.size")
        return entry

    def load_body(self, inode: dict[str, int], validate_shared: bool) -> dict:
        size = inode["xattr_size"]
        if size < 12:
            raise Reject("EINTEGRITY", "ibody.header")
        if size == 12:
            raise Reject("EOPNOTSUPP", "ibody.header-only")
        raw = self.read(self.body_offset(inode), size, "ibody")
        shared_count = raw[4]
        header_size = 12 + shared_count * 4
        if header_size > size:
            raise Reject("EINTEGRITY", "ibody.shared-count")
        shared_ids = [
            struct.unpack_from("<I", raw, 12 + index * 4)[0]
            for index in range(shared_count)
        ]
        inline = []
        cursor = header_size
        while cursor < size:
            entry, total = self.parse_entry(raw, cursor, size, "inline")
            inline.append(entry)
            cursor += total
        if validate_shared:
            for shared_id in shared_ids:
                self.shared_entry(shared_id)
        return {
            "header_size": header_size,
            "inline": inline,
            "raw": raw,
            "shared_ids": shared_ids,
        }


class MountBudget:
    def __init__(self, limit: int):
        self.limit = limit
        self.resident = 0
        self.lock = threading.Lock()

    def reserve(self, amount: int) -> bool:
        with self.lock:
            if amount > self.limit - self.resident:
                return False
            self.resident += amount
            return True

    def release(self, amount: int) -> None:
        with self.lock:
            if amount > self.resident:
                raise RuntimeError("cache budget underflow")
            self.resident -= amount


class BodyCache:
    EMPTY = "EMPTY"
    INFLIGHT = "INFLIGHT"
    READY = "READY"
    FAILED = "FAILED"

    def __init__(self, spec: dict, budget: MountBudget):
        self.limit = spec["model"]["entry_body_limit_bytes"]
        self.budget = budget
        self.cv = threading.Condition()
        self.state = self.EMPTY
        self.body: dict | None = None
        self.error: Reject | None = None
        self.waiters = 0
        self.loads = 0
        self.charged = 0
        self.closing = False
        self.wait_for_waiters = 0

    def acquire(self, reader: Reader, inode: dict[str, int]) -> dict:
        owner = False
        reserved = False
        with self.cv:
            if self.closing:
                raise Reject("ENXIO", "cache.closing")
            if self.state == self.READY:
                if self.body is None:
                    raise RuntimeError("READY cache has no body")
                return self.body
            if self.state == self.INFLIGHT:
                self.waiters += 1
                try:
                    while self.state == self.INFLIGHT:
                        self.cv.wait()
                    if self.state == self.READY:
                        if self.body is None:
                            raise RuntimeError("published cache has no body")
                        return self.body
                    if self.state != self.FAILED or self.error is None:
                        raise RuntimeError("cache waiter has no typed result")
                    raise Reject(self.error.errno_name, self.error.point)
                finally:
                    self.waiters -= 1
                    if self.state == self.FAILED and self.waiters == 0:
                        self.error = None
                        self.state = self.EMPTY
                    if self.waiters == 0:
                        self.cv.notify_all()
            if self.state == self.FAILED:
                if self.waiters:
                    return reader.load_body(inode, validate_shared=True)
                self.error = None
                self.state = self.EMPTY
            if inode["xattr_size"] <= self.limit:
                reserved = self.budget.reserve(inode["xattr_size"])
            if not reserved:
                return reader.load_body(inode, validate_shared=True)
            self.state = self.INFLIGHT
            self.charged = inode["xattr_size"]
            owner = True
        if not owner:
            raise RuntimeError("cache claim lost owner state")
        try:
            if self.wait_for_waiters:
                deadline = time.monotonic() + 5
                while True:
                    with self.cv:
                        if self.waiters >= self.wait_for_waiters:
                            break
                    if time.monotonic() >= deadline:
                        raise RuntimeError("cache waiters did not reach owner barrier")
                    time.sleep(0.001)
            body = reader.load_body(inode, validate_shared=True)
            with self.cv:
                self.loads += 1
                self.body = body
                self.error = None
                self.state = self.READY
                self.cv.notify_all()
            return body
        except Reject as error:
            with self.cv:
                self.loads += 1
                self.budget.release(self.charged)
                self.charged = 0
                self.error = error
                self.state = self.FAILED
                if self.waiters == 0:
                    self.error = None
                    self.state = self.EMPTY
                self.cv.notify_all()
            raise

    def invalidate(self) -> None:
        with self.cv:
            while self.state == self.INFLIGHT or self.waiters:
                self.cv.wait()
            if self.state == self.READY:
                self.body = None
                self.budget.release(self.charged)
                self.charged = 0
            self.error = None
            self.state = self.EMPTY

    def close(self) -> None:
        with self.cv:
            self.closing = True
            while self.state == self.INFLIGHT or self.waiters:
                self.cv.wait()
            if self.state == self.READY:
                self.body = None
                self.budget.release(self.charged)
                self.charged = 0
            self.error = None
            self.state = self.EMPTY


def body_for_operation(
    reader: Reader, inode: dict[str, int], variant: str, cache: BodyCache | None
) -> dict:
    if variant == "baseline":
        return reader.load_body(inode, validate_shared=False)
    if variant != "candidate" or cache is None:
        raise RuntimeError("invalid B19a oracle variant")
    return cache.acquire(reader, inode)


def entries_for_body(reader: Reader, body: dict) -> list[dict]:
    entries = list(body["inline"])
    entries.extend(reader.shared_entry(shared_id) for shared_id in body["shared_ids"])
    return entries


def getxattr(
    reader: Reader,
    inode: dict[str, int],
    name: bytes,
    variant: str,
    cache: BodyCache | None,
) -> bytes:
    body = body_for_operation(reader, inode, variant, cache)
    for entry in body["inline"]:
        if entry["index"] == 1 and entry["name"] == name:
            return bytes(entry["value"])
    for shared_id in body["shared_ids"]:
        entry = reader.shared_entry(shared_id)
        if entry["index"] == 1 and entry["name"] == name:
            return bytes(entry["value"])
    raise Reject("ENOATTR", "xattr.not-found")


def listxattr(
    reader: Reader,
    inode: dict[str, int],
    variant: str,
    cache: BodyCache | None,
) -> list[bytes]:
    return [
        bytes(entry["name"])
        for entry in entries_for_body(
            reader, body_for_operation(reader, inode, variant, cache)
        )
        if entry["index"] == 1
    ]


def expected_value(label: str, length: int) -> bytes:
    seed = (label + "|").encode("ascii")
    return (seed * (length // len(seed) + 1))[:length]


def encode_dirblock(names: list[bytes], size: int = 128) -> bytes:
    first_name = len(names) * 12
    payload = bytearray(size)
    cursor = first_name
    for index, name in enumerate(names):
        if cursor + len(name) > size:
            raise ValueError("synthetic directory block is too small")
        struct.pack_into("<QHBB", payload, index * 12, index + 1, cursor, 1, 0)
        payload[cursor : cursor + len(name)] = name
        cursor += len(name)
    return bytes(payload)


def expect_reject(function, errno_name: str, point: str) -> dict[str, str]:
    try:
        function()
    except Reject as error:
        if error.errno_name != errno_name or error.point != point:
            raise SystemExit(
                f"oracle control expected {errno_name} at {point}, got {error}"
            )
        return {"errno": error.errno_name, "point": error.point}
    raise SystemExit(f"oracle control unexpectedly passed: {errno_name} at {point}")


def run_selftest(output: Path) -> dict:
    controls = {}
    first = Reader.parse_dirblock(
        encode_dirblock([b"alpha", b"boundary-last"], 96),
        0,
    )
    second = Reader.parse_dirblock(
        encode_dirblock([b"target", b"zeta"], 80), 1
    )
    Reader.validate_boundary(first, second)
    combined = first + second
    target_nid = Reader.resolve_entries(combined, b"target")
    matches = [entry for entry in combined if entry["name"] == b"target"]
    if (
        len(matches) != 1
        or matches[0]["block"] != 1
        or int(matches[0]["nid"]) != target_nid
    ):
        raise SystemExit("target was not resolved after a block boundary")
    controls["target-after-boundary"] = "PASS"
    controls["short-final-block"] = "PASS"

    controls["duplicate-within-block"] = expect_reject(
        lambda: Reader.parse_dirblock(encode_dirblock([b"dup", b"dup"]), 0),
        "EINTEGRITY",
        "directory.order",
    )
    left = Reader.parse_dirblock(encode_dirblock([b"alpha", b"dup"]), 0)
    right = Reader.parse_dirblock(encode_dirblock([b"dup", b"zeta"]), 1)
    controls["duplicate-across-blocks"] = expect_reject(
        lambda: Reader.validate_boundary(left, right),
        "EINTEGRITY",
        "directory.boundary-order",
    )
    controls["missing-target"] = expect_reject(
        lambda: Reader.resolve_entries(combined, b"missing"),
        "ENOENT",
        "path.missing",
    )
    duplicate_entries = combined + [{**combined[-1], "nid": 999, "name": b"target"}]
    controls["duplicate-resolution"] = expect_reject(
        lambda: Reader.resolve_entries(duplicate_entries, b"target"),
        "EINTEGRITY",
        "path.duplicate",
    )
    backwards = Reader.parse_dirblock(encode_dirblock([b"aardvark", b"beta"]), 1)
    controls["backwards-boundary"] = expect_reject(
        lambda: Reader.validate_boundary(first, backwards),
        "EINTEGRITY",
        "directory.boundary-order",
    )
    malformed = bytearray(encode_dirblock([b"alpha", b"bravo"]))
    struct.pack_into("<H", malformed, 20, len(malformed) + 1)
    controls["crossing-name"] = expect_reject(
        lambda: Reader.parse_dirblock(bytes(malformed), 0),
        "EINTEGRITY",
        "directory.name-bounds",
    )
    report = {"controls": controls, "status": "PASS"}
    output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="ascii")
    return report


def operation_result(
    image: Path, spec: dict, variant: str, operation: str
) -> dict[str, object]:
    budget = MountBudget(spec["model"]["cache_budget_bytes"])
    cache = BodyCache(spec, budget) if variant == "candidate" else None
    try:
        with Reader(image) as reader:
            inode = reader.resolve(spec["fixture"]["target"])
            reader.stats.reset()
            if operation == "get-inline":
                value = getxattr(reader, inode, b"inline", variant, cache)
                return {"status": "PASS", "value_sha256": hashlib.sha256(value).hexdigest()}
            if operation == "get-shared":
                value = getxattr(reader, inode, b"shared", variant, cache)
                return {"status": "PASS", "value_sha256": hashlib.sha256(value).hexdigest()}
            if operation == "list-user":
                names = listxattr(reader, inode, variant, cache)
                return {"names": [name.decode("ascii") for name in names], "status": "PASS"}
            raise RuntimeError(f"unknown operation: {operation}")
    except Reject as error:
        return {"errno": spec["errno"].get(error.errno_name, -1), "point": error.point, "status": error.errno_name}
    finally:
        if cache is not None:
            cache.close()
        if budget.resident != 0:
            raise RuntimeError("operation cache budget leaked")


def run_correctness(fixtures: Path, spec: dict, output: Path) -> dict:
    valid = fixtures / "valid.erofs"
    with Reader(valid) as reader:
        root = reader.inode(reader.root_nid)
        root_entries = reader.directory_entries(root)
        target_entries = [entry for entry in root_entries if entry["name"] == b"target.bin"]
        if len(target_entries) != 1 or int(target_entries[0]["block"]) == 0:
            raise SystemExit("real target did not resolve after the first block")
        if [entry["name"] for entry in root_entries] != sorted(
            entry["name"] for entry in root_entries
        ):
            raise SystemExit("real directory order is unstable")
        boundary = []
        for left, right in zip(root_entries, root_entries[1:]):
            if left["block"] != right["block"]:
                boundary.append(
                    {
                        "left": bytes(left["name"]).decode("ascii"),
                        "right": bytes(right["name"]).decode("ascii"),
                    }
                )
        inode = reader.resolve(spec["fixture"]["target"])
        body = reader.load_body(inode, validate_shared=True)
        if not body["inline"] or not body["shared_ids"]:
            raise SystemExit("real target lacks inline/shared xattr coverage")
        directory = {
            "block_count": (root["size"] + reader.block_size - 1) // reader.block_size,
            "boundaries": boundary,
            "entry_count": len(root_entries),
            "target_block": target_entries[0]["block"],
            "target_entry": target_entries[0]["entry"],
        }

    expected_inline = expected_value(
        "inline-value-p15-022", spec["fixture"]["inline_value_bytes"]
    )
    expected_shared = expected_value(
        "shared-value-p15-022", spec["fixture"]["shared_value_bytes"]
    )
    cases = []
    for variant in ("baseline", "candidate"):
        inline = operation_result(valid, spec, variant, "get-inline")
        shared = operation_result(valid, spec, variant, "get-shared")
        listed = operation_result(valid, spec, variant, "list-user")
        if inline.get("value_sha256") != hashlib.sha256(expected_inline).hexdigest():
            raise SystemExit(f"{variant} inline value mismatch")
        if shared.get("value_sha256") != hashlib.sha256(expected_shared).hexdigest():
            raise SystemExit(f"{variant} shared value mismatch")
        if sorted(listed.get("names", [])) != ["inline", "shared"]:
            raise SystemExit(f"{variant} list result mismatch: {listed}")
        cases.extend(
            [
                {"id": "valid-inline", "variant": variant, "result": inline},
                {"id": "valid-shared", "variant": variant, "result": shared},
                {"id": "valid-list", "variant": variant, "result": listed},
            ]
        )

    damaged = [
        ("corrupt-shared-count.erofs", "get-inline", "EINTEGRITY"),
        ("corrupt-shared-id.erofs", "get-shared", "EINTEGRITY"),
        ("corrupt-inline-name.erofs", "list-user", "EINTEGRITY"),
    ]
    for filename, operation, expected in damaged:
        for variant in ("baseline", "candidate"):
            result = operation_result(fixtures / filename, spec, variant, operation)
            if result["status"] != expected or int(result.get("errno", -1)) < 0:
                raise SystemExit(f"{filename} {variant} mismatch: {result}")
            cases.append({"id": filename, "variant": variant, "result": result})

    budget = MountBudget(spec["model"]["cache_budget_bytes"])
    cache = BodyCache(spec, budget)
    cache.wait_for_waiters = spec["concurrency"]["workers"] - 1
    with Reader(valid) as reader:
        inode = reader.resolve(spec["fixture"]["target"])
        barrier = threading.Barrier(spec["concurrency"]["workers"])

        def worker(_: int) -> str:
            barrier.wait()
            value = getxattr(reader, inode, b"inline", "candidate", cache)
            return hashlib.sha256(value).hexdigest()

        with ThreadPoolExecutor(max_workers=spec["concurrency"]["workers"]) as executor:
            hashes = list(executor.map(worker, range(spec["concurrency"]["workers"])))
        if len(set(hashes)) != 1 or cache.loads != 1 or cache.state != BodyCache.READY:
            raise SystemExit("candidate owner/waiter publication mismatch")
        first_charge = budget.resident
        cache.invalidate()
        if budget.resident != 0 or cache.state != BodyCache.EMPTY:
            raise SystemExit("candidate invalidation did not release cache")
        cache.wait_for_waiters = 0
        getxattr(reader, inode, b"inline", "candidate", cache)
        if cache.loads != 2 or budget.resident != first_charge:
            raise SystemExit("candidate re-establishment mismatch")
    cache.close()
    if budget.resident != 0:
        raise SystemExit("candidate reclaim/close leaked budget")

    report = {
        "cases": cases,
        "concurrency": {
            "body_loads": 1,
            "completed": spec["concurrency"]["workers"],
            "workers": spec["concurrency"]["workers"],
        },
        "directory": directory,
        "fixture_sha256": {
            path.name: sha256(path) for path in sorted(fixtures.glob("*.erofs"))
        },
        "lifecycle": {
            "establish": "PASS",
            "hit": "PASS",
            "invalidate": "PASS",
            "reclaim": "PASS",
            "resident_after": budget.resident,
        },
        "status": "PASS",
    }
    output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="ascii")
    return report


def run_sample(
    image: Path, spec: dict, variant: str, sample: int, output: Path
) -> dict:
    budget = MountBudget(spec["model"]["cache_budget_bytes"])
    cache = BodyCache(spec, budget) if variant == "candidate" else None
    operations = spec["benchmark"]["operations"]
    with Reader(image) as reader:
        inode = reader.resolve(spec["fixture"]["target"])

        def one_loop() -> None:
            for operation in operations:
                if operation == "get-inline":
                    getxattr(reader, inode, b"inline", variant, cache)
                elif operation == "get-shared":
                    getxattr(reader, inode, b"shared", variant, cache)
                elif operation == "list-user":
                    listxattr(reader, inode, variant, cache)
                else:
                    raise RuntimeError(f"unknown benchmark operation: {operation}")

        for _ in range(spec["benchmark"]["warmup_loops"]):
            one_loop()
        reader.stats.reset()
        started = time.monotonic_ns()
        for _ in range(spec["benchmark"]["loops_per_sample"]):
            one_loop()
        elapsed = time.monotonic_ns() - started
        reads = reader.stats.snapshot()
    loads = cache.loads if cache is not None else 0
    resident_before_close = budget.resident
    if cache is not None:
        cache.close()
    if budget.resident != 0:
        raise SystemExit("benchmark cache budget leaked")
    report = {
        "cache_body_loads": loads,
        "cache_resident_before_close": resident_before_close,
        "elapsed_ns": elapsed,
        "fixture_sha256": sha256(image),
        "host": os.uname().sysname + " " + os.uname().release,
        "loops": spec["benchmark"]["loops_per_sample"],
        "operations": operations,
        "reads": reads,
        "sample": sample,
        "status": "PASS",
        "variant": variant,
        "warmup_loops": spec["benchmark"]["warmup_loops"],
    }
    output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="ascii")
    print(json.dumps(report, sort_keys=True))
    return report


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--spec", type=Path, required=True)
    subparsers = parser.add_subparsers(dest="command", required=True)
    selftest = subparsers.add_parser("selftest")
    selftest.add_argument("--output", type=Path, required=True)
    correctness = subparsers.add_parser("correctness")
    correctness.add_argument("--fixtures", type=Path, required=True)
    correctness.add_argument("--output", type=Path, required=True)
    sample = subparsers.add_parser("sample")
    sample.add_argument("--image", type=Path, required=True)
    sample.add_argument("--variant", choices=("baseline", "candidate"), required=True)
    sample.add_argument("--sample", type=int, required=True)
    sample.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    spec = load_spec(args.spec)
    if args.command == "selftest":
        report = run_selftest(args.output)
        print(json.dumps(report, sort_keys=True))
    elif args.command == "correctness":
        report = run_correctness(args.fixtures, spec, args.output)
        print(json.dumps(report, sort_keys=True))
    else:
        run_sample(args.image, spec, args.variant, args.sample, args.output)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())