#!/usr/bin/env python3
from __future__ import annotations

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


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


def sha256_path(path: Path) -> str:
    return sha256_bytes(path.read_bytes())


def load_spec(path: Path) -> dict[str, object]:
    spec = json.loads(path.read_text(encoding="ascii"))
    if (
        spec.get("schema") != 1
        or spec.get("batch") != "B28"
        or spec.get("candidate") != "P15-086"
        or spec.get("test") != "TC176-stream-runtime"
    ):
        raise SystemExit("invalid B28 partial fixture spec")
    if set(spec.get("codecs", {})) != {"lz4", "lzma", "deflate", "zstd"}:
        raise SystemExit("B28 codec set changed")
    return spec


def make_sources(root: Path, spec: dict[str, object]) -> dict[str, bytes]:
    sources: dict[str, bytes] = {}
    lz4_spec = spec["sources"]["lz4"]
    sources["lz4"] = b"".join(
        bytes([segment["byte"]]) * segment["length"]
        for segment in lz4_spec["segments"]
    )
    stream_spec = spec["sources"]["stream"]
    sources["stream"] = b"".join(
        f"P15-086-{index % 64:02d}:alpha-beta-gamma-delta:{(index * 17) % 256:02x}\n".encode(
            "ascii"
        )
        for index in range(stream_spec["line_count"])
    )
    for name, content in sources.items():
        expected = spec["sources"][name]
        if len(content) != expected["size"] or sha256_bytes(content) != expected["sha256"]:
            raise SystemExit(f"B28 {name} source identity changed")
        directory = root / name
        directory.mkdir(parents=True)
        payload = directory / "payload.bin"
        payload.write_bytes(content)
        os.utime(payload, (0, 0))
        os.utime(directory, (0, 0))
    return sources


def run(argv: list[str], cwd: Path, expected: set[int] = {0}) -> subprocess.CompletedProcess[str]:
    completed = subprocess.run(
        argv,
        cwd=cwd,
        check=False,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        timeout=120,
    )
    if completed.returncode not in expected:
        raise SystemExit(
            f"command failed ({completed.returncode}): {' '.join(argv)}\n{completed.stdout}"
        )
    return completed


def verify_fsck(image: Path, expected: bytes, label: str) -> None:
    with tempfile.TemporaryDirectory(prefix=f"b28-{label}-fsck-") as temporary:
        destination = Path(temporary)
        run(["fsck.erofs", f"--extract={destination}", str(image)], image.parent)
        if (destination / "payload.bin").read_bytes() != expected:
            raise SystemExit(f"B28 {label} fsck extraction mismatch")


def build_codec(
    codec: str,
    codec_spec: dict[str, object],
    sources: dict[str, bytes],
    source_root: Path,
    images: Path,
) -> list[dict[str, object]]:
    valid = images / f"{codec}-valid.erofs"
    run(
        [
            "mkfs.erofs",
            *codec_spec["mkfs_args"],
            str(valid),
            str(source_root / codec_spec["source"]),
        ],
        images,
    )
    valid_data = valid.read_bytes()
    actual_image_sha256 = sha256_bytes(valid_data)
    source = sources[codec_spec["source"]]
    verify_fsck(valid, source, f"{codec}-valid")

    extent = codec_spec["extent"]
    start = extent["physical_offset"]
    end = start + extent["physical_length"]
    block = valid_data[start:end]
    leading = next((index for index, value in enumerate(block) if value), len(block))
    stream = block[leading:]
    if (
        leading != extent["leading_zero_bytes"]
        or len(stream) != extent["stream_bytes"]
        or not (extent["prefix_consumed"] < codec_spec["corruption_start"] < len(stream))
    ):
        raise SystemExit(f"{codec} pcluster or corruption boundary changed")

    corrupted_data = bytearray(valid_data)
    corruption = start + leading + codec_spec["corruption_start"]
    corrupted_data[corruption:end] = b"\0" * (end - corruption)
    if corrupted_data == valid_data:
        raise SystemExit(f"{codec} corruption mutation changed no bytes")
    corrupted = images / f"{codec}-corrupt.erofs"
    corrupted.write_bytes(corrupted_data)
    with tempfile.TemporaryDirectory(prefix=f"b28-{codec}-corrupt-") as temporary:
        failed = run(
            ["fsck.erofs", f"--extract={temporary}", str(corrupted)],
            images,
            set(range(1, 256)),
        )

    return [
        {
            "class": "valid",
            "codec": codec,
            "expected_errno": 0,
            "gate_image_sha256": codec_spec["gate_image_sha256"],
            "path": valid.name,
            "sha256": actual_image_sha256,
        },
        {
            "class": "corrupt",
            "codec": codec,
            "expected_full_errno": 97,
            "expected_prefix_errno": 0 if codec_spec["decision"] == "GO" else 97,
            "fsck_exit": failed.returncode,
            "path": corrupted.name,
            "sha256": sha256_path(corrupted),
        },
    ]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--spec", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    spec = load_spec(args.spec)
    if args.output.exists():
        raise SystemExit(f"refusing existing B28 output: {args.output}")
    for tool in ("mkfs.erofs", "fsck.erofs"):
        if shutil.which(tool) is None:
            raise SystemExit(f"{tool} is required")

    args.output.mkdir(parents=True)
    source_root = args.output / "sources"
    images = args.output / "images"
    images.mkdir()
    sources = make_sources(source_root, spec)
    records = []
    for codec in sorted(spec["codecs"]):
        records.extend(
            build_codec(codec, spec["codecs"][codec], sources, source_root, images)
        )
    (args.output / "spec.json").write_bytes(args.spec.read_bytes())
    index = {
        "batch": "B28",
        "candidate": "P15-086",
        "fixture_count": len(records),
        "fixtures": records,
        "schema": 1,
        "source_sha256": {
            name: sha256_bytes(content) for name, content in sorted(sources.items())
        },
        "status": "READY",
        "test": "TC176-stream-runtime",
    }
    (args.output / "fixture-index.json").write_text(
        json.dumps(index, indent=2, sort_keys=True) + "\n", encoding="ascii"
    )
    print(json.dumps(index, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()