#!/usr/bin/env python3
"""Build assertion-driven structured variants for the G1 manual tests."""

from __future__ import annotations

import argparse
import hashlib
import shutil
import struct
from pathlib import Path

from erofs_fixture import ErofsImage, FEATURE_INCOMPAT_48BIT, SUPER


S_IFMT = 0o170000
S_IFDIR = 0o040000
S_IFREG = 0o100000
S_IFLNK = 0o120000
S_IFCHR = 0o020000
S_IFBLK = 0o060000
S_IFIFO = 0o010000
FLAT_PLAIN = 0
FLAT_INLINE = 2
I_NLINK_1 = 0x10
NULL_ADDR_48 = (1 << 48) - 1
LARGE_FILE_SIZE = (1 << 32) + 4097


def inode_for(image: ErofsImage, path: str):
    _, entry = image.resolve_root_entry(path)
    return image.inode(entry.nid)


def assert_inode(
    image: ErofsImage,
    path: str,
    *,
    inode_size: int | None = None,
    layout: int | None = None,
    file_type: int | None = None,
):
    inode = inode_for(image, path)
    if inode_size is not None and inode.inode_size != inode_size:
        raise ValueError(f"{path}: inode size {inode.inode_size}, expected {inode_size}")
    if layout is not None and inode.layout != layout:
        raise ValueError(f"{path}: layout {inode.layout}, expected {layout}")
    if file_type is not None and inode.mode & S_IFMT != file_type:
        raise ValueError(f"{path}: mode {inode.mode:#o}, expected type {file_type:#o}")
    return inode


def save_checked(image: ErofsImage, path: Path) -> None:
    image.update_checksum()
    image.validate_superblock()
    image.save(path)


def select_48bit_root(image: ErofsImage) -> int:
    root_nid = image.root_nid
    if root_nid == 0 or root_nid >= 1 << 48:
        raise ValueError(f"root NID {root_nid} is unsuitable for the 48-bit fixture")
    image.put_u32(SUPER + 80, image.feature_incompat | FEATURE_INCOMPAT_48BIT)
    image.put_u16(SUPER + 14, 0)
    image.put_u64(SUPER + 112, root_nid)
    return root_nid


def rewrite_without_dot(image: ErofsImage, path: str) -> tuple[int, int, int]:
    inode = assert_inode(
        image, path, inode_size=32, layout=FLAT_INLINE, file_type=S_IFDIR
    )
    entries = image.directory_entries(inode)
    names = [entry.name for entry in entries]
    if names.count(b".") != 1 or b".." not in names:
        raise ValueError(f"{path}: expected one dot and an explicit dotdot entry")
    kept = [entry for entry in entries if entry.name != b"."]
    data_offset = inode.offset + inode.inode_size + inode.xattr_size
    data = bytearray(len(kept) * 12)
    names_blob = bytearray()
    for index, entry in enumerate(kept):
        name_offset = len(data) + len(names_blob)
        file_type = image.data[entry.offset + 10]
        struct.pack_into("<QHBB", data, index * 12, entry.nid, name_offset, file_type, 0)
        names_blob.extend(entry.name)
    rebuilt = data + names_blob
    if len(rebuilt) >= inode.size:
        raise ValueError(f"{path}: rebuilt directory did not shrink")
    image.data[data_offset : data_offset + inode.size] = bytes(inode.size)
    image.data[data_offset : data_offset + len(rebuilt)] = rebuilt
    image.put_u32(inode.offset + 8, len(rebuilt))
    image.put_u16(inode.offset, inode.inode_format | I_NLINK_1)
    updated = image.inode(inode.nid)
    updated_names = [entry.name for entry in image.directory_entries(updated)]
    if b"." in updated_names or b".." not in updated_names:
        raise AssertionError(f"{path}: dot omission rewrite failed")
    return inode.nid, inode.size, len(rebuilt)


def image_hash(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def write_hashes(output: Path) -> None:
    with (output / "IMAGE-SHA256SUMS").open("w", encoding="ascii") as sums:
        for path in sorted(output.glob("*.erofs")):
            sums.write(f"{image_hash(path)}  {path.name}\n")


def copy_validated(source: Path, destination: Path) -> ErofsImage:
    image = ErofsImage.load(source)
    image.validate_superblock()
    shutil.copy2(source, destination)
    return image


def make(args: argparse.Namespace) -> None:
    output = args.output
    if output.exists():
        raise ValueError(f"output already exists: {output}")
    output.mkdir(parents=True)

    compact = copy_validated(args.compact, output / "compact.erofs")
    extended = copy_validated(args.extended, output / "extended.erofs")
    flat = copy_validated(args.flat, output / "flat.erofs")
    inline = copy_validated(args.inline, output / "inline.erofs")
    evidence = []

    compact_root = compact.inode(compact.root_nid)
    if compact_root.inode_size != 32:
        raise ValueError("compact image root is not a compact inode")
    explicit_dir = assert_inode(
        compact, "/dotdir", inode_size=32, layout=FLAT_INLINE, file_type=S_IFDIR
    )
    explicit_names = [entry.name for entry in compact.directory_entries(explicit_dir)]
    if b"." not in explicit_names or b".." not in explicit_names:
        raise ValueError("compact /dotdir does not contain explicit dot entries")
    if explicit_dir.inode_format & I_NLINK_1:
        raise ValueError("compact /dotdir unexpectedly advertises dot omission")
    evidence.append(
        f"compact root_nid={compact.root_nid} inode_size=32 "
        f"blocks={compact.blocks} inos={compact.u64(SUPER + 16)} "
        f"dotdir_nid={explicit_dir.nid} dotdir_i_format={explicit_dir.inode_format:#06x} "
        f"dotdir_names={','.join(name.decode('ascii') for name in explicit_names)}"
    )

    small = assert_inode(compact, "/small.txt", inode_size=32, file_type=S_IFREG)
    single = assert_inode(compact, "/single.txt", inode_size=32, file_type=S_IFREG)
    hard_a = assert_inode(compact, "/hard-a.txt", inode_size=32, file_type=S_IFREG)
    hard_b = assert_inode(compact, "/hard-b.txt", inode_size=32, file_type=S_IFREG)
    if hard_a.nid != hard_b.nid or compact.u16(hard_a.offset + 6) != 2:
        raise ValueError("compact hard-link fixture does not encode nlink=2")
    if compact.u16(single.offset + 6) != 1 or single.inode_format & I_NLINK_1:
        raise ValueError("compact single-link baseline is not explicit nlink=1")
    evidence.append(
        f"compact small_nid={small.nid} single_nid={single.nid} "
        f"single_i_format={single.inode_format:#06x} single_i_nb=1 "
        f"hardlink_nid={hard_a.nid} hardlink_i_nb=2"
    )

    expected_rdev = 0x543ABC21
    special_fields = []
    for path, expected_type in (
        ("/char-large", S_IFCHR),
        ("/block-large", S_IFBLK),
        ("/fifo", S_IFIFO),
    ):
        inode = assert_inode(
            compact, path, inode_size=32, file_type=expected_type
        )
        raw = compact.u32(inode.offset + 16)
        if expected_type in (S_IFCHR, S_IFBLK) and raw != expected_rdev:
            raise ValueError(f"{path}: raw rdev {raw:#x}, expected {expected_rdev:#x}")
        special_fields.append(f"{path[1:]}:nid={inode.nid}:raw={raw:#010x}")
    evidence.append("compact special=" + ",".join(special_fields))

    bad_dirent = compact.clone()
    _, invalid_entry = bad_dirent.resolve_root_entry("/invalid-target.txt")
    invalid_nid = (bad_dirent.blocks << bad_dirent.block_bits) // 32 + 1024
    bad_dirent.put_u64(invalid_entry.offset, invalid_nid)
    save_checked(bad_dirent, output / "invalid-dirent-nid.erofs")
    evidence.append(
        f"invalid-dirent entry_offset={invalid_entry.offset} "
        f"old_nid={invalid_entry.nid} new_nid={invalid_nid} crc=valid"
    )

    nlink1 = compact.clone()
    single = assert_inode(nlink1, "/single.txt", inode_size=32, file_type=S_IFREG)
    nlink1.put_u16(single.offset, single.inode_format | I_NLINK_1)
    nlink1.put_u16(single.offset + 6, 0x1234)
    save_checked(nlink1, output / "compact-nlink1.erofs")
    evidence.append(
        f"nlink1 nid={single.nid} i_format={single.inode_format | I_NLINK_1:#06x} "
        "i_nb=0x1234 crc=valid"
    )

    dot_omitted = compact.clone()
    dot_nid, old_size, new_size = rewrite_without_dot(dot_omitted, "/dotdir")
    dot_root = select_48bit_root(dot_omitted)
    save_checked(dot_omitted, output / "compact-dot-omitted.erofs")
    evidence.append(
        f"dot-omitted nid={dot_nid} i_format_bit4=1 old_size={old_size} "
        f"new_size={new_size} ondisk_names=..,child.txt "
        f"feature=0x80 rootnid_8b={dot_root} blocks_hi=0 crc=valid"
    )

    root8 = compact.clone()
    root_nid = select_48bit_root(root8)
    save_checked(root8, output / "root8-48bit.erofs")
    if root8.root_nid != root_nid or root8.blocks != compact.blocks:
        raise AssertionError("48-bit root selector self-check failed")
    evidence.append(
        f"root8 feature={root8.feature_incompat:#x} rootnid_8b={root_nid} "
        f"blocks_hi={root8.u16(SUPER + 14)} blocks_lo={root8.u32(SUPER + 36)}"
    )

    fallback = compact.clone()
    fallback_root = fallback.root_nid
    fallback.put_u32(
        SUPER + 80, fallback.feature_incompat | FEATURE_INCOMPAT_48BIT
    )
    fallback.put_u64(SUPER + 112, 0)
    save_checked(fallback, output / "fallback-48bit-root2.erofs")
    if fallback.root_nid != fallback_root or fallback.blocks != compact.blocks:
        raise AssertionError("48-bit fallback selector self-check failed")
    evidence.append(
        f"fallback feature={fallback.feature_incompat:#x} "
        f"rootnid_2b={fallback_root} rootnid_8b=0 blocks_lo={fallback.blocks}"
    )

    normal = assert_inode(
        extended, "/large-file.bin", inode_size=64, file_type=S_IFREG
    )
    evidence.append(
        f"extended large-file_nid={normal.nid} inode_off={normal.offset} "
        f"size={normal.size} inode_size=64"
    )

    large_hole = extended.clone()
    huge = assert_inode(
        large_hole, "/huge-sparse.dat", inode_size=64, file_type=S_IFREG
    )
    large_hole.put_u16(huge.offset, (huge.inode_format & ~(7 << 1)) | (FLAT_PLAIN << 1))
    large_hole.put_u16(huge.offset + 6, (NULL_ADDR_48 >> 32) & 0xFFFF)
    large_hole.put_u32(huge.offset + 16, NULL_ADDR_48 & 0xFFFFFFFF)
    large_hole.put_u64(huge.offset + 8, LARGE_FILE_SIZE)
    save_checked(large_hole, output / "extended-large-hole.erofs")
    evidence.append(
        f"extended-large-hole nid={huge.nid} inode_off={huge.offset} "
        f"size={LARGE_FILE_SIZE} layout=0 startblk=0xffffffffffff crc=valid"
    )

    oversize = extended.clone()
    oversize_inode = assert_inode(
        oversize, "/oversize.dat", inode_size=64, file_type=S_IFREG
    )
    oversize.put_u64(oversize_inode.offset + 8, 1 << 63)
    save_checked(oversize, output / "extended-size-bit63.erofs")
    evidence.append(
        f"extended-size-bit63 nid={oversize_inode.nid} "
        f"inode_off={oversize_inode.offset} i_size=0x8000000000000000 crc=valid"
    )

    for path in (
        "/small.txt",
        "/medium.dat",
        "/large.bin",
        "/data.txt",
        "/random-test.bin",
        "/huge-data.bin",
    ):
        inode = assert_inode(flat, path, layout=FLAT_PLAIN, file_type=S_IFREG)
        evidence.append(
            f"flat path={path} nid={inode.nid} layout={inode.layout} size={inode.size}"
        )
    longlink = assert_inode(
        flat, "/longlink", layout=FLAT_INLINE, file_type=S_IFLNK
    )
    if longlink.size <= 60:
        raise ValueError("longlink target is not longer than 60 bytes")
    evidence.append(
        f"flat path=/longlink nid={longlink.nid} layout={longlink.layout} "
        f"size={longlink.size}"
    )

    inline_zero = inline.clone()
    empty = assert_inode(inline_zero, "/empty.txt", inode_size=32, file_type=S_IFREG)
    inline_zero.put_u16(
        empty.offset, (empty.inode_format & ~(7 << 1)) | (FLAT_INLINE << 1)
    )
    save_checked(inline_zero, output / "inline-zero.erofs")
    empty_after = assert_inode(
        inline_zero,
        "/empty.txt",
        inode_size=32,
        layout=FLAT_INLINE,
        file_type=S_IFREG,
    )
    if empty_after.size != 0:
        raise AssertionError("inline zero fixture size changed")
    evidence.append(
        f"inline-zero nid={empty_after.nid} i_format={empty_after.inode_format:#06x} size=0"
    )

    inline_expectations = {
        "/tiny.txt": FLAT_INLINE,
        "/inline.txt": FLAT_INLINE,
        "/tailpacked.dat": FLAT_INLINE,
        "/file-4095.dat": FLAT_PLAIN,
        "/file-4096.dat": FLAT_PLAIN,
        "/file-4097.dat": FLAT_INLINE,
        "/link1": FLAT_INLINE,
        "/brokenlink": FLAT_INLINE,
    }
    for path, layout in inline_expectations.items():
        expected_type = S_IFLNK if path in ("/link1", "/brokenlink") else S_IFREG
        inode = assert_inode(inline, path, layout=layout, file_type=expected_type)
        evidence.append(
            f"inline path={path} nid={inode.nid} layout={inode.layout} size={inode.size}"
        )

    inline_cross = inline.clone()
    cross = assert_inode(
        inline_cross,
        "/inline.txt",
        inode_size=32,
        layout=FLAT_INLINE,
        file_type=S_IFREG,
    )
    if inline_cross.u16(cross.offset + 2) != 0:
        raise ValueError("inline.txt unexpectedly has inline xattrs")
    chosen_count = None
    chosen_offset = None
    for count in range(1, 65536):
        xattr_size = 12 + 4 * (count - 1)
        data_offset = cross.offset + cross.inode_size + xattr_size
        block_offset = data_offset % inline_cross.block_size
        if block_offset + cross.size > inline_cross.block_size:
            chosen_count = count
            chosen_offset = data_offset
            break
    if chosen_count is None or chosen_offset is None:
        raise AssertionError("could not construct cross-block inline range")
    inline_cross.put_u16(cross.offset + 2, chosen_count)
    save_checked(inline_cross, output / "inline-cross-block.erofs")
    evidence.append(
        f"inline-cross nid={cross.nid} inode_off={cross.offset} "
        f"xattr_icount=0->{chosen_count} data_blockoff={chosen_offset % inline_cross.block_size} "
        f"size={cross.size} crc=valid"
    )

    (output / "fixture-evidence.txt").write_text(
        "\n".join(evidence) + "\n", encoding="ascii"
    )
    write_hashes(output)


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


def main() -> None:
    make(parse_args())


if __name__ == "__main__":
    main()