#!/usr/bin/env python3
"""Build and self-check fixtures for TC154-TC156."""

from __future__ import annotations

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


SUPER = 1024
EROFS_MAGIC = 0xE0F5E1E2
FEATURE_COMPAT_SB_CHKSUM = 0x1
EROFS_INODE_FLAT_PLAIN = 0
EROFS_INODE_FLAT_INLINE = 2
EROFS_INODE_COMPRESSED_FULL = 1
Z_EROFS_ADVISE_EXTENTS = 0x1
UINT64_MAX = (1 << 64) - 1
INT64_MAX = (1 << 63) - 1
FNV1_32_INIT = 33554467
FNV_32_PRIME = 0x01000193


class FixtureError(RuntimeError):
    pass


class Inode:
    def __init__(
        self,
        nid: int,
        offset: int,
        inode_format: int,
        inode_size: int,
        xattr_size: int,
        layout: int,
        size: int,
        start_block: int,
    ) -> None:
        self.nid = nid
        self.offset = offset
        self.inode_format = inode_format
        self.inode_size = inode_size
        self.xattr_size = xattr_size
        self.layout = layout
        self.size = size
        self.start_block = start_block


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

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

    def clone(self) -> "ErofsImage":
        return ErofsImage(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)

    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:
        if not 9 <= self.block_bits <= 16:
            raise FixtureError(f"invalid block bits: {self.block_bits}")
        return 1 << self.block_bits

    @property
    def super_size(self) -> int:
        return 128 + self.data[SUPER + 13] * 16

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

    @property
    def feature_compat(self) -> int:
        return self.u32(SUPER + 8)

    @property
    def root_nid(self) -> int:
        feature_incompat = self.u32(SUPER + 80)
        root_nid_8b = self.u64(SUPER + 112)
        if feature_incompat & 0x80 and root_nid_8b != 0:
            return root_nid_8b
        return self.u16(SUPER + 14)

    def calculated_checksum(self) -> int:
        window = bytearray(self.data[SUPER : self.checksum_end])
        struct.pack_into("<I", window, 4, 0)
        checksum = 0xFFFFFFFF
        polynomial = 0x82F63B78
        for byte in window:
            checksum ^= byte
            for _ in range(8):
                checksum = (checksum >> 1) ^ (
                    polynomial if checksum & 1 else 0
                )
        return checksum & 0xFFFFFFFF

    def checksum_valid(self) -> bool:
        if not self.feature_compat & FEATURE_COMPAT_SB_CHKSUM:
            return True
        return self.u32(SUPER + 4) == self.calculated_checksum()

    def update_checksum(self) -> None:
        if not self.feature_compat & FEATURE_COMPAT_SB_CHKSUM:
            raise FixtureError("base image lacks superblock checksum support")
        self.put_u32(SUPER + 4, 0)
        self.put_u32(SUPER + 4, self.calculated_checksum())
        if not self.checksum_valid():
            raise FixtureError("updated superblock checksum does not verify")

    def validate_superblock(self) -> None:
        if len(self.data) < SUPER + 144:
            raise FixtureError("image is too short for an EROFS superblock")
        if self.u32(SUPER) != EROFS_MAGIC:
            raise FixtureError(f"bad EROFS magic: {self.u32(SUPER):#x}")
        if SUPER + self.super_size > len(self.data):
            raise FixtureError("declared superblock is truncated")
        if not self.checksum_valid():
            raise FixtureError("superblock checksum is invalid")

    def inode(self, nid: int) -> Inode:
        metadata = self.u32(SUPER + 40) << self.block_bits
        offset = metadata + (nid << 5)
        if offset > len(self.data) - 32:
            raise FixtureError(f"nid {nid} lies outside the image")
        inode_format = self.u16(offset)
        inode_size = 64 if inode_format & 1 else 32
        if offset > len(self.data) - inode_size:
            raise FixtureError(f"nid {nid} has a truncated inode")
        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 Inode(
            nid=nid,
            offset=offset,
            inode_format=inode_format,
            inode_size=inode_size,
            xattr_size=xattr_size,
            layout=(inode_format >> 1) & 7,
            size=size,
            start_block=self.u32(offset + 16),
        )

    def directory_data(self, inode: Inode) -> bytes:
        if inode.size == 0:
            raise FixtureError("empty directory cannot contain a target")
        if inode.layout == EROFS_INODE_FLAT_PLAIN:
            offset = inode.start_block << self.block_bits
            end = offset + inode.size
            if end > len(self.data):
                raise FixtureError("plain directory data is truncated")
            return bytes(self.data[offset:end])
        if inode.layout != EROFS_INODE_FLAT_INLINE:
            raise FixtureError(f"unsupported directory layout {inode.layout}")
        tail_start = ((inode.size + self.block_size - 1) // self.block_size - 1)
        tail_start *= self.block_size
        full_offset = inode.start_block << self.block_bits
        inline_offset = inode.offset + inode.inode_size + inode.xattr_size
        full = self.data[full_offset : full_offset + tail_start]
        tail = self.data[inline_offset : inline_offset + inode.size - tail_start]
        if len(full) + len(tail) != inode.size:
            raise FixtureError("inline directory data is truncated")
        return bytes(full + tail)

    def directory_entries(self, inode: Inode) -> list[tuple[bytes, int]]:
        data = self.directory_data(inode)
        entries: list[tuple[bytes, int]] = []
        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_nameoff = struct.unpack_from("<H", block, 8)[0]
            if (
                first_nameoff < 12
                or first_nameoff % 12 != 0
                or first_nameoff >= len(block)
            ):
                raise FixtureError("invalid first directory name offset")
            count = first_nameoff // 12
            previous = 0
            for index in range(count):
                entry_offset = index * 12
                nid = struct.unpack_from("<Q", block, entry_offset)[0]
                nameoff = struct.unpack_from("<H", block, entry_offset + 8)[0]
                endoff = (
                    struct.unpack_from("<H", block, entry_offset + 20)[0]
                    if index + 1 < count
                    else len(block)
                )
                if (
                    nameoff < first_nameoff
                    or nameoff <= previous
                    or endoff <= nameoff
                    or endoff > len(block)
                ):
                    raise FixtureError("invalid directory name range")
                name = block[nameoff:endoff].split(b"\0", 1)[0]
                if not name:
                    raise FixtureError("empty directory name")
                entries.append((name, nid))
                previous = nameoff
        return entries

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

    def inode_bytes(self, inode: Inode) -> bytes:
        return bytes(self.data[inode.offset : inode.offset + inode.inode_size])

    def generation(self, inode: Inode) -> int:
        seed = fnv1_32(
            bytes(self.data[SUPER : SUPER + self.super_size]), FNV1_32_INIT
        )
        generation = fnv1_32(struct.pack("<Q", inode.nid), seed)
        generation = fnv1_32(self.inode_bytes(inode), generation)
        return generation or 1

    def save(self, path: Path) -> None:
        self.validate_superblock()
        path.write_bytes(self.data)


def fnv1_32(data: bytes, initial: int) -> int:
    value = initial
    for byte in data:
        value = (value * FNV_32_PRIME) & 0xFFFFFFFF
        value ^= byte
    return value


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


def run_mkfs(source: Path, image: Path, uuid: str, *extra: str) -> None:
    command = [
        "mkfs.erofs",
        "-d0",
        "-x-1",
        "-T0",
        "--all-time",
        "--all-root",
        "--workers=1",
        "--sort=path",
        "-U",
        uuid,
        *extra,
        str(image),
        str(source),
    ]
    subprocess.run(command, check=True)


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 sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def make_fixtures(output: Path) -> None:
    if shutil.which("mkfs.erofs") is None:
        raise FixtureError("mkfs.erofs is required")
    if output.exists() and any(output.iterdir()):
        raise FixtureError(f"output directory is not empty: {output}")
    output.mkdir(parents=True, exist_ok=True)
    source = output / "source"
    nfs_source = source / "nfs"
    extent_source = source / "extent"
    extended_source = source / "extended"
    nfs_source.mkdir(parents=True)
    extent_source.mkdir(parents=True)
    extended_source.mkdir(parents=True)

    for index in range(192):
        (nfs_source / f"filler-{index:03d}.txt").write_text(
            f"filler {index:03d}\n", encoding="ascii"
        )
    (nfs_source / "zz-identity.txt").write_text(
        "stable inode identity\n", encoding="ascii"
    )
    (extent_source / "extent.bin").write_bytes(b"E" * (1024 * 1024))
    (extended_source / "extended-time.txt").write_text(
        "extended timestamp\n", encoding="ascii"
    )
    normalize_times(source)

    nfs_base_path = output / ".nfs-base.erofs"
    extent_base_path = output / ".extent-base.erofs"
    extended_base_path = output / ".extended-base.erofs"
    run_mkfs(
        nfs_source,
        nfs_base_path,
        "66666666-7777-4888-9999-aaaaaaaaaa54",
        "-E",
        "force-inode-compact",
    )
    run_mkfs(
        extent_source,
        extent_base_path,
        "66666666-7777-4888-9999-aaaaaaaaaa55",
        "-E",
        "legacy-compress,force-inode-extended",
        "-z",
        "lz4",
        "-C4096",
    )
    run_mkfs(
        extended_source,
        extended_base_path,
        "66666666-7777-4888-9999-aaaaaaaaaa56",
        "-E",
        "force-inode-extended",
    )

    evidence: list[str] = []
    manifest: dict[str, object] = {}

    nfs_a = ErofsImage.load(nfs_base_path)
    identity_a = nfs_a.resolve_root_entry("zz-identity.txt")
    if identity_a.inode_size != 32:
        raise FixtureError("NFS identity target is not a compact inode")
    if identity_a.offset < nfs_a.checksum_end:
        raise FixtureError("NFS identity inode overlaps the checksum window")
    nfs_a_path = output / "nfs-inode-a.erofs"
    nfs_a.save(nfs_a_path)

    nfs_b = nfs_a.clone()
    identity_b = nfs_b.resolve_root_entry("zz-identity.txt")
    old_mtime = nfs_b.u32(identity_b.offset + 12)
    nfs_b.put_u32(identity_b.offset + 12, old_mtime + 1)
    nfs_b_path = output / "nfs-inode-b.erofs"
    nfs_b.save(nfs_b_path)
    nfs_differences = [
        offset
        for offset, (left, right) in enumerate(zip(nfs_a.data, nfs_b.data))
        if left != right
    ]
    mtime_offsets = range(identity_b.offset + 12, identity_b.offset + 16)
    if not nfs_differences or any(
        offset not in mtime_offsets for offset in nfs_differences
    ):
        raise FixtureError("NFS replacement changed data outside i_mtime")
    if nfs_a.data[: nfs_a.checksum_end] != nfs_b.data[: nfs_b.checksum_end]:
        raise FixtureError("NFS replacement changed the superblock block")
    if nfs_a.data[SUPER : SUPER + nfs_a.super_size] != nfs_b.data[
        SUPER : SUPER + nfs_b.super_size
    ]:
        raise FixtureError("NFS replacement changed the declared superblock")
    if nfs_a.inode_bytes(identity_a) == nfs_b.inode_bytes(identity_b):
        raise FixtureError("NFS replacement did not change inode metadata")
    generation_a = nfs_a.generation(identity_a)
    generation_b = nfs_b.generation(identity_b)
    if generation_a == generation_b:
        raise FixtureError("per-inode generation did not change")

    nfs_corrupt = nfs_a.clone()
    corrupt_inode = nfs_corrupt.resolve_root_entry("zz-identity.txt")
    nfs_corrupt.put_u16(
        corrupt_inode.offset, nfs_corrupt.u16(corrupt_inode.offset) | 0x8000
    )
    nfs_corrupt_path = output / "nfs-inode-corrupt.erofs"
    nfs_corrupt.save(nfs_corrupt_path)
    if nfs_a.data[: nfs_a.checksum_end] != nfs_corrupt.data[: nfs_a.checksum_end]:
        raise FixtureError("corrupt inode fixture changed the superblock block")

    evidence.append(
        "nfs "
        f"nid={identity_a.nid} inode_offset={identity_a.offset} "
        f"checksum_end={nfs_a.checksum_end} mtime={old_mtime}->{old_mtime + 1} "
        f"generation={generation_a}->{generation_b} superblock_block=identical"
    )
    manifest["nfs"] = {
        "path": "/zz-identity.txt",
        "nid": identity_a.nid,
        "inode_offset": identity_a.offset,
        "generation_a": generation_a,
        "generation_b": generation_b,
        "super_size": nfs_a.super_size,
        "checksum_end": nfs_a.checksum_end,
    }

    extent = ErofsImage.load(extent_base_path)
    extent_inode = extent.resolve_root_entry("extent.bin")
    if (
        extent_inode.inode_size != 64
        or extent_inode.layout != EROFS_INODE_COMPRESSED_FULL
    ):
        raise FixtureError("extent target is not an extended full-index inode")
    header = align(
        extent_inode.offset + extent_inode.inode_size + extent_inode.xattr_size,
        8,
    )
    physical_base_offset = align(header + 8, 4)
    record_offset = physical_base_offset + 8
    if record_offset + 8 > len(extent.data):
        raise FixtureError("extent records lie outside the base image")
    root_inode = extent.inode(extent.root_nid)
    if not (
        record_offset + 8 <= root_inode.offset
        or root_inode.offset + root_inode.inode_size <= header
    ):
        raise FixtureError("extent conversion overlaps the root inode")
    struct.pack_into(
        "<IHH", extent.data, header, 2, Z_EROFS_ADVISE_EXTENTS, 0
    )
    physical_base = UINT64_MAX - 4095
    extent.put_u64(physical_base_offset, physical_base)
    extent.put_u32(record_offset, 8192)
    extent.put_u32(record_offset + 4, 4096)
    extent.update_checksum()
    extent_path = output / "extent-pa-wrap.erofs"
    extent.save(extent_path)
    cluster_size = 1 << (extent.block_bits + (extent.data[header + 7] & 15))
    if extent.u32(header) != 2 or extent.u16(header + 6) != 0:
        raise FixtureError("extent header does not declare exactly two records")
    if extent.u16(header + 4) != Z_EROFS_ADVISE_EXTENTS:
        raise FixtureError("extent advise does not select 4-byte records")
    if cluster_size != 4096 or extent_inode.size <= cluster_size:
        raise FixtureError("extent fixture lacks a second 4 KiB logical cluster")
    if extent.u64(physical_base_offset) + extent.u32(record_offset) <= UINT64_MAX:
        raise FixtureError("extent fixture does not overflow pa + plen")
    evidence.append(
        "extent-wrap "
        f"nid={extent_inode.nid} inode_offset={extent_inode.offset} "
        f"header={header} base_offset={physical_base_offset} "
        f"record_offset={record_offset} recsize=4 pa=0x{physical_base:016x} "
        "plen0=8192 plen1=4096 logical_probe=4096 overflow=yes"
    )
    manifest["extent_wrap"] = {
        "path": "/extent.bin",
        "nid": extent_inode.nid,
        "inode_offset": extent_inode.offset,
        "header_offset": header,
        "physical_base_offset": physical_base_offset,
        "record_offset": record_offset,
        "record_size": 4,
        "cluster_size": cluster_size,
        "physical_base": physical_base,
        "plen": [8192, 4096],
        "logical_probe": 4096,
    }

    compact_epoch_range = nfs_a.clone()
    range_inode = compact_epoch_range.resolve_root_entry("zz-identity.txt")
    compact_epoch_range.put_u64(SUPER + 24, INT64_MAX)
    compact_epoch_range.put_u32(range_inode.offset + 12, 1)
    compact_epoch_range.update_checksum()
    compact_epoch_range_path = output / "compact-epoch-range.erofs"
    compact_epoch_range.save(compact_epoch_range_path)

    compact_epoch_wrap = nfs_a.clone()
    wrap_root = compact_epoch_wrap.inode(compact_epoch_wrap.root_nid)
    compact_epoch_wrap.put_u64(SUPER + 24, UINT64_MAX)
    compact_epoch_wrap.put_u32(wrap_root.offset + 12, 1)
    compact_epoch_wrap.update_checksum()
    compact_epoch_wrap_path = output / "compact-epoch-wrap.erofs"
    compact_epoch_wrap.save(compact_epoch_wrap_path)

    compact_nsec = nfs_a.clone()
    compact_nsec.put_u32(SUPER + 32, 1_000_000_000)
    compact_nsec.update_checksum()
    compact_nsec_path = output / "compact-nsec-invalid.erofs"
    compact_nsec.save(compact_nsec_path)

    extended_base = ErofsImage.load(extended_base_path)
    extended_inode = extended_base.resolve_root_entry("extended-time.txt")
    if extended_inode.inode_size != 64:
        raise FixtureError("extended timestamp target is not extended")
    extended_nsec = extended_base.clone()
    extended_nsec_inode = extended_nsec.resolve_root_entry("extended-time.txt")
    extended_nsec.put_u32(extended_nsec_inode.offset + 40, 1_000_000_000)
    extended_nsec.update_checksum()
    extended_nsec_path = output / "extended-nsec-invalid.erofs"
    extended_nsec.save(extended_nsec_path)

    extended_seconds = extended_base.clone()
    extended_seconds_inode = extended_seconds.resolve_root_entry(
        "extended-time.txt"
    )
    extended_seconds.put_u64(extended_seconds_inode.offset + 32, INT64_MAX + 1)
    extended_seconds.update_checksum()
    extended_seconds_path = output / "extended-seconds-range.erofs"
    extended_seconds.save(extended_seconds_path)

    if compact_epoch_range.u64(SUPER + 24) + compact_epoch_range.u32(
        range_inode.offset + 12
    ) != INT64_MAX + 1:
        raise FixtureError("compact range fixture fields do not self-check")
    if compact_epoch_wrap.u64(SUPER + 24) + compact_epoch_wrap.u32(
        wrap_root.offset + 12
    ) <= UINT64_MAX:
        raise FixtureError("compact wrap fixture fields do not overflow")
    if compact_nsec.u32(SUPER + 32) != 1_000_000_000:
        raise FixtureError("compact nanoseconds fixture is incorrect")
    if extended_nsec.u32(extended_nsec_inode.offset + 40) != 1_000_000_000:
        raise FixtureError("extended nanoseconds fixture is incorrect")
    if extended_seconds.u64(extended_seconds_inode.offset + 32) != INT64_MAX + 1:
        raise FixtureError("extended seconds fixture is incorrect")

    evidence.extend(
        [
            "compact-epoch-range "
            f"nid={range_inode.nid} inode_offset={range_inode.offset} "
            f"epoch={INT64_MAX} delta=1 sum={INT64_MAX + 1}",
            "compact-epoch-wrap "
            f"root_nid={wrap_root.nid} inode_offset={wrap_root.offset} "
            f"epoch={UINT64_MAX} delta=1 overflow=yes",
            "compact-nsec "
            f"fixed_nsec={compact_nsec.u32(SUPER + 32)}",
            "extended-nsec "
            f"nid={extended_nsec_inode.nid} inode_offset={extended_nsec_inode.offset} "
            f"nsec={extended_nsec.u32(extended_nsec_inode.offset + 40)}",
            "extended-seconds "
            f"nid={extended_seconds_inode.nid} "
            f"inode_offset={extended_seconds_inode.offset} "
            f"seconds={extended_seconds.u64(extended_seconds_inode.offset + 32)}",
        ]
    )
    manifest["timestamps"] = {
        "compact_path": "/zz-identity.txt",
        "compact_nid": range_inode.nid,
        "extended_path": "/extended-time.txt",
        "extended_nid": extended_inode.nid,
        "time_t_max": INT64_MAX,
        "nsec_limit": 1_000_000_000,
    }

    for base in (nfs_base_path, extent_base_path, extended_base_path):
        base.unlink()
    shutil.rmtree(source)
    (output / "fixture-evidence.txt").write_text(
        "\n".join(evidence) + "\n", encoding="ascii"
    )
    (output / "fixture-manifest.json").write_text(
        json.dumps(manifest, 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(path)}  {path.name}\n")
    print((output / "fixture-evidence.txt").read_text(encoding="ascii"), end="")
    print((output / "SHA256SUMS").read_text(encoding="ascii"), end="")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command", required=True)
    make_parser = subparsers.add_parser("make")
    make_parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.command == "make":
        make_fixtures(args.output)


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