#!/usr/bin/env python3
"""Generate deterministic B21 superblock trust-order fixtures."""

from __future__ import annotations

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


SUPER = 1024
FEATURE_COMPAT_SB_CHKSUM = 0x00000001
FEATURE_INCOMPAT_XATTR_PREFIXES = 0x00000040
UNKNOWN_INCOMPAT = 0x80000000
CRC32C_POLY = 0x82F63B78


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


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"
    )


class Image:
    def __init__(self, data: bytes | bytearray):
        self.data = bytearray(data)
        if len(self.data) < SUPER + 144:
            raise ValueError("image is shorter than the extended superblock")
        if self.u32(SUPER) != 0xE0F5E1E2:
            raise ValueError("image has the wrong EROFS magic")
        if not self.u32(SUPER + 8) & FEATURE_COMPAT_SB_CHKSUM:
            raise ValueError("image does not declare a superblock checksum")

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

    def clone(self) -> "Image":
        return Image(self.data)

    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)

    @property
    def block_bits(self) -> int:
        return self.data[SUPER + 12]

    @property
    def block_size(self) -> int:
        return 1 << self.block_bits

    @property
    def checksum_end(self) -> int:
        span = self.block_size
        if span > SUPER:
            span -= SUPER
        return SUPER + span

    def checksum_valid(self) -> bool:
        expected = self.u32(SUPER + 4)
        canonical = bytearray(self.data[SUPER : self.checksum_end])
        struct.pack_into("<I", canonical, 4, 0)
        kernel = crc32c(self.data[SUPER + 8 : self.checksum_end], 0x5045B54A)
        return expected == crc32c(canonical) == kernel

    def update_checksum(self) -> None:
        self.put_u32(SUPER + 4, 0)
        checksum = crc32c(self.data[SUPER : self.checksum_end])
        self.put_u32(SUPER + 4, checksum)
        if not self.checksum_valid():
            raise AssertionError("updated B21 checksum does not verify")

    def inode_offset(self, nid: int) -> int:
        return (self.u32(SUPER + 40) << self.block_bits) + (nid << 5)

    def fix_packed_inode_mode(self) -> None:
        packed_nid = self.u64(SUPER + 96)
        if packed_nid == 0:
            raise ValueError("B21 seed did not materialize a packed prefix carrier")
        self.put_u16(self.inode_offset(packed_nid) + 4, 0o100644)


def load_spec(path: Path) -> dict:
    spec = json.loads(path.read_text(encoding="ascii"))
    if spec.get("schema") != 1 or spec.get("batch") != "B21":
        raise ValueError("invalid B21 fixture spec identity")
    if not isinstance(spec.get("cases"), list) or not spec["cases"]:
        raise ValueError("B21 fixture spec has no cases")
    return spec


def create_source(root: Path) -> None:
    root.mkdir(parents=True)
    control = root / "control.txt"
    prefix = root / "prefix.bin"
    control.write_bytes(b"B21 superblock control\n")
    prefix.write_bytes(b"B21 prefix control\n")
    control.chmod(0o644)
    prefix.chmod(0o644)
    os.setxattr(prefix, b"user.company.branch.leaf", b"B21-prefix-value")
    for path in (control, prefix, root):
        os.utime(path, (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 for B21 fixtures")
    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",
        "-Eforce-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)
    image.fix_packed_inode_mode()
    image.update_checksum()
    output.write_bytes(image.data)
    if not image.u32(SUPER + 80) & FEATURE_INCOMPAT_XATTR_PREFIXES:
        raise ValueError("B21 seed does not declare xattr prefixes")
    return {"version": version}


def mutate(image: Image, mutation: str) -> list[dict[str, int | str]]:
    changes: list[dict[str, int | str]] = []

    def record(field: str, offset: int, size: int) -> None:
        changes.append({"field": field, "offset": offset, "size": size})

    if mutation == "none":
        return changes
    if mutation in {"unknown-feature", "feature-prefix"}:
        offset = SUPER + 80
        image.put_u32(offset, image.u32(offset) | UNKNOWN_INCOMPAT)
        record("feature_incompat", offset, 4)
    if mutation == "blocks":
        offset = SUPER + 36
        image.put_u32(offset, 0xFFFFFFFF)
        record("blocks_lo", offset, 4)
    elif mutation == "extslots":
        offset = SUPER + 13
        image.data[offset] = 0xFF
        record("sb_extslots", offset, 1)
    elif mutation == "dirblk":
        offset = SUPER + 90
        image.data[offset] = 1
        record("dirblkbits", offset, 1)
    elif mutation in {"prefix-offset", "feature-prefix"}:
        offset = SUPER + 92
        image.put_u32(offset, 0xFFFFFFFF)
        record("xattr_prefix_start", offset, 4)
    elif mutation != "unknown-feature":
        raise ValueError(f"unknown B21 mutation: {mutation}")
    return changes


def generate(spec_path: Path, output: Path, work: Path) -> int:
    spec = load_spec(spec_path)
    if output.exists() or work.exists():
        raise SystemExit("B21 output and work paths must not already exist")
    output.mkdir(parents=True)
    work.mkdir(parents=True)
    seed_path = work / "seed.erofs"
    build = build_seed(spec, seed_path, work / "seed-work")
    seed = Image.load(seed_path)
    seed_hash = sha256(seed_path)
    cases = []
    for item in spec["cases"]:
        image = seed.clone()
        changes = mutate(image, item["mutation"])
        if item["authenticated"] and item["mutation"] != "none":
            image.update_checksum()
        path = output / f"{item['id']}.erofs"
        path.write_bytes(image.data)
        checksum_valid = image.checksum_valid()
        if checksum_valid != item["authenticated"]:
            raise AssertionError(f"checksum state mismatch for {item['id']}")
        cases.append(
            {
                **item,
                "changes": changes,
                "checksum_valid": checksum_valid,
                "filename": path.name,
                "sha256": sha256(path),
                "size": path.stat().st_size,
            }
        )
    index = {
        "batch": "B21",
        "build": build,
        "case_count": len(cases),
        "cases": cases,
        "damaged_count": sum(item["class"] == "damaged" for item in cases),
        "legal_count": sum(item["class"] == "legal" for item in cases),
        "schema": 1,
        "seed_sha256": seed_hash,
    }
    write_json(output / "fixture-index.json", index)
    print(json.dumps({"status": "PASS", **index}, sort_keys=True))
    return 0


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--spec", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--work", type=Path, required=True)
    args = parser.parse_args()
    return generate(args.spec, args.output, args.work)


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