#!/usr/bin/env python3
"""Build and self-check deterministic fixtures for the G3 manual set."""

from __future__ import annotations

import argparse
import hashlib
import os
from pathlib import Path
import shutil
import stat
import subprocess

from erofs_fixture import ErofsImage, SUPER


FLAT_PLAIN = 0
HUGE_DIRECTORY_BLOCKS = (1 << 31) + 1
S_IFMT = 0o170000


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


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


def deterministic_bytes(length: int) -> bytes:
    return bytes(
        ((index * 131 + 17) ^ (index >> 3)) & 0xFF
        for index in range(length)
    )


def make_source(source: Path) -> None:
    testdir = source / "testdir"
    (testdir / "subdir").mkdir(parents=True)
    for index in range(5):
        (testdir / f"dir-{index:02d}").mkdir()
    (source / "parent" / "child" / "grandchild").mkdir(parents=True)

    (testdir / "file.txt").write_text("repo22 G3 file payload\n", encoding="ascii")
    for index in range(12):
        (testdir / f"regular-{index:02d}.txt").write_text(
            f"regular {index:02d}\n", encoding="ascii"
        )
    for name in ("File.txt", "FILE.TXT", "FiLe.TxT"):
        (testdir / name).write_text(f"case variant {name}\n", encoding="ascii")
    for index in range(1, 4):
        (testdir / f"file{index}.txt").write_text(
            f"cache file {index}\n", encoding="ascii"
        )
    (testdir / "subdir" / "child.txt").write_text("child\n", encoding="ascii")
    (testdir / "script.sh").write_text("#!/bin/sh\nexit 0\n", encoding="ascii")
    (testdir / "rootonly.txt").write_text("restricted\n", encoding="ascii")
    (testdir / "writeonly.txt").write_text("write only mode\n", encoding="ascii")
    (testdir / "noexec.txt").write_text("not executable\n", encoding="ascii")
    long_name = "long-" + "x" * 250
    (testdir / long_name).write_text("255 byte name\n", encoding="ascii")
    long_target = "subdir/" + "y" * 112
    (testdir / "subdir" / ("y" * 112)).write_text(
        "long target\n", encoding="ascii"
    )

    os.symlink("file.txt", testdir / "shortlink")
    os.symlink("subdir/child.txt", testdir / "relative-link")
    os.symlink("missing-target", testdir / "broken-link")
    os.symlink(long_target, testdir / "long-link")

    (source / "parent" / "file.txt").write_text("parent file\n", encoding="ascii")
    (source / "parent" / "child" / "grandchild" / "marker.txt").write_text(
        "grandchild\n", encoding="ascii"
    )
    (source / "pager.bin").write_bytes(deterministic_bytes(5 * 4096 + 731))

    os.chmod(testdir / "script.sh", 0o755)
    os.chmod(testdir / "rootonly.txt", 0o600)
    os.chmod(testdir / "writeonly.txt", 0o200)
    os.chmod(testdir / "noexec.txt", 0o644)
    set_epoch(source)


def directory_data_offset(image: ErofsImage, inode) -> int:
    if inode.layout == 2:
        return inode.offset + inode.inode_size + inode.xattr_size
    if inode.layout == FLAT_PLAIN:
        return inode.start_block << image.block_bits
    raise ValueError(f"unsupported directory layout {inode.layout}")


def directory_entries(image: ErofsImage, inode):
    if inode.size == 0 or inode.size > image.block_size:
        raise ValueError("G3 qualifier accepts only one-block path components")
    return image.directory_entries(inode)


def resolve(image: ErofsImage, path: str):
    inode = image.inode(image.root_nid)
    for component in path.strip("/").split("/"):
        if not component:
            continue
        entry = next(
            (
                item
                for item in directory_entries(image, inode)
                if item.name.decode("ascii") == component
            ),
            None,
        )
        if entry is None:
            raise ValueError(f"path not found: {path}")
        inode = image.inode(entry.nid)
    return inode


def validate_image(image_path: Path, source: Path) -> list[str]:
    image = ErofsImage.load(image_path)
    image.validate_superblock()
    evidence = [
        f"image={image_path.name} provider_bytes={len(image.data)} "
        f"block_size={image.block_size} blocks={image.blocks} "
        f"root_nid={image.root_nid} checksum_valid={image.checksum_valid()}"
    ]
    paths = [
        "/testdir",
        "/testdir/file.txt",
        "/testdir/script.sh",
        "/testdir/rootonly.txt",
        "/testdir/writeonly.txt",
        "/testdir/shortlink",
        "/testdir/long-link",
        "/pager.bin",
        "/parent/child/grandchild",
    ]
    for path in paths:
        source_path = source / path.lstrip("/")
        source_stat = source_path.lstat()
        inode = resolve(image, path)
        if inode.mode & S_IFMT != stat.S_IFMT(source_stat.st_mode):
            raise ValueError(f"{path}: inode type differs from source")
        if inode.mode & 0o7777 != source_stat.st_mode & 0o7777:
            raise ValueError(f"{path}: inode permissions differ from source")
        expected_size = None
        if source_path.is_symlink():
            expected_size = len(os.readlink(source_path).encode("ascii"))
        elif source_path.is_file():
            expected_size = source_stat.st_size
        if expected_size is not None and inode.size != expected_size:
            raise ValueError(f"{path}: size {inode.size}, expected {expected_size}")
        evidence.append(
            f"path={path} nid={inode.nid} inode_offset={inode.offset} "
            f"layout={inode.layout} mode={inode.mode:#07o} size={inode.size} "
            f"start_block={inode.start_block}"
        )

    source_names = sorted(item.name for item in (source / "testdir").iterdir())
    testdir = resolve(image, "/testdir")
    image_names = sorted(
        entry.name.decode("ascii")
        for entry in directory_entries(image, testdir)
        if entry.name not in (b".", b"..")
    )
    if source_names != image_names:
        raise ValueError("/testdir names differ from source")
    evidence.append(
        f"path=/testdir real_entries={len(image_names)} "
        f"directory_size={testdir.size} layout={testdir.layout}"
    )
    return evidence


def write_source_hashes(source: Path, output: Path) -> None:
    with output.open("w", encoding="ascii") as sums:
        for path in sorted(source.rglob("*")):
            if path.is_file() and not path.is_symlink():
                sums.write(f"{sha256(path)}  {path.relative_to(source)}\n")


def make_vfs(output: Path) -> None:
    if output.exists():
        raise ValueError(f"output already exists: {output}")
    source = output / "source"
    output.mkdir(parents=True)
    make_source(source)
    plain = output / "vfs-plain.erofs"
    lz4 = output / "vfs-lz4.erofs"
    common = [
        "mkfs.erofs",
        "-d0",
        "-T0",
        "--all-time",
        "--all-root",
        "--workers=1",
    ]
    subprocess.run(
        common
        + [
            "-x-1",
            "-E",
            "noinline_data",
            "-U",
            "33333333-4444-5555-6666-777777777771",
            str(plain),
            str(source),
        ],
        check=True,
    )
    subprocess.run(
        common
        + [
            "-z",
            "lz4",
            "-U",
            "33333333-4444-5555-6666-777777777772",
            str(lz4),
            str(source),
        ],
        check=True,
    )
    evidence = validate_image(plain, source) + validate_image(lz4, source)
    (output / "fixture-evidence.txt").write_text(
        "\n".join(evidence) + "\n", encoding="ascii"
    )
    (output / "expected-testdir.txt").write_text(
        "\n".join(sorted(item.name for item in (source / "testdir").iterdir()))
        + "\n",
        encoding="ascii",
    )
    write_source_hashes(source, output / "SOURCE-SHA256SUMS")
    with (output / "IMAGE-SHA256SUMS").open("w", encoding="ascii") as sums:
        for path in (lz4, plain):
            sums.write(f"{sha256(path)}  {path.name}\n")
    print((output / "fixture-evidence.txt").read_text(encoding="ascii"), end="")
    print((output / "IMAGE-SHA256SUMS").read_text(encoding="ascii"), end="")


def make_large_prefix(source: Path, output: Path, evidence_path: Path) -> None:
    image = ErofsImage.load(source)
    image.validate_superblock()
    huge = resolve(image, "/huge")
    expected_size = HUGE_DIRECTORY_BLOCKS * image.block_size
    if huge.inode_size != 64 or huge.layout != FLAT_PLAIN:
        raise ValueError("TC153 directory is not extended FLAT_PLAIN/Layout0")
    if huge.size != expected_size:
        raise ValueError(f"TC153 directory size {huge.size}, expected {expected_size}")
    if huge.start_block == 0:
        raise ValueError("TC153 directory has no plain-data start block")
    provider_blocks = huge.start_block + HUGE_DIRECTORY_BLOCKS
    if provider_blocks > 0xFFFFFFFF:
        raise ValueError("TC153 sparse provider does not fit blocks_lo")
    image.put_u32(SUPER + 36, provider_blocks)
    image.update_checksum()
    image.validate_superblock()
    output.parent.mkdir(parents=True, exist_ok=True)
    image.save(output)

    checked = ErofsImage.load(output)
    checked_huge = resolve(checked, "/huge")
    if (
        checked.blocks != provider_blocks
        or checked_huge.layout != FLAT_PLAIN
        or checked_huge.size != expected_size
    ):
        raise AssertionError("TC153 sparse prefix self-check failed")
    provider_bytes = provider_blocks * checked.block_size
    evidence_path.write_text(
        "tc153-sparse-prefix "
        f"nid={checked_huge.nid} inode_offset={checked_huge.offset} "
        f"layout=flat-plain start_block={checked_huge.start_block} "
        f"directory_blocks={HUGE_DIRECTORY_BLOCKS} "
        f"last_block={HUGE_DIRECTORY_BLOCKS - 1} "
        f"blocks_lo={provider_blocks} provider_bytes={provider_bytes} "
        f"prefix_bytes={len(checked.data)} checksum_valid={checked.checksum_valid()} "
        f"sha256={sha256(output)}\n",
        encoding="ascii",
    )
    print(evidence_path.read_text(encoding="ascii"), end="")


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


def main() -> None:
    args = parse_args()
    if args.command == "make-vfs":
        make_vfs(args.output)
    else:
        make_large_prefix(args.source, args.output, args.evidence)


if __name__ == "__main__":
    main()