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

import argparse
from dataclasses import asdict
import json
from pathlib import Path
import shutil

from scripts.kernel_artifact import ARTIFACT_NAME, KernelArtifactError, verify_kernel_dependency


TARGET_RELATIVE = Path("opt/matrix-image-bootstrap/axp313a")
DEFAULT_RESERVE_BYTES = 256 * 1024 * 1024
EXPECTED_FILES = {ARTIFACT_NAME, "SHA256SUMS", "METADATA.json"}


class ImageKernelPayloadError(ValueError):
    pass


def _root(path: Path) -> Path:
    resolved = Path(path).resolve()
    if not resolved.is_dir() or resolved == Path(resolved.anchor):
        raise ImageKernelPayloadError("image root must be an explicit mounted or isolated directory")
    return resolved


def verify_staged_payload(root: Path, dependency: Path) -> dict[str, object]:
    root = _root(root)
    dependency = Path(dependency).resolve()
    expected = verify_kernel_dependency(dependency)
    target = root / TARGET_RELATIVE
    if not target.is_dir() or {path.name for path in target.iterdir()} != EXPECTED_FILES:
        raise ImageKernelPayloadError("staged image kernel payload files are not exact")
    for name in ("SHA256SUMS", "METADATA.json"):
        if (target / name).read_bytes() != (dependency / name).read_bytes():
            raise ImageKernelPayloadError(f"staged image kernel {name} differs from the dependency")
    actual = verify_kernel_dependency(target)
    if actual != expected:
        raise ImageKernelPayloadError("staged image kernel payload identity changed")
    return {"target": TARGET_RELATIVE.as_posix(), **asdict(actual)}


def stage_payload(
    root: Path,
    dependency: Path,
    *,
    reserve_bytes: int = DEFAULT_RESERVE_BYTES,
) -> dict[str, object]:
    root = _root(root)
    dependency = Path(dependency).resolve()
    if type(reserve_bytes) is not int or reserve_bytes < 0:
        raise ImageKernelPayloadError("reserve_bytes must be a non-negative integer")
    info = verify_kernel_dependency(dependency)
    required = info.artifact_bytes + reserve_bytes
    if shutil.disk_usage(root).free < required:
        raise ImageKernelPayloadError("copied image root filesystem lacks room for the offline kernel payload")
    target = root / TARGET_RELATIVE
    if target.exists():
        raise ImageKernelPayloadError("image kernel payload target already exists")
    created = False
    try:
        target.mkdir(parents=True, mode=0o755)
        created = True
        for name in EXPECTED_FILES:
            shutil.copy2(dependency / name, target / name)
        return verify_staged_payload(root, dependency)
    except BaseException:
        if created:
            shutil.rmtree(target, ignore_errors=True)
        raise


def main() -> int:
    parser = argparse.ArgumentParser(description="Stage or verify the offline kernel in a copied IMG root")
    parser.add_argument("action", choices=("stage", "verify"))
    parser.add_argument("--root", type=Path, required=True)
    parser.add_argument("--dependency", type=Path, required=True)
    parser.add_argument("--reserve-bytes", type=int, default=DEFAULT_RESERVE_BYTES)
    args = parser.parse_args()
    if args.action == "stage":
        result = stage_payload(args.root, args.dependency, reserve_bytes=args.reserve_bytes)
    else:
        result = verify_staged_payload(args.root, args.dependency)
    print(json.dumps(result, ensure_ascii=False, sort_keys=True, indent=2))
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (ImageKernelPayloadError, KernelArtifactError) as exc:
        print(json.dumps({"ok": False, "error": str(exc)}, ensure_ascii=False))
        raise SystemExit(1)