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

import argparse
from dataclasses import asdict
import hashlib
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")
APP_RELATIVE = Path("app/MSCBOOT.TGZ")
KERNEL_RELATIVE = Path("axp313a")
DEFAULT_RESERVE_BYTES = 256 * 1024 * 1024
KERNEL_FILES = {ARTIFACT_NAME, "SHA256SUMS", "METADATA.json"}


class ImageBootstrapPayloadError(ValueError):
    pass


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


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(4 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def verify_staged_payload(root: Path, bundle: Path, dependency: Path) -> dict[str, object]:
    root = _root(root)
    bundle = Path(bundle).resolve()
    dependency = Path(dependency).resolve()
    expected_kernel = verify_kernel_dependency(dependency)
    target = root / TARGET_RELATIVE
    app_target = target / APP_RELATIVE
    kernel_target = target / KERNEL_RELATIVE
    if not target.is_dir() or {path.name for path in target.iterdir()} != {"app", "axp313a"}:
        raise ImageBootstrapPayloadError("staged image bootstrap payload directories are not exact")
    if not app_target.is_file() or {path.name for path in app_target.parent.iterdir()} != {"MSCBOOT.TGZ"}:
        raise ImageBootstrapPayloadError("staged image application payload files are not exact")
    if app_target.stat().st_size != bundle.stat().st_size or _sha256(app_target) != _sha256(bundle):
        raise ImageBootstrapPayloadError("staged image application payload differs from the build bundle")
    if not kernel_target.is_dir() or {path.name for path in kernel_target.iterdir()} != KERNEL_FILES:
        raise ImageBootstrapPayloadError("staged image kernel payload files are not exact")
    for name in ("SHA256SUMS", "METADATA.json"):
        if (kernel_target / name).read_bytes() != (dependency / name).read_bytes():
            raise ImageBootstrapPayloadError(f"staged image kernel {name} differs from the dependency")
    actual_kernel = verify_kernel_dependency(kernel_target)
    if actual_kernel != expected_kernel:
        raise ImageBootstrapPayloadError("staged image kernel payload identity changed")
    return {
        "target": TARGET_RELATIVE.as_posix(),
        "bundle_path": (TARGET_RELATIVE / APP_RELATIVE).as_posix(),
        "bundle_bytes": app_target.stat().st_size,
        "bundle_sha256": _sha256(app_target),
        "kernel": asdict(actual_kernel),
    }


def stage_payload(
    root: Path,
    bundle: Path,
    dependency: Path,
    *,
    reserve_bytes: int = DEFAULT_RESERVE_BYTES,
) -> dict[str, object]:
    root = _root(root)
    bundle = Path(bundle).resolve()
    dependency = Path(dependency).resolve()
    if not bundle.is_file():
        raise ImageBootstrapPayloadError("image application bundle is missing")
    if type(reserve_bytes) is not int or reserve_bytes < 0:
        raise ImageBootstrapPayloadError("reserve_bytes must be a non-negative integer")
    verify_kernel_dependency(dependency)
    payload_bytes = bundle.stat().st_size + sum((dependency / name).stat().st_size for name in KERNEL_FILES)
    if shutil.disk_usage(root).free < payload_bytes + reserve_bytes:
        raise ImageBootstrapPayloadError("copied image root filesystem lacks room for both offline payloads")
    target = root / TARGET_RELATIVE
    if target.exists():
        raise ImageBootstrapPayloadError("image bootstrap payload target already exists")
    created = False
    try:
        (target / APP_RELATIVE.parent).mkdir(parents=True, mode=0o755)
        (target / KERNEL_RELATIVE).mkdir(mode=0o755)
        created = True
        shutil.copy2(bundle, target / APP_RELATIVE)
        for name in KERNEL_FILES:
            shutil.copy2(dependency / name, target / KERNEL_RELATIVE / name)
        return verify_staged_payload(root, bundle, dependency)
    except BaseException:
        if created:
            shutil.rmtree(target, ignore_errors=True)
        raise


def main() -> int:
    parser = argparse.ArgumentParser(description="Stage or verify rootfs image bootstrap payloads")
    parser.add_argument("action", choices=("stage", "verify"))
    parser.add_argument("--root", type=Path, required=True)
    parser.add_argument("--bundle", 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.bundle, args.dependency, reserve_bytes=args.reserve_bytes)
    else:
        result = verify_staged_payload(args.root, args.bundle, args.dependency)
    print(json.dumps(result, ensure_ascii=False, sort_keys=True, indent=2))
    return 0


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