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

import argparse
from dataclasses import asdict, dataclass
import json
import os
from pathlib import Path
import tarfile

from scripts.kernel_artifact import ARTIFACT_NAME, REQUIRED_BOOT_FILES, verify_kernel_dependency
from scripts.render_dual_kernel_boot import BootScriptError, render


DEFAULT_SAFETY_BYTES = 32 * 1024 * 1024
CONTROL_FILE_COUNT = 2  # matrix-original.SHA256SUMS and matrix-kernel-good


class BootSpaceError(ValueError):
    pass


@dataclass(frozen=True)
class BootSpaceBudget:
    required_bytes: int
    safety_bytes: int
    total_bytes: int
    available_bytes: int
    cluster_bytes: int


def _allocated(size: int, cluster_bytes: int) -> int:
    if type(size) is not int or size < 0:
        raise BootSpaceError("boot file size must be a non-negative integer")
    if type(cluster_bytes) is not int or cluster_bytes <= 0:
        raise BootSpaceError("boot cluster size must be a positive integer")
    return max(1, (size + cluster_bytes - 1) // cluster_bytes) * cluster_bytes


def calculate_budget(
    *,
    candidate_sizes: dict[str, int],
    boot_cmd: bytes,
    boot_scr: bytes,
    cluster_bytes: int,
    available_bytes: int,
    safety_bytes: int = DEFAULT_SAFETY_BYTES,
) -> BootSpaceBudget:
    if set(candidate_sizes) != REQUIRED_BOOT_FILES:
        raise BootSpaceError("candidate boot file set is not exact")
    if type(available_bytes) is not int or available_bytes < 0:
        raise BootSpaceError("available boot bytes must be a non-negative integer")
    if type(safety_bytes) is not int or safety_bytes < 0:
        raise BootSpaceError("boot safety bytes must be a non-negative integer")
    try:
        managed_cmd = render(boot_cmd.decode("utf-8")).encode("utf-8")
    except (UnicodeError, BootScriptError) as exc:
        raise BootSpaceError("original boot.cmd cannot be rendered safely") from exc
    # mkimage adds a 64-byte legacy image header. Count both temporary and final
    # managed scripts so the estimate remains safe even when FAT cannot reuse the
    # original chains during replacement.
    managed_scr_size = len(managed_cmd) + 64
    allocations = list(candidate_sizes.values())
    allocations.extend((len(boot_cmd), len(boot_scr)))  # rollback copies
    allocations.extend((len(managed_cmd), managed_scr_size))  # mktemp files
    allocations.extend((len(managed_cmd), managed_scr_size))  # final replacements
    allocations.extend(1 for _ in range(CONTROL_FILE_COUNT))
    required = sum(_allocated(size, cluster_bytes) for size in allocations)
    total = required + safety_bytes
    return BootSpaceBudget(required, safety_bytes, total, available_bytes, cluster_bytes)


def candidate_sizes_from_dependency(dependency: Path) -> dict[str, int]:
    dependency = Path(dependency).resolve()
    verify_kernel_dependency(dependency)
    with tarfile.open(dependency / ARTIFACT_NAME, mode="r:gz") as archive:
        normalized = {
            member.name.removeprefix("./"): member.size
            for member in archive.getmembers()
            if member.isfile()
        }
    return {name: normalized[name] for name in REQUIRED_BOOT_FILES}


def candidate_sizes_from_artifact_root(root: Path) -> dict[str, int]:
    root = Path(root).resolve()
    result: dict[str, int] = {}
    for name in REQUIRED_BOOT_FILES:
        path = root / name
        if not path.is_file():
            raise BootSpaceError(f"candidate boot file is missing: {name}")
        result[name] = path.stat().st_size
    return result


def check_mounted_boot(
    boot_root: Path,
    artifact_root: Path,
    *,
    safety_bytes: int = DEFAULT_SAFETY_BYTES,
) -> BootSpaceBudget:
    boot_root = Path(boot_root).resolve()
    if not boot_root.is_dir() or boot_root == Path(boot_root.anchor):
        raise BootSpaceError("boot root must be an explicit mounted directory")
    boot_cmd = (boot_root / "boot.cmd").read_bytes()
    boot_scr = (boot_root / "boot.scr").read_bytes()
    stats = os.statvfs(boot_root)
    cluster_bytes = stats.f_frsize or stats.f_bsize
    available_bytes = stats.f_bavail * cluster_bytes
    budget = calculate_budget(
        candidate_sizes=candidate_sizes_from_artifact_root(artifact_root),
        boot_cmd=boot_cmd,
        boot_scr=boot_scr,
        cluster_bytes=cluster_bytes,
        available_bytes=available_bytes,
        safety_bytes=safety_bytes,
    )
    if budget.available_bytes < budget.total_bytes:
        raise BootSpaceError(
            "boot filesystem lacks room for the candidate kernel and rollback files: "
            f"available={budget.available_bytes}, required={budget.required_bytes}, "
            f" safety={budget.safety_bytes}, total={budget.total_bytes}"
        )
    return budget


def main() -> int:
    parser = argparse.ArgumentParser(description="Validate safe /boot space before kernel installation")
    parser.add_argument("--boot-root", type=Path, required=True)
    parser.add_argument("--artifact-root", type=Path, required=True)
    parser.add_argument("--safety-bytes", type=int, default=DEFAULT_SAFETY_BYTES)
    args = parser.parse_args()
    result = check_mounted_boot(
        args.boot_root,
        args.artifact_root,
        safety_bytes=args.safety_bytes,
    )
    print(json.dumps(asdict(result), ensure_ascii=False, sort_keys=True))
    return 0


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