#!/usr/bin/env python3
"""Generate deterministic EROFS explicit-extent fixtures for Pre15 B23."""

from __future__ import annotations

import argparse
from dataclasses import dataclass
import hashlib
import json
import os
from pathlib import Path
import shutil
import stat
import struct
import subprocess


SUPER = 1024
MAGIC = 0xE0F5E1E2
FEATURE_COMPAT_SB_CHKSUM = 0x00000001
FEATURE_INCOMPAT_METABOX = 0x00000100
METABOX_NID_BIT = 1 << 63
LAYOUT_FLAT_PLAIN = 0
LAYOUT_COMPRESSED_FULL = 1
Z_EROFS_ADVISE_EXTENTS = 0x0001
BLOCK_SIZE = 4096
CARRIER_BYTES = 512 * 1024
CRC32C_POLYNOMIAL = 0x82F63B78
FIXED_UUID = "23232323-1515-4234-8123-000000000023"


class FixtureError(RuntimeError):
    pass


@dataclass(frozen=True)
class Inode:
    nid: int
    offset: int
    inode_size: int
    layout: int
    size: int
    start_block: int


@dataclass(frozen=True)
class DirectoryEntry:
    name: bytes
    nid: int
    offset: int


def align(value: int, alignment: int) -> int:
    return (value + alignment - 1) & -alignment


def sha256_bytes(data: bytes | bytearray) -> str:
    return hashlib.sha256(data).hexdigest()


def crc32c(data: bytes | bytearray) -> int:
    checksum = 0xFFFFFFFF
    for byte in data:
        checksum ^= byte
        for _ in range(8):
            checksum = (checksum >> 1) ^ (
                CRC32C_POLYNOMIAL if checksum & 1 else 0
            )
    return checksum & 0xFFFFFFFF


class Image:
    def __init__(self, data: bytes | bytearray) -> None:
        self.data = bytearray(data)
        self.validate_super()

    @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_u8(self, offset: int, value: int) -> None:
        self.data[offset] = value

    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 put_u64(self, offset: int, value: int) -> None:
        struct.pack_into("<Q", 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 meta_offset(self) -> int:
        return self.u32(SUPER + 40) << self.block_bits

    @property
    def root_nid(self) -> int:
        return self.u16(SUPER + 14)

    @property
    def blocks(self) -> int:
        return self.u32(SUPER + 36)

    @property
    def declared_size(self) -> int:
        return self.blocks << self.block_bits

    def validate_super(self) -> None:
        if len(self.data) < SUPER + 144 or self.u32(SUPER) != MAGIC:
            raise FixtureError("invalid EROFS superblock")
        if self.block_size != BLOCK_SIZE:
            raise FixtureError("B23 requires a 4096-byte EROFS block")
        if self.declared_size > len(self.data):
            raise FixtureError("declared EROFS image exceeds provider bytes")
        if self.u32(SUPER + 8) & FEATURE_COMPAT_SB_CHKSUM:
            window = bytearray(self.data[SUPER : self.block_size])
            expected = struct.unpack_from("<I", window, 4)[0]
            struct.pack_into("<I", window, 4, 0)
            if crc32c(window) != expected:
                raise FixtureError("EROFS superblock checksum mismatch")

    def update_checksum(self) -> None:
        if not self.u32(SUPER + 8) & FEATURE_COMPAT_SB_CHKSUM:
            raise FixtureError("B23 seed lacks the checksum feature")
        self.put_u32(SUPER + 4, 0)
        self.put_u32(SUPER + 4, crc32c(self.data[SUPER : self.block_size]))

    def inode(self, nid: int) -> Inode:
        offset = self.meta_offset + (nid << 5)
        if offset + 32 > self.declared_size:
            raise FixtureError(f"inode {nid} exceeds the primary metadata")
        inode_format = self.u16(offset)
        inode_size = 64 if inode_format & 1 else 32
        size = self.u64(offset + 8) if inode_size == 64 else self.u32(offset + 8)
        return Inode(
            nid=nid,
            offset=offset,
            inode_size=inode_size,
            layout=(inode_format >> 1) & 7,
            size=size,
            start_block=self.u32(offset + 16),
        )

    def inode_data(self, inode: Inode) -> tuple[bytes, int]:
        if inode.layout == LAYOUT_FLAT_PLAIN:
            start = inode.start_block << self.block_bits
            end = start + inode.size
            if end > self.declared_size:
                raise FixtureError("plain inode data exceeds the image")
            return bytes(self.data[start:end]), start
        if inode.layout != 2:
            raise FixtureError(f"unsupported host inode layout {inode.layout}")
        start = inode.offset + inode.inode_size
        end = start + inode.size
        if end > self.declared_size:
            raise FixtureError("inline inode data exceeds the image")
        return bytes(self.data[start:end]), start

    def directory_entries(self, inode: Inode) -> list[DirectoryEntry]:
        data, data_base = self.inode_data(inode)
        entries: list[DirectoryEntry] = []
        for block_start in range(0, len(data), self.block_size):
            block = data[block_start : block_start + self.block_size]
            if len(block) < 12:
                raise FixtureError("short directory block")
            first_name = struct.unpack_from("<H", block, 8)[0]
            if first_name < 12 or first_name % 12 or first_name > len(block):
                raise FixtureError("invalid directory entry table")
            count = first_name // 12
            for index in range(count):
                item = index * 12
                name_start = struct.unpack_from("<H", block, item + 8)[0]
                name_end = (
                    struct.unpack_from("<H", block, item + 20)[0]
                    if index + 1 < count
                    else len(block)
                )
                name = block[name_start:name_end].split(b"\0", 1)[0]
                entries.append(
                    DirectoryEntry(
                        name=name,
                        nid=struct.unpack_from("<Q", block, item)[0],
                        offset=data_base + block_start + item,
                    )
                )
        return entries

    def resolve_root(self, name: str) -> tuple[DirectoryEntry, Inode]:
        encoded = name.encode("ascii")
        for entry in self.directory_entries(self.inode(self.root_nid)):
            if entry.name == encoded:
                return entry, self.inode(entry.nid)
        raise FixtureError(f"root entry not found: {name}")

    def ensure_size(self, end: int) -> None:
        rounded = align(end, self.block_size)
        if len(self.data) < rounded:
            self.data.extend(bytes(rounded - len(self.data)))
        self.put_u32(SUPER + 36, rounded >> self.block_bits)


def normalize_times(root: Path) -> None:
    for path in sorted(root.rglob("*"), reverse=True):
        os.utime(path, (0, 0), follow_symlinks=False)
    os.utime(root, (0, 0), follow_symlinks=False)


def payload_block() -> bytes:
    return bytes((index * 37 + 11) & 0xFF for index in range(BLOCK_SIZE))


def make_seed(output: Path, expected_version: str) -> Image:
    source = output / ".source"
    source.mkdir()
    (source / "target.bin").write_bytes(b"placeholder\n")
    (source / "payload.bin").write_bytes(payload_block())
    (source / "carrier.bin").write_bytes(bytes(CARRIER_BYTES))
    normalize_times(source)
    version_output = subprocess.run(
        ["mkfs.erofs", "-V"], check=True, text=True,
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
    ).stdout.strip()
    version = version_output.splitlines()[0] if version_output else ""
    if version != expected_version:
        raise FixtureError(f"mkfs version changed: {version!r}")
    seed = output / ".seed.erofs"
    command = [
        "mkfs.erofs", "-d0", "-x-1", "-T0", "--all-time", "--all-root",
        "--workers=1", "--sort=path", "-U", FIXED_UUID,
        "-E", "force-inode-extended", str(seed), str(source),
    ]
    subprocess.run(command, check=True)
    shutil.rmtree(source)
    image = Image.load(seed)
    _, payload = image.resolve_root("payload.bin")
    _, carrier = image.resolve_root("carrier.bin")
    if payload.layout != LAYOUT_FLAT_PLAIN or payload.size != BLOCK_SIZE:
        raise FixtureError("payload seed is not one plain block")
    if carrier.layout != LAYOUT_FLAT_PLAIN or carrier.size != CARRIER_BYTES:
        raise FixtureError("carrier seed is not the expected plain file")
    return image


def write_inode(image: Image, offset: int, size: int) -> None:
    image.data[offset : offset + 64] = bytes(64)
    image.put_u16(offset, 1 | (LAYOUT_COMPRESSED_FULL << 1))
    image.put_u16(offset + 4, stat.S_IFREG | 0o444)
    image.put_u64(offset + 8, size)
    image.put_u32(offset + 16, 1)
    image.put_u32(offset + 44, 1)


def logical_starts(case: dict[str, object]) -> tuple[list[int], int, list[int]]:
    count = int(case["count"])
    mutation = case.get("mutation")
    if case.get("logical_mode") == "high32":
        starts = [0, 1 << 32, (1 << 32) + BLOCK_SIZE]
        return starts, (1 << 32) + 2 * BLOCK_SIZE, starts[1:]
    if mutation == "descending-lstart":
        starts = [0, 2 * BLOCK_SIZE, 3 * BLOCK_SIZE]
        return starts, 4 * BLOCK_SIZE, []
    starts = [index * BLOCK_SIZE for index in range(count)]
    size = count * BLOCK_SIZE
    probes = sorted({0, ((count // 2) * BLOCK_SIZE), size - BLOCK_SIZE})
    return starts, size, probes


def extent_span(record_size: int, count: int) -> int:
    return 72 + (8 if record_size == 4 else 0) + record_size * count


def write_extent_table(
    image: Image,
    storage_base: int,
    inode_offset: int,
    record_size: int,
    starts: list[int],
    file_size: int,
    payload_offset: int,
) -> dict[str, int]:
    write_inode(image, storage_base + inode_offset, file_size)
    header = align(inode_offset + 64, 8)
    table = align(header + 8, record_size)
    advise = Z_EROFS_ADVISE_EXTENTS | ({4: 0, 8: 1, 16: 2, 32: 3}[record_size] << 1)
    struct.pack_into(
        "<IHH", image.data, storage_base + header,
        len(starts) & 0xFFFFFFFF, advise, len(starts) >> 32,
    )
    records = table
    if record_size == 4:
        image.put_u64(storage_base + table, payload_offset)
        records += 8
    for index, lstart in enumerate(starts):
        offset = storage_base + records + index * record_size
        image.data[offset : offset + record_size] = bytes(record_size)
        image.put_u32(offset, BLOCK_SIZE)
        if record_size >= 8:
            image.put_u32(offset + 4, payload_offset & 0xFFFFFFFF)
        if record_size >= 16:
            image.put_u32(offset + 8, payload_offset >> 32)
            image.put_u32(offset + 12, lstart & 0xFFFFFFFF)
        if record_size == 32:
            image.put_u32(offset + 16, lstart >> 32)
    return {
        "header": header,
        "table": table,
        "records": records,
        "end": records + record_size * len(starts),
    }


def enable_metabox(image: Image) -> None:
    if image.meta_offset != 0:
        raise FixtureError("B23 metabox seed expects metadata block zero")
    metadata = bytes(image.data[: image.block_size])
    relocated = align(len(image.data), image.block_size)
    image.ensure_size(relocated + image.block_size)
    image.data[relocated : relocated + image.block_size] = metadata
    image.put_u32(SUPER + 40, relocated >> image.block_bits)
    image.put_u8(SUPER + 13, 1)
    image.put_u32(
        SUPER + 80, image.u32(SUPER + 80) | FEATURE_INCOMPAT_METABOX
    )


def save_case(
    output: Path,
    seed: Image,
    case: dict[str, object],
    payload_offset: int,
) -> dict[str, object]:
    image = seed.clone()
    backing = str(case["backing"])
    record_size = int(case["record_size"])
    starts, file_size, probes = logical_starts(case)
    stored_count = len(starts)
    mutation = case.get("mutation")
    if backing == "primary":
        inode_absolute = align(len(image.data), image.block_size) + 4000
        inode_offset = inode_absolute
        storage_base = 0
        stored_end = inode_absolute + extent_span(record_size, stored_count)
        image.ensure_size(stored_end)
        target_entry, _ = image.resolve_root("target.bin")
        nid_delta = inode_absolute - image.meta_offset
        if nid_delta < 0 or nid_delta % 32:
            raise FixtureError("primary synthetic inode is not NID aligned")
        target_nid = nid_delta >> 5
        image.put_u64(target_entry.offset, target_nid)
        carrier_inode_offset = None
        carrier_size = None
    elif backing == "metabox":
        enable_metabox(image)
        target_entry, _ = image.resolve_root("target.bin")
        _, carrier = image.resolve_root("carrier.bin")
        storage_base = carrier.start_block << image.block_bits
        span = extent_span(record_size, stored_count)
        inode_offset = ((CARRIER_BYTES - span) // 32) * 32
        carrier_size = inode_offset + span
        if inode_offset < 4096 or carrier_size > CARRIER_BYTES:
            raise FixtureError("metabox explicit table does not fit the carrier")
        target_nid = METABOX_NID_BIT | (inode_offset >> 5)
        image.put_u64(target_entry.offset, target_nid)
        carrier_inode_offset = carrier.offset
        image.put_u64(carrier.offset + 8, carrier_size)
        image.put_u64(SUPER + 128, carrier.nid)
    else:
        raise FixtureError(f"unknown backing: {backing}")

    layout = write_extent_table(
        image, storage_base, inode_offset, record_size, starts, file_size,
        payload_offset,
    )
    image.update_checksum()
    baseline = image.clone()
    mutation_record = None
    if mutation == "missing-tail":
        image.put_u32(storage_base + layout["header"], 2)
        mutation_record = {
            "field": "map_header.h_extents_lo",
            "offset": storage_base + layout["header"],
            "size": 4,
        }
    elif mutation == "huge-count":
        image.put_u16(storage_base + layout["header"] + 6, 256)
        mutation_record = {
            "field": "map_header.h_extents_hi",
            "offset": storage_base + layout["header"] + 6,
            "size": 2,
        }
    elif mutation == "descending-lstart":
        changed = storage_base + layout["records"] + 2 * record_size + 12
        image.put_u32(changed, BLOCK_SIZE)
        mutation_record = {
            "field": "extent[2].lstart_lo",
            "offset": changed,
            "size": 4,
        }
        starts[2] = BLOCK_SIZE
    elif mutation == "tail-minus-one":
        if carrier_inode_offset is None or carrier_size is None:
            raise FixtureError("tail-minus-one requires metabox backing")
        image.put_u64(carrier_inode_offset + 8, carrier_size - 1)
        mutation_record = {
            "field": "metabox_carrier.i_size",
            "offset": carrier_inode_offset + 8,
            "size": 8,
        }
    elif mutation is not None:
        raise FixtureError(f"unknown mutation: {mutation}")

    name = str(case["name"])
    image_path = output / f"{name}.erofs"
    image_path.write_bytes(image.data)
    baseline_name = None
    if mutation_record is not None:
        baseline_name = f"{name}-baseline.erofs"
        (output / baseline_name).write_bytes(baseline.data)

    declared_count = stored_count
    if mutation == "missing-tail":
        declared_count = 2
    elif mutation == "huge-count":
        declared_count = 1 + (256 << 32)
    table_bytes = record_size * declared_count
    scan_calls = 1 if record_size <= 8 else (
        (table_bytes + 65535) // 65536
        if case["outcome"] == "PASS" else 1
    )
    return {
        "backing": backing,
        "baseline": baseline_name,
        "carrier_inode_offset": carrier_inode_offset,
        "carrier_size": carrier_size,
        "declared_count": declared_count,
        "expected_errno": 0 if case["outcome"] == "PASS" else 97,
        "expected_scan_calls": scan_calls,
        "file_size": file_size,
        "header_image_offset": storage_base + layout["header"],
        "header_logical_offset": layout["header"],
        "image": image_path.name,
        "image_sha256": sha256_bytes(image.data),
        "inode_image_offset": storage_base + inode_offset,
        "inode_logical_offset": inode_offset,
        "lstarts": starts,
        "mutation": mutation_record,
        "payload_offset": payload_offset,
        "probe_offsets": probes,
        "record_image_offset": storage_base + layout["records"],
        "record_logical_offset": layout["records"],
        "record_size": record_size,
        "stored_count": stored_count,
        "table_end_logical": layout["end"],
        "target_dirent_offset": target_entry.offset,
        "target_marker": case["target_marker"],
        "target_nid": target_nid,
    }


def generate(spec_path: Path, output: Path) -> None:
    spec = json.loads(spec_path.read_text(encoding="ascii"))
    if spec.get("schema") != 1 or spec.get("batch") != "B23":
        raise FixtureError("invalid B23 spec identity")
    if output.exists() and any(output.iterdir()):
        raise FixtureError(f"output is not empty: {output}")
    output.mkdir(parents=True, exist_ok=True)
    seed = make_seed(output, str(spec["mkfs_version"]))
    _, payload = seed.resolve_root("payload.bin")
    payload_offset = payload.start_block << seed.block_bits
    if seed.data[payload_offset : payload_offset + BLOCK_SIZE] != payload_block():
        raise FixtureError("plain payload bytes changed")

    cases: dict[str, object] = {}
    for case in spec["cases"]:
        cases[str(case["name"])] = save_case(
            output, seed, case, payload_offset
        )
    (output / ".seed.erofs").unlink()
    manifest = {
        "batch": "B23",
        "cases": cases,
        "chunk_size": spec["chunk_size"],
        "mkfs_version": spec["mkfs_version"],
        "payload_sha256": sha256_bytes(payload_block()),
        "schema": 1,
        "test": spec["test"],
    }
    (output / "fixture-manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=True, indent=2, sort_keys=True) + "\n",
        encoding="ascii",
    )
    with (output / "SHA256SUMS").open("w", encoding="ascii") as sums:
        for path in sorted(output.glob("*.erofs")):
            sums.write(f"{sha256_bytes(path.read_bytes())}  {path.name}\n")
    print(
        f"B23 fixtures generated cases={len(cases)} "
        f"images={len(list(output.glob('*.erofs')))}"
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--spec", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    generate(args.spec, args.output)


if __name__ == "__main__":
    try:
        main()
    except (FixtureError, OSError, subprocess.CalledProcessError) as error:
        raise SystemExit(f"B23 fixture generation failed: {error}") from error