#!/usr/bin/env python3
"""Build deterministic B17 EROFS xattr images from one audited seed."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path
import shutil
import struct
import subprocess


SUPER_OFFSET = 1024
FEATURE_COMPAT_XATTR_FILTER = 0x00000004
ACL_FILTER_BITS = (1 << 21) | (1 << 30)
CRC32C_POLY = 0x82F63B78


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 write_json(path: Path, value: object) -> None:
    path.write_text(
        json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="ascii"
    )


def crc32c(data: bytes | bytearray, 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


class Image:
    def __init__(self, data: bytes | bytearray):
        self.data = bytearray(data)
        if len(self.data) < SUPER_OFFSET + 144:
            raise ValueError("seed is shorter than the EROFS superblock")
        if self.u32(SUPER_OFFSET) != 0xE0F5E1E2:
            raise ValueError("seed has the wrong EROFS magic")
        self.block_bits = self.data[SUPER_OFFSET + 12]
        self.block_size = 1 << self.block_bits
        self.meta_blkaddr = self.u32(SUPER_OFFSET + 40)
        self.xattr_blkaddr = self.u32(SUPER_OFFSET + 44)
        self.root_nid = self.u16(SUPER_OFFSET + 14)
        self.packed_nid = self.u64(SUPER_OFFSET + 96)

    @classmethod
    def load(cls, path: Path) -> "Image":
        return cls(path.read_bytes())

    def u16(self, offset: int) -> int:
        return struct.unpack_from("<H", self.data, offset)[0]

    def u32(self, offset: int) -> int:
        return struct.unpack_from("<I", self.data, offset)[0]

    def u64(self, offset: int) -> int:
        return struct.unpack_from("<Q", self.data, offset)[0]

    def put_u16(self, offset: int, value: int) -> None:
        struct.pack_into("<H", self.data, offset, value)

    def put_u32(self, offset: int, value: int) -> None:
        struct.pack_into("<I", self.data, offset, value)

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

    def inline_data_offset(self, inode: dict[str, int]) -> int:
        if inode["layout"] != 2:
            raise ValueError("fixture seed expected flat-inline metadata")
        return inode["offset"] + inode["inode_size"] + inode["xattr_size"]

    def directory_entries(self, inode: dict[str, int]) -> list[tuple[bytes, int]]:
        start = self.inline_data_offset(inode)
        data = self.data[start : start + inode["size"]]
        first_name = struct.unpack_from("<H", data, 8)[0]
        if first_name == 0 or first_name % 12 != 0 or first_name > len(data):
            raise ValueError("seed root directory is malformed")
        count = first_name // 12
        entries = []
        for index in range(count):
            entry = index * 12
            nid = struct.unpack_from("<Q", data, entry)[0]
            name_start = struct.unpack_from("<H", data, entry + 8)[0]
            if index + 1 < count:
                name_end = struct.unpack_from("<H", data, entry + 20)[0]
            else:
                name_end = len(data)
            name = bytes(data[name_start:name_end]).split(b"\0", 1)[0]
            entries.append((name, nid))
        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
            nid = next(
                (candidate for name, candidate in self.directory_entries(inode)
                 if name == component),
                None,
            )
            if nid is None:
                raise ValueError(f"seed path is absent: {path}")
            inode = self.inode(nid)
        return inode

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

    def first_inline_entry(self, path: str) -> tuple[dict[str, int], int]:
        inode = self.resolve(path)
        body = self.body_offset(inode)
        shared_count = self.data[body + 4]
        entry = body + 12 + shared_count * 4
        if entry + 4 > body + inode["xattr_size"]:
            raise ValueError(f"seed path has no inline xattr entry: {path}")
        return inode, entry

    def first_shared_entry(self, path: str) -> tuple[dict[str, int], int, int]:
        inode = self.resolve(path)
        body = self.body_offset(inode)
        if self.data[body + 4] == 0:
            raise ValueError(f"seed path has no shared xattr id: {path}")
        shared_id = self.u32(body + 12)
        entry = (self.xattr_blkaddr << self.block_bits) + shared_id * 4
        return inode, body, entry

    def prefix_record_offset(self) -> int:
        if self.packed_nid == 0:
            raise ValueError("seed has no packed prefix carrier")
        packed = self.inode(self.packed_nid)
        logical = self.u32(SUPER_OFFSET + 92) << 2
        if logical + 2 > packed["size"]:
            raise ValueError("seed prefix header is outside packed data")
        return self.inline_data_offset(packed) + logical

    def update_checksum(self) -> None:
        checksum = SUPER_OFFSET + 4
        self.put_u32(checksum, 0)
        end = SUPER_OFFSET + self.block_size - SUPER_OFFSET
        if end > len(self.data):
            raise ValueError("seed does not contain the checksummed block")
        self.put_u32(checksum, crc32c(self.data[SUPER_OFFSET:end]))


def load_spec(path: Path) -> dict:
    value = json.loads(path.read_text(encoding="ascii"))
    if value.get("schema") != 1 or value.get("batch") != "B17":
        raise SystemExit("invalid B17 fixture spec identity")
    cases = value.get("cases")
    if not isinstance(cases, list) or not cases:
        raise SystemExit("B17 fixture spec has no cases")
    identifiers = [case.get("id") for case in cases]
    if any(not isinstance(item, str) for item in identifiers):
        raise SystemExit("B17 fixture case has an invalid id")
    if len(identifiers) != len(set(identifiers)):
        raise SystemExit("B17 fixture case ids are not unique")
    return value


def create_source(root: Path) -> None:
    root.mkdir(parents=True)
    root.chmod(0o755)
    files = {
        "acl.bin": b"acl\n",
        "inline.bin": b"inline\n",
        "prefix.bin": b"prefix\n",
        "shared-0.bin": b"shared-0\n",
        "shared-1.bin": b"shared-1\n",
        "shared-2.bin": b"shared-2\n",
        "shared-3.bin": b"shared-3\n",
    }
    for name, payload in files.items():
        path = root / name
        path.write_bytes(payload)
        path.chmod(0o644)
    os.setxattr(root / "acl.bin", b"user.aclx", bytes([0xA5]) * 24)
    os.setxattr(root / "inline.bin", b"user.alpha", b"value\0binary\xff")
    os.setxattr(
        root / "prefix.bin", b"user.company.branch.leaf", b"prefix-value"
    )
    for index in range(4):
        os.setxattr(
            root / f"shared-{index}.bin", b"user.shared", b"shared-value"
        )
    for path in sorted(root.iterdir()):
        os.utime(path, (0, 0), follow_symlinks=False)
    os.utime(root, (0, 0), follow_symlinks=False)


def build_seed(spec: dict, output: Path, work: Path) -> dict[str, object]:
    mkfs = shutil.which("mkfs.erofs")
    if mkfs is None:
        raise SystemExit("mkfs.erofs is required to rebuild the B17 seed")
    version_run = subprocess.run(
        [mkfs, "-V"], check=False, text=True,
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
    )
    version = version_run.stdout.splitlines()[0] if version_run.stdout else ""
    if version != spec["seed"]["mkfs_version"]:
        raise SystemExit(
            f"mkfs.erofs version differs: expected {spec['seed']['mkfs_version']!r}, "
            f"got {version!r}"
        )
    source = work / "source"
    create_source(source)
    command = [
        mkfs,
        "-d0",
        "-T0",
        "--all-time",
        "--all-root",
        "--workers=1",
        "--sort=path",
        f"-U{spec['seed']['uuid']}",
        "-x2",
        "-Exattr-name-filter,force-inode-extended",
        "--xattr-prefix=user.company.",
        str(output),
        str(source),
    ]
    completed = subprocess.run(
        command, check=False, text=True,
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
    )
    if completed.returncode != 0:
        raise SystemExit(
            f"mkfs.erofs failed with {completed.returncode}:\n{completed.stdout}"
        )
    image = Image.load(output)
    packed = image.inode(image.packed_nid)
    image.put_u16(packed["offset"] + 4, 0o100644)
    image.update_checksum()
    output.write_bytes(image.data)
    return {"version": version, "command": command, "stdout": completed.stdout}


def configure_acl(image: Image, direct: list[dict[str, object]]) -> None:
    inode, entry = image.first_inline_entry("/acl.bin")
    if image.data[entry] != 4 or image.data[entry + 1] != 1:
        raise ValueError("seed ACL staging entry shape changed")
    if image.u16(entry + 2) != 24:
        raise ValueError("seed ACL staging value shape changed")
    acl = b"".join(
        [
            struct.pack("<I", 2),
            struct.pack("<HHI", 1, 6, 0xFFFFFFFF),
            struct.pack("<HHI", 4, 0, 0xFFFFFFFF),
            struct.pack("<HHI", 32, 0, 0xFFFFFFFF),
        ]
    )
    image.data[entry] = 0
    image.data[entry + 1] = 2
    image.put_u16(entry + 2, len(acl))
    image.data[entry + 4 : entry + 4 + len(acl)] = acl
    body = image.body_offset(inode)
    image.put_u32(body, image.u32(body) & ~ACL_FILTER_BITS)
    direct.extend(
        [
            {"field": "acl.empty-suffix-entry", "offset": entry, "bytes": 32},
            {"field": "acl.h_name_filter", "offset": body, "bytes": 4},
        ]
    )


def apply_mutation(image: Image, name: str) -> list[dict[str, object]]:
    direct: list[dict[str, object]] = []
    if name == "none":
        return direct
    if name == "acl-empty-suffix":
        configure_acl(image, direct)
    elif name in {"filter-reserved-one", "filter-reserved-255"}:
        value = 1 if name.endswith("one") else 255
        offset = SUPER_OFFSET + 104
        image.data[offset] = value
        direct.append({"field": "super.xattr_filter_reserved", "offset": offset, "bytes": 1})
    elif name in {"filter-feature-off", "filter-feature-off-reserved-one"}:
        offset = SUPER_OFFSET + 8
        image.put_u32(offset, image.u32(offset) & ~FEATURE_COMPAT_XATTR_FILTER)
        direct.append({"field": "super.feature_compat", "offset": offset, "bytes": 4})
        if name.endswith("reserved-one"):
            reserved = SUPER_OFFSET + 104
            image.data[reserved] = 1
            direct.append({"field": "super.xattr_filter_reserved", "offset": reserved, "bytes": 1})
    elif name == "unknown-filter-acl-present":
        configure_acl(image, direct)
        inode = image.resolve("/acl.bin")
        body = image.body_offset(inode)
        image.put_u32(body, image.u32(body) | ACL_FILTER_BITS)
        image.data[SUPER_OFFSET + 104] = 1
        direct.extend(
            [
                {"field": "ibody.h_name_filter", "offset": body, "bytes": 4},
                {"field": "super.xattr_filter_reserved", "offset": SUPER_OFFSET + 104, "bytes": 1},
            ]
        )
    elif name in {"inline-name-nul", "short-name-index", "inline-value-size"}:
        _, entry = image.first_inline_entry("/inline.bin")
        if name == "inline-name-nul":
            offset = entry + 4 + image.data[entry] // 2
            image.data[offset] = 0
            direct.append({"field": "inline.e_name", "offset": offset, "bytes": 1})
        elif name == "short-name-index":
            image.data[entry + 1] = 7
            direct.append({"field": "inline.e_name_index", "offset": entry + 1, "bytes": 1})
        else:
            image.put_u16(entry + 2, 0xFFFF)
            direct.append({"field": "inline.e_value_size", "offset": entry + 2, "bytes": 2})
    elif name in {"shared-name-nul", "shared-value-size"}:
        _, _, entry = image.first_shared_entry("/shared-0.bin")
        if name == "shared-name-nul":
            offset = entry + 4 + image.data[entry] // 2
            image.data[offset] = 0
            direct.append({"field": "shared.e_name", "offset": offset, "bytes": 1})
        else:
            image.put_u16(entry + 2, 0xFFFF)
            direct.append({"field": "shared.e_value_size", "offset": entry + 2, "bytes": 2})
    elif name == "shared-id-offset":
        _, body, _ = image.first_shared_entry("/shared-0.bin")
        image.put_u32(body + 12, 0xFFFFFFFF)
        direct.append({"field": "ibody.h_shared_xattrs[0]", "offset": body + 12, "bytes": 4})
    elif name in {"long-prefix-id",}:
        _, entry = image.first_inline_entry("/prefix.bin")
        image.data[entry + 1] = 0xFF
        direct.append({"field": "inline.e_name_index", "offset": entry + 1, "bytes": 1})
    elif name in {"prefix-infix-nul", "prefix-base-index", "prefix-record-truncated"}:
        prefix = image.prefix_record_offset()
        if name == "prefix-infix-nul":
            offset = prefix + 4
            image.data[offset] = 0
            direct.append({"field": "prefix.infix", "offset": offset, "bytes": 1})
        elif name == "prefix-base-index":
            image.data[prefix + 2] = 0
            direct.append({"field": "prefix.base_index", "offset": prefix + 2, "bytes": 1})
        else:
            image.put_u16(prefix, 255)
            direct.append({"field": "prefix.record_length", "offset": prefix, "bytes": 2})
    elif name == "prefix-offset-bounds":
        offset = SUPER_OFFSET + 92
        image.put_u32(offset, 0x100)
        direct.append({"field": "super.xattr_prefix_start", "offset": offset, "bytes": 4})
    elif name in {"header-shared-count", "header-only", "ibody-bounds", "inline-entry-truncated"}:
        inode = image.resolve("/inline.bin")
        body = image.body_offset(inode)
        if name == "header-shared-count":
            image.data[body + 4] = (inode["xattr_size"] - 12) // 4 + 1
            direct.append({"field": "ibody.h_shared_count", "offset": body + 4, "bytes": 1})
        else:
            count = {"header-only": 1, "ibody-bounds": 0xFFFF, "inline-entry-truncated": 3}[name]
            image.put_u16(inode["offset"] + 2, count)
            direct.append({"field": "inode.i_xattr_icount", "offset": inode["offset"] + 2, "bytes": 2})
    else:
        raise ValueError(f"unknown B17 mutation: {name}")
    return direct


def validate_seed_shape(image: Image) -> None:
    feature = image.u32(SUPER_OFFSET + 8)
    if not feature & FEATURE_COMPAT_XATTR_FILTER:
        raise ValueError("seed does not declare the xattr filter feature")
    if image.data[SUPER_OFFSET + 104] != 0:
        raise ValueError("seed xattr filter reserved byte is not zero")
    _, inline = image.first_inline_entry("/inline.bin")
    if bytes(image.data[inline + 4 : inline + 4 + image.data[inline]]) != b"alpha":
        raise ValueError("seed inline xattr changed")
    shared_inode, shared_body, shared = image.first_shared_entry("/shared-0.bin")
    if image.data[shared_body + 4] != 1 or shared_inode["xattr_size"] < 16:
        raise ValueError("seed shared xattr header changed")
    if bytes(image.data[shared + 4 : shared + 4 + image.data[shared]]) != b"shared":
        raise ValueError("seed shared xattr entry changed")
    _, long_entry = image.first_inline_entry("/prefix.bin")
    if not image.data[long_entry + 1] & 0x80:
        raise ValueError("seed long-prefix entry changed")
    prefix = image.prefix_record_offset()
    length = image.u16(prefix)
    if image.data[prefix + 2] != 1 or bytes(image.data[prefix + 3 : prefix + 2 + length]) != b"company.":
        raise ValueError("seed long-prefix record changed")
    _, acl = image.first_inline_entry("/acl.bin")
    if image.data[acl] != 4 or image.u16(acl + 2) != 24:
        raise ValueError("seed ACL staging entry changed")


def command_rebuild(spec_path: Path, seed_path: Path, output: Path, work: Path) -> int:
    spec = load_spec(spec_path)
    if output.exists() or work.exists():
        raise SystemExit("rebuild output and work paths must not exist")
    work.mkdir(parents=True)
    details = build_seed(spec, output, work)
    actual = sha256(output)
    expected = spec["seed"]["sha256"]
    tracked = sha256(seed_path)
    if actual != expected or tracked != expected:
        raise SystemExit(
            f"B17 seed reproducibility failed: generated={actual} tracked={tracked} expected={expected}"
        )
    validate_seed_shape(Image.load(output))
    print(json.dumps({"status": "PASS", "seed_sha256": actual, **details}, sort_keys=True))
    return 0


def command_create_seed(spec_path: Path, output: Path, work: Path) -> int:
    spec = load_spec(spec_path)
    if output.exists() or work.exists():
        raise SystemExit("seed output and work paths must not exist")
    work.mkdir(parents=True)
    details = build_seed(spec, output, work)
    validate_seed_shape(Image.load(output))
    actual = sha256(output)
    print(json.dumps({"status": "PASS", "seed_sha256": actual, **details}, sort_keys=True))
    return 0


def command_generate(spec_path: Path, seed_path: Path, output: Path) -> int:
    spec = load_spec(spec_path)
    if output.exists():
        raise SystemExit(f"refusing to overwrite fixture output: {output}")
    expected_seed = spec["seed"]["sha256"]
    actual_seed = sha256(seed_path)
    if actual_seed != expected_seed:
        raise SystemExit(
            f"B17 seed hash differs: expected {expected_seed}, got {actual_seed}"
        )
    seed = Image.load(seed_path)
    validate_seed_shape(seed)
    output.mkdir(parents=True)
    records = []
    aggregate = hashlib.sha256()
    for case in spec["cases"]:
        image = Image(seed.data)
        before = bytes(image.data)
        direct = apply_mutation(image, case["mutation"])
        image.update_checksum()
        destination = output / f"{case['id']}.erofs"
        destination.write_bytes(image.data)
        changed = [index for index, pair in enumerate(zip(before, image.data)) if pair[0] != pair[1]]
        record = {
            "id": case["id"],
            "class": case["class"],
            "mutation": case["mutation"],
            "path": destination.name,
            "bytes": destination.stat().st_size,
            "sha256": sha256(destination),
            "direct_mutations": direct,
            "changed_byte_count_with_checksum": len(changed),
        }
        records.append(record)
        aggregate.update(case["id"].encode("ascii"))
        aggregate.update(b"\0")
        aggregate.update(record["sha256"].encode("ascii"))
        aggregate.update(b"\n")
    manifest = {
        "schema": 1,
        "batch": "B17",
        "seed_sha256": actual_seed,
        "fixture_count": len(records),
        "legal_count": sum(record["class"] == "legal" for record in records),
        "damaged_count": sum(record["class"] == "damaged" for record in records),
        "fixture_set_sha256": aggregate.hexdigest(),
        "fixtures": records,
    }
    write_json(output / "fixture-index.json", manifest)
    print(
        json.dumps(
            {
                "status": "PASS",
                "fixture_count": manifest["fixture_count"],
                "fixture_set_sha256": manifest["fixture_set_sha256"],
            },
            sort_keys=True,
        )
    )
    return 0


def main() -> int:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)
    create = subparsers.add_parser("create-seed")
    create.add_argument("--spec", type=Path, required=True)
    create.add_argument("--output", type=Path, required=True)
    create.add_argument("--work", type=Path, required=True)
    rebuild = subparsers.add_parser("rebuild-seed")
    rebuild.add_argument("--spec", type=Path, required=True)
    rebuild.add_argument("--seed", type=Path, required=True)
    rebuild.add_argument("--output", type=Path, required=True)
    rebuild.add_argument("--work", type=Path, required=True)
    generate = subparsers.add_parser("generate")
    generate.add_argument("--spec", type=Path, required=True)
    generate.add_argument("--seed", type=Path, required=True)
    generate.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    if args.command == "create-seed":
        return command_create_seed(args.spec, args.output, args.work)
    if args.command == "rebuild-seed":
        return command_rebuild(args.spec, args.seed, args.output, args.work)
    return command_generate(args.spec, args.seed, args.output)


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