from __future__ import annotations

import hashlib
import io
import os
import tarfile
import tempfile
import unittest
from dataclasses import replace
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock, patch

from uvm.config import Settings
from uvm.errors import UvmError
from uvm.install import Installer
from uvm.integrity import load_manifest
from uvm.state import StateStore
from uvm.system import CommandRunner


class InstallerArchitectureTests(unittest.TestCase):
    def test_installs_matching_native_artifacts_with_a_verified_manifest(self) -> None:
        for machine, architecture in (
            ("x86_64", "x86_64"),
            ("aarch64", "aarch64"),
            ("arm64", "aarch64"),
        ):
            with (
                self.subTest(machine=machine),
                tempfile.TemporaryDirectory() as temporary_directory,
                patch.dict(os.environ, {"UVM_BASE": temporary_directory}, clear=True),
                patch("uvm.config.os.uname", return_value=SimpleNamespace(machine=machine)),
            ):
                settings = Settings.from_environment()
                version = settings.firecracker_version
                binaries = {
                    "firecracker": f"firecracker for {architecture}".encode(),
                    "jailer": f"jailer for {architecture}".encode(),
                }
                archive_buffer = io.BytesIO()
                with tarfile.open(fileobj=archive_buffer, mode="w:gz") as archive:
                    for name, content in binaries.items():
                        member = tarfile.TarInfo(
                            f"release-{version}-{architecture}/{name}-{version}-{architecture}"
                        )
                        member.size = len(content)
                        member.mode = 0o755
                        archive.addfile(member, io.BytesIO(content))

                archive_bytes = archive_buffer.getvalue()
                kernel_bytes = f"kernel for {architecture}".encode()
                rootfs_bytes = f"rootfs for {architecture}".encode()
                settings = replace(
                    settings,
                    firecracker_sha256=hashlib.sha256(archive_bytes).hexdigest(),
                    kernel_sha256=hashlib.sha256(kernel_bytes).hexdigest(),
                    rootfs_sha256=hashlib.sha256(rootfs_bytes).hexdigest(),
                )
                guest_base = (
                    "https://s3.amazonaws.com/spec.ccfc.min/img/quickstart_guide/"
                    f"{architecture}"
                )
                payloads = {
                    "https://github.com/firecracker-microvm/firecracker/releases/download/"
                    f"{version}/firecracker-{version}-{architecture}.tgz": archive_bytes,
                    f"{guest_base}/kernels/vmlinux.bin": kernel_bytes,
                    f"{guest_base}/rootfs/bionic.rootfs.ext4": rootfs_bytes,
                }
                runner = Mock(spec=CommandRunner)
                installer = Installer(settings, runner, StateStore(settings), emit=None)
                with (
                    patch("uvm.install.require_root"),
                    patch("uvm.install.check_kvm"),
                    patch(
                        "uvm.install.urllib.request.urlopen",
                        side_effect=lambda url, timeout: io.BytesIO(payloads[url]),
                    ) as download,
                ):
                    firecracker, assets = installer.install()

                self.assertEqual(firecracker.read_bytes(), binaries["firecracker"])
                self.assertEqual(settings.jailer_binary.read_bytes(), binaries["jailer"])
                self.assertEqual(assets.kernel.read_bytes(), kernel_bytes)
                self.assertEqual(assets.rootfs.read_bytes(), rootfs_bytes)
                self.assertEqual(download.call_count, 3)
                self.assertEqual({call.args[0] for call in download.call_args_list}, set(payloads))
                manifest = load_manifest(settings.integrity_manifest_path)
                self.assertTrue(manifest.verified)
                self.assertEqual(manifest.checksums, {
                    name: hashlib.sha256(content).hexdigest()
                    for name, content in {
                        **binaries, "kernel": kernel_bytes, "rootfs": rootfs_bytes,
                    }.items()
                })

    def test_arm64_install_requires_checksums_before_host_commands_or_downloads(self) -> None:
        with tempfile.TemporaryDirectory() as temporary_directory:
            settings = Settings(base=Path(temporary_directory), architecture="aarch64")
            runner = Mock(spec=CommandRunner)
            installer = Installer(settings, runner, StateStore(settings), emit=None)
            with (
                patch("uvm.install.require_root"),
                patch("uvm.install.check_kvm") as check_kvm,
                patch("uvm.install.urllib.request.urlopen") as download,
                self.assertRaisesRegex(UvmError, "Firecracker archive checksum is required"),
            ):
                installer.install()

            runner.run.assert_not_called()
            check_kvm.assert_not_called()
            download.assert_not_called()
            self.assertFalse(settings.integrity_manifest_path.exists())