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

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


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


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


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


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 = json.loads(args.spec.read_text(encoding="ascii"))
    if (
        spec.get("schema") != 1
        or spec.get("batch") != "B32"
        or spec.get("test") != "TC168-cache-inflight"
    ):
        raise SystemExit("invalid B32 cache fixture spec")
    if args.output.exists():
        raise SystemExit(f"refusing existing B32 output: {args.output}")
    if shutil.which("mkfs.erofs") is None:
        raise SystemExit("mkfs.erofs is required")

    source = (b"Pre15-B32-cache-state:0123456789abcdef\n" * 2000)[
        : spec["source"]["size"]
    ]
    if sha256_bytes(source) != spec["source"]["sha256"]:
        raise SystemExit("B32 deterministic source identity changed")

    source_dir = args.output / "source"
    images = args.output / "images"
    source_dir.mkdir(parents=True)
    images.mkdir()
    payload = source_dir / "payload.bin"
    payload.write_bytes(source)
    os.utime(payload, (0, 0))
    os.utime(source_dir, (0, 0))

    valid = images / "lzma-valid.erofs"
    run(["mkfs.erofs", *spec["mkfs_args"], str(valid), str(source_dir)], images)
    valid_data = valid.read_bytes()
    if len(valid_data) != 8192:
        raise SystemExit(f"unexpected B32 image size: {len(valid_data)}")
    pcluster_start = 4096
    pcluster_end = 8192
    pcluster = valid_data[pcluster_start:pcluster_end]
    leading = next(
        (index for index, value in enumerate(pcluster) if value), len(pcluster)
    )
    stream = pcluster[leading:]
    if leading == len(pcluster) or len(stream) < 2:
        raise SystemExit("B32 LZMA stream boundary is empty")

    truncated_data = bytearray(valid_data)
    truncated_data[pcluster_start + leading : pcluster_end] = b"\0" + stream[:-1]
    if truncated_data == valid_data:
        raise SystemExit("B32 truncated mutation changed no bytes")
    truncated = images / "lzma-truncated.erofs"
    truncated.write_bytes(truncated_data)

    records = []
    for image_class, path, expected_errno in (
        ("valid", valid, 0),
        ("truncated", truncated, 97),
    ):
        records.append(
            {
                "class": image_class,
                "expected_errno": expected_errno,
                "path": path.name,
                "sha256": sha256_path(path),
            }
        )
    index = {
        "batch": "B32",
        "decoded_size": len(source),
        "fixture_count": len(records),
        "fixtures": records,
        "leading_zero_bytes": leading,
        "physical_length": len(pcluster),
        "physical_offset": pcluster_start,
        "schema": 1,
        "source_sha256": sha256_path(payload),
        "status": "READY",
        "test": "TC168-cache-inflight",
    }
    (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()