#!/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 load_spec(path: Path) -> dict[str, object]:
    spec = json.loads(path.read_text(encoding="ascii"))
    if (
        spec.get("schema") != 1
        or spec.get("batch") != "B27"
        or spec.get("candidate") != "P15-083"
        or spec.get("test") != "TC176-stream-runtime"
    ):
        raise SystemExit("invalid B27 stream fixture spec")
    if set(spec.get("codecs", {})) != {"lzma", "deflate", "zstd"}:
        raise SystemExit("B27 codec set changed")
    return spec


def make_source(root: Path, spec: dict[str, object]) -> bytes:
    root.mkdir(parents=True)
    source_spec = spec["source"]
    content = b"".join(
        f"P15-083-{index % 64:02d}:alpha-beta-gamma-delta:{(index * 17) % 256:02x}\n".encode("ascii")
        for index in range(source_spec["line_count"])
    )
    if len(content) != source_spec["size"] or sha256_bytes(content) != source_spec["sha256"]:
        raise SystemExit("B27 deterministic payload identity changed")
    payload = root / "payload.bin"
    payload.write_bytes(content)
    os.utime(payload, (0, 0))
    os.utime(root, (0, 0))
    return content


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{completed.stdout}"
        )


def build_codec(
    codec: str,
    codec_spec: dict[str, object],
    spec: dict[str, object],
    source_dir: Path,
    images: Path,
) -> list[dict[str, object]]:
    valid = images / f"{codec}-valid.erofs"
    run(["mkfs.erofs", *codec_spec["mkfs_args"], str(valid), str(source_dir)], images)
    valid_data = valid.read_bytes()
    if (
        len(valid_data) != codec_spec["expected_image_size"]
        or sha256_bytes(valid_data) != codec_spec["expected_image_sha256"]
    ):
        raise SystemExit(f"{codec} valid image is not the G04 reproducible fixture")

    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"]:
        raise SystemExit(f"{codec} pcluster padding/stream boundary changed")

    tail = bytes.fromhex(spec["tail_bytes_hex"])
    if not tail or any(value == 0 for value in tail) or leading <= len(tail):
        raise SystemExit("B27 trailing mutation is not nonzero or lacks pcluster room")
    variants: dict[str, bytearray] = {}
    tail_data = bytearray(valid_data)
    tail_data[start + leading - len(tail) : end] = stream + tail
    variants["tail"] = tail_data

    truncated_data = bytearray(valid_data)
    truncated_data[start + leading : end] = b"\0" + stream[:-1]
    variants["truncated"] = truncated_data

    corruption_start = codec_spec["corruption_start"]
    if not (codec_spec["gate"]["partial_consumed"] < corruption_start < len(stream)):
        raise SystemExit(f"{codec} corruption is not after partial oracle consumption")
    corrupt_data = bytearray(valid_data)
    corrupt_data[start + leading : end] = (
        stream[:corruption_start] + b"\0" * (len(stream) - corruption_start)
    )
    if corrupt_data == valid_data:
        raise SystemExit(f"{codec} corruption mutation changed no bytes")
    variants["corrupt"] = corrupt_data

    records = [
        {
            "class": "valid",
            "codec": codec,
            "expected_errno": 0,
            "path": valid.name,
            "sha256": sha256_path(valid),
        }
    ]
    for kind, data in variants.items():
        path = images / f"{codec}-{kind}.erofs"
        path.write_bytes(data)
        records.append(
            {
                "class": kind,
                "codec": codec,
                "expected_errno": 0 if kind == "corrupt-partial" else 97,
                "path": path.name,
                "sha256": sha256_path(path),
            }
        )
    return records


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 B27 output: {args.output}")
    if shutil.which("mkfs.erofs") is None:
        raise SystemExit("mkfs.erofs is required")
    args.output.mkdir(parents=True)
    source_dir = args.output / "source"
    images = args.output / "images"
    images.mkdir()
    make_source(source_dir, spec)
    records = []
    for codec in sorted(spec["codecs"]):
        records.extend(
            build_codec(codec, spec["codecs"][codec], spec, source_dir, images)
        )
    index = {
        "batch": "B27",
        "candidate": "P15-083",
        "fixture_count": len(records),
        "fixtures": records,
        "schema": 1,
        "source_sha256": spec["source"]["sha256"],
        "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()