#!/usr/bin/env python3
"""Generate and verify deterministic repo22 G7 stress fixtures."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path
import subprocess
from typing import Any


BLOCK_SIZE = 4096
MIB = 1024 * 1024
SPARSE_SIZE = 4 * 1024 * MIB + BLOCK_SIZE + 1
DEEP_LEVELS = 128
MANY_SMALL_COUNT = 12000
LARGE_DIRECTORY_COUNT = 12000
CONCURRENT_FILE_COUNT = 16
CONCURRENT_FILE_SIZE = 2 * MIB
PRESSURE_FILE_SIZE = 96 * MIB
SEQUENTIAL_FILE_SIZE = 256 * MIB
RANDOM_FILE_SIZE = 64 * MIB
FIXED_UUIDS = {
    "boundaries.erofs": "00000000-0000-0000-0000-000000000071",
    "workloads.erofs": "00000000-0000-0000-0000-000000000072",
}
SPARSE_MARKERS = (
    (0, b"repo22-g7-start"),
    (BLOCK_SIZE - 8, b"g7-block-edge"),
    (2**31 - 8, b"g7-2gib-edge"),
    (2**32 - 8, b"g7-4gib-edge"),
    (SPARSE_SIZE - 16, b"repo22-g7-end!!"),
)
SPARSE_RANGES = (
    ("start", 0, 64),
    ("block-edge", BLOCK_SIZE - 32, 64),
    ("one-gib-hole", 1024 * MIB, BLOCK_SIZE),
    ("two-gib-edge", 2**31 - 32, 64),
    ("four-gib-edge", 2**32 - 32, 64),
    ("eof-edge", SPARSE_SIZE - 64, 64),
)


class FixtureError(RuntimeError):
    pass


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        while chunk := source.read(MIB):
            digest.update(chunk)
    return digest.hexdigest()


def run(command: list[str], log: Path | None = None) -> str:
    result = subprocess.run(
        command,
        check=False,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )
    if log is not None:
        log.write_text(
            "$ " + " ".join(command) + "\n" + result.stdout,
            encoding="utf-8",
        )
    if result.returncode != 0:
        raise FixtureError(
            f"command failed ({result.returncode}): {' '.join(command)}\n"
            f"{result.stdout}"
        )
    return result.stdout


def repeated_bytes(label: str, size: int) -> bytes:
    token = (label + "\n").encode("ascii")
    return (token * ((size + len(token) - 1) // len(token)))[:size]


def write_pattern(
    path: Path,
    size: int,
    seed: str,
    random_period: int,
) -> None:
    remaining = size
    block_index = 0
    with path.open("wb") as output:
        while remaining:
            amount = min(BLOCK_SIZE, remaining)
            if block_index % random_period == 0:
                block = hashlib.shake_256(
                    f"repo22-g7:{seed}:{block_index}".encode("ascii")
                ).digest(amount)
            else:
                block = repeated_bytes(
                    f"repo22-g7:{seed}:{block_index % 97:02d}", amount
                )
            output.write(block)
            remaining -= amount
            block_index += 1


def create_sparse(path: Path) -> None:
    with path.open("wb") as output:
        output.truncate(SPARSE_SIZE)
        for offset, marker in SPARSE_MARKERS:
            output.seek(offset)
            output.write(marker)


def create_boundaries(root: Path) -> dict[str, Any]:
    root.mkdir(parents=True)
    (root / "empty.bin").write_bytes(b"")

    maximum = root / "maximum"
    maximum.mkdir()
    create_sparse(maximum / "sparse-boundary.bin")

    deep = root / "deep"
    deep.mkdir()
    current = deep
    components = []
    for level in range(DEEP_LEVELS):
        component = f"d{level:03d}"
        components.append(component)
        current /= component
        current.mkdir()
    deep_payload = current / "payload.bin"
    deep_payload.write_bytes(repeated_bytes("repo22-g7-deep", BLOCK_SIZE))
    deep_path = str(Path("deep", *components, "payload.bin"))

    long_directory = root / "longname"
    long_directory.mkdir()
    long_name = "n" * 255
    (long_directory / long_name).write_bytes(
        repeated_bytes("repo22-g7-long-name", BLOCK_SIZE)
    )

    many_small = root / "many-small"
    many_small.mkdir()
    for index in range(MANY_SMALL_COUNT):
        shard = many_small / f"shard-{index // 1000:02d}"
        shard.mkdir(exist_ok=True)
        (shard / f"file-{index:05d}.bin").write_bytes(
            repeated_bytes(f"repo22-g7-small:{index:05d}", 1024)
        )

    large_directory = root / "large-dir"
    large_directory.mkdir()
    for index in range(LARGE_DIRECTORY_COUNT):
        (large_directory / f"entry-{index:05d}.txt").write_bytes(
            f"repo22-g7-large-dir:{index:05d}\n".encode("ascii")
        )

    return {
        "deep_path": deep_path,
        "long_name": long_name,
    }


def create_workloads(root: Path) -> None:
    root.mkdir(parents=True)
    concurrent = root / "concurrent"
    concurrent.mkdir()
    for index in range(CONCURRENT_FILE_COUNT):
        write_pattern(
            concurrent / f"reader-{index:02d}.bin",
            CONCURRENT_FILE_SIZE,
            f"concurrent-{index:02d}",
            8,
        )

    pressure = root / "pressure"
    pressure.mkdir()
    write_pattern(
        pressure / "pressure.bin",
        PRESSURE_FILE_SIZE,
        "pressure",
        8,
    )

    sequential = root / "sequential"
    sequential.mkdir()
    write_pattern(
        sequential / "sequential.bin",
        SEQUENTIAL_FILE_SIZE,
        "sequential",
        4,
    )

    random_directory = root / "random"
    random_directory.mkdir()
    write_pattern(
        random_directory / "random.bin",
        RANDOM_FILE_SIZE,
        "random",
        1,
    )


def normalize_modes(root: Path) -> None:
    for path in sorted(root.rglob("*")):
        os.chmod(path, 0o755 if path.is_dir() else 0o644)
    os.chmod(root, 0o755)


def inventory(source_root: Path, output_root: Path) -> list[dict[str, Any]]:
    records = []
    for path in sorted(item for item in source_root.rglob("*") if item.is_file()):
        records.append(
            {
                "sha256": sha256(path),
                "size": path.stat().st_size,
                "path": str(path.relative_to(output_root)),
            }
        )
    return records


def inventory_text(records: list[dict[str, Any]]) -> str:
    lines = ["sha256\tsize\tpath"]
    lines.extend(
        f"{record['sha256']}\t{record['size']}\t{record['path']}"
        for record in records
    )
    return "\n".join(lines) + "\n"


def checksum_text(records: list[dict[str, Any]]) -> str:
    return "".join(
        f"{record['sha256']}  {record['path']}\n" for record in records
    )


def create_images(output_root: Path) -> list[dict[str, Any]]:
    images = output_root / "images"
    logs = output_root / "logs"
    images.mkdir()
    logs.mkdir()
    records = []
    for image_name, source_name in (
        ("boundaries.erofs", "boundaries"),
        ("workloads.erofs", "workloads"),
    ):
        image = images / image_name
        command = [
            "mkfs.erofs",
            "--workers=1",
            "--sort=path",
            "--all-root",
            "-T0",
            "--all-time",
            f"-U{FIXED_UUIDS[image_name]}",
            "-z",
            "lz4",
            str(image),
            str(output_root / "sources" / source_name),
        ]
        run(command, logs / f"mkfs-{source_name}.log")
        run(
            ["fsck.erofs", "-d0", str(image)],
            logs / f"fsck-{source_name}.log",
        )
        records.append(
            {
                "sha256": sha256(image),
                "size": image.stat().st_size,
                "path": str(image.relative_to(output_root)),
                "command": command,
            }
        )
    return records


def write_metadata(
    output_root: Path,
    source_records: list[dict[str, Any]],
    image_records: list[dict[str, Any]],
    boundary_metadata: dict[str, Any],
) -> None:
    source_inventory = inventory_text(source_records)
    image_checksums = checksum_text(image_records)
    (output_root / "SOURCE-INVENTORY.tsv").write_text(
        source_inventory, encoding="utf-8"
    )
    (output_root / "SOURCE-SHA256SUMS").write_text(
        checksum_text(source_records), encoding="utf-8"
    )
    (output_root / "SHA256SUMS").write_text(
        image_checksums, encoding="utf-8"
    )
    (output_root / "DEEP-PATH").write_text(
        boundary_metadata["deep_path"] + "\n", encoding="ascii"
    )
    (output_root / "LONG-NAME").write_text(
        boundary_metadata["long_name"] + "\n", encoding="ascii"
    )
    (output_root / "SPARSE-RANGES.tsv").write_text(
        "".join(
            f"{label}\t{offset}\t{length}\n"
            for label, offset, length in SPARSE_RANGES
        ),
        encoding="ascii",
    )
    manifest = {
        "format": 1,
        "block_size": BLOCK_SIZE,
        "deep_levels": DEEP_LEVELS,
        "deep_path": boundary_metadata["deep_path"],
        "long_name_bytes": len(boundary_metadata["long_name"].encode("ascii")),
        "many_small_count": MANY_SMALL_COUNT,
        "large_directory_count": LARGE_DIRECTORY_COUNT,
        "concurrent_file_count": CONCURRENT_FILE_COUNT,
        "concurrent_file_size": CONCURRENT_FILE_SIZE,
        "pressure_file_size": PRESSURE_FILE_SIZE,
        "sequential_file_size": SEQUENTIAL_FILE_SIZE,
        "random_file_size": RANDOM_FILE_SIZE,
        "sparse_size": SPARSE_SIZE,
        "sparse_markers": [
            {"offset": offset, "hex": marker.hex()}
            for offset, marker in SPARSE_MARKERS
        ],
        "sparse_ranges": [
            {"label": label, "offset": offset, "length": length}
            for label, offset, length in SPARSE_RANGES
        ],
        "source_count": len(source_records),
        "source_inventory_sha256": hashlib.sha256(
            source_inventory.encode("utf-8")
        ).hexdigest(),
        "image_count": len(image_records),
        "images": image_records,
    }
    (output_root / "fixture-manifest.json").write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )


def create(output_root: Path) -> None:
    if output_root.exists():
        raise FixtureError(f"output already exists: {output_root}")
    sources = output_root / "sources"
    sources.mkdir(parents=True)
    boundary_metadata = create_boundaries(sources / "boundaries")
    create_workloads(sources / "workloads")
    normalize_modes(sources)
    source_records = inventory(sources, output_root)
    image_records = create_images(output_root)
    write_metadata(
        output_root,
        source_records,
        image_records,
        boundary_metadata,
    )
    verify(output_root)


def require(condition: bool, message: str) -> None:
    if not condition:
        raise FixtureError(message)


def verify(output_root: Path) -> None:
    manifest_path = output_root / "fixture-manifest.json"
    require(manifest_path.is_file(), "missing fixture-manifest.json")
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    sources = output_root / "sources"
    source_records = inventory(sources, output_root)
    source_text = inventory_text(source_records)
    require(
        source_text == (output_root / "SOURCE-INVENTORY.tsv").read_text(
            encoding="utf-8"
        ),
        "source inventory mismatch",
    )
    require(
        checksum_text(source_records)
        == (output_root / "SOURCE-SHA256SUMS").read_text(encoding="utf-8"),
        "source checksum list mismatch",
    )
    require(len(source_records) == manifest["source_count"], "source count mismatch")
    require(
        hashlib.sha256(source_text.encode("utf-8")).hexdigest()
        == manifest["source_inventory_sha256"],
        "source inventory hash mismatch",
    )

    sparse = sources / "boundaries" / "maximum" / "sparse-boundary.bin"
    require(sparse.stat().st_size == SPARSE_SIZE, "sparse file size mismatch")
    require(
        sparse.stat().st_blocks * 512 < MIB,
        "sparse source unexpectedly consumes at least 1 MiB",
    )
    with sparse.open("rb") as source:
        for offset, marker in SPARSE_MARKERS:
            source.seek(offset)
            require(source.read(len(marker)) == marker, f"marker mismatch at {offset}")

    deep_path = (output_root / "DEEP-PATH").read_text(encoding="ascii").strip()
    require(deep_path == manifest["deep_path"], "deep path metadata mismatch")
    require((sources / "boundaries" / deep_path).is_file(), "deep payload missing")
    require(
        len(Path(deep_path).parts) - 2 == DEEP_LEVELS,
        "deep directory level mismatch",
    )
    long_name = (output_root / "LONG-NAME").read_text(encoding="ascii").strip()
    require(len(long_name.encode("ascii")) == 255, "long name is not 255 bytes")
    require(
        (sources / "boundaries" / "longname" / long_name).is_file(),
        "long-name source missing",
    )
    require(
        sum(1 for path in (sources / "boundaries" / "many-small").rglob("*") if path.is_file())
        == MANY_SMALL_COUNT,
        "many-small count mismatch",
    )
    require(
        sum(1 for path in (sources / "boundaries" / "large-dir").iterdir() if path.is_file())
        == LARGE_DIRECTORY_COUNT,
        "large-directory count mismatch",
    )
    require(
        sum(1 for path in (sources / "workloads" / "concurrent").iterdir() if path.is_file())
        == CONCURRENT_FILE_COUNT,
        "concurrent source count mismatch",
    )

    image_records = []
    logs = output_root / "logs"
    for image in sorted((output_root / "images").glob("*.erofs")):
        run(["fsck.erofs", "-d0", str(image)], logs / f"verify-{image.stem}.log")
        image_records.append(
            {
                "sha256": sha256(image),
                "size": image.stat().st_size,
                "path": str(image.relative_to(output_root)),
            }
        )
    require(len(image_records) == manifest["image_count"], "image count mismatch")
    expected_images = [
        {key: record[key] for key in ("sha256", "size", "path")}
        for record in manifest["images"]
    ]
    require(image_records == expected_images, "image inventory mismatch")
    require(
        checksum_text(image_records)
        == (output_root / "SHA256SUMS").read_text(encoding="utf-8"),
        "image checksum list mismatch",
    )
    sparse_dump = run(
        [
            "dump.erofs",
            "--path=/maximum/sparse-boundary.bin",
            str(output_root / "images" / "boundaries.erofs"),
        ]
    )
    require(f"Size: {SPARSE_SIZE} " in sparse_dump, "image sparse size mismatch")
    print(
        "verified "
        f"sources={len(source_records)} images={len(image_records)} "
        f"inventory_sha256={manifest['source_inventory_sha256']}"
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)
    for command in ("create", "verify"):
        subparser = subparsers.add_parser(command)
        subparser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    try:
        if args.command == "create":
            create(args.output.resolve())
        else:
            verify(args.output.resolve())
    except FixtureError as error:
        print(f"error: {error}")
        return 1
    return 0


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