from __future__ import annotations

import unittest
from ipaddress import IPv4Address, IPv4Network
from pathlib import Path
from unittest.mock import patch

from uvm.config import Settings
from uvm.domain import VmRecord
from uvm.errors import FirecrackerError
from uvm.firecracker.api import FirecrackerClient
from uvm.firecracker.config import build_config, vcpu_count


class RecordingFirecrackerClient(FirecrackerClient):
    def __init__(self) -> None:
        self.calls: list[tuple[str, str, object]] = []

    def _request(self, method: str, path: str, body: object) -> None:
        self.calls.append((method, path, body))


class FirecrackerConfigTests(unittest.TestCase):
    def test_uses_arm_boot_console_arguments_only_on_arm64(self) -> None:
        vm = VmRecord(
            id="vm-test",
            cpu=1,
            ram_mib=512,
            guest_ip="10.42.0.2",
            gateway="10.42.0.1",
            tap="uvm-test",
            mac="02:fc:00:00:00:01",
            socket="/tmp/firecracker.sock",
            config="/tmp/config.json",
            log="/tmp/firecracker.log",
        )
        for architecture in ("x86_64", "aarch64", "arm64"):
            with self.subTest(architecture=architecture):
                config = build_config(
                    Settings(architecture=architecture),
                    vm,
                    Path("/images/vmlinux"),
                    Path("/vm/rootfs.ext4"),
                )

                arguments = config["boot-source"]["boot_args"].split()
                self.assertEqual("keep_bootcon" in arguments, architecture != "x86_64")
                self.assertIn("console=ttyS0", arguments)
                self.assertIn("ip=10.42.0.2::10.42.0.1:255.255.255.0::eth0:off", arguments)
                client = RecordingFirecrackerClient()
                client.configure_and_start(config)
                self.assertEqual(client.calls[-2], ("PUT", "/entropy", {}))
                self.assertEqual(client.calls[-1], ("PUT", "/actions", {"action_type": "InstanceStart"}))

    def test_builds_a_writable_vm_disk_configuration(self) -> None:
        vm = VmRecord(
            id="vm-test",
            cpu=0.5,
            ram_mib=512,
            guest_ip="10.42.0.2",
            gateway="10.42.0.1",
            tap="uvm-test",
            mac="02:fc:00:00:00:01",
            socket="/tmp/firecracker.sock",
            config="/tmp/config.json",
            log="/tmp/firecracker.log",
            disk="/vm/rootfs.ext4",
        )
        config = build_config(
            Settings(),
            vm,
            Path("/images/vmlinux"),
            Path("/vm/rootfs.ext4"),
        )

        self.assertEqual(vcpu_count(0.5), 1)
        self.assertEqual(config["machine-config"]["mem_size_mib"], 512)
        self.assertEqual(config["drives"][0]["path_on_host"], "/vm/rootfs.ext4")
        self.assertIn("ip=10.42.0.2::10.42.0.1", config["boot-source"]["boot_args"])
        self.assertNotIn("password", str(config))

    def test_configures_firecracker_in_required_order_before_start(self) -> None:
        vm = VmRecord(
            id="vm-test",
            cpu=1,
            ram_mib=512,
            guest_ip="10.42.0.2",
            gateway="10.42.0.1",
            tap="uvm-test",
            mac="02:fc:00:00:00:01",
            socket="/tmp/firecracker.sock",
            config="/tmp/config.json",
            log="/tmp/firecracker.log",
            disk="/vm/rootfs.ext4",
        )
        config = build_config(
            Settings(), vm, Path("/images/vmlinux"), Path("/vm/rootfs.ext4")
        )
        client = RecordingFirecrackerClient()

        client.configure_and_start(config)

        self.assertEqual(
            [path for _method, path, _body in client.calls],
            [
                "/machine-config",
                "/boot-source",
                "/drives/rootfs",
                "/network-interfaces/eth0",
                "/entropy",
                "/actions",
            ],
        )
        self.assertEqual(
            client.calls[0][2],
            {"vcpu_count": 1, "mem_size_mib": 512, "smt": False},
        )
        self.assertEqual(client.calls[-1][2], {"action_type": "InstanceStart"})

    def test_entropy_failure_prevents_guest_boot(self) -> None:
        config = {
            "machine-config": {"vcpu_count": 1, "mem_size_mib": 512},
            "boot-source": {},
            "drives": [{}],
            "network-interfaces": [{}],
            "entropy": {},
        }
        client = RecordingFirecrackerClient()

        def request(method: str, path: str, body: object) -> None:
            client.calls.append((method, path, body))
            if path == "/entropy":
                raise FirecrackerError("entropy device unavailable")

        with patch.object(client, "_request", side_effect=request):
            with self.assertRaisesRegex(FirecrackerError, "entropy device unavailable"):
                client.configure_and_start(config)

        self.assertEqual(client.calls[-1], ("PUT", "/entropy", {}))
        self.assertNotIn("/actions", [path for _method, path, _body in client.calls])

        for entropy in (None, [], False):
            with self.subTest(entropy=entropy):
                client.calls.clear()
                config["entropy"] = entropy
                with self.assertRaisesRegex(FirecrackerError, "no valid entropy device"):
                    client.configure_and_start(config)
                self.assertEqual(client.calls, [])

    def test_uses_the_configured_network_netmask_in_guest_boot_arguments(self) -> None:
        settings = Settings(
            network=IPv4Network("10.50.0.0/16"),
            gateway=IPv4Address("10.50.0.1"),
        )
        vm = VmRecord(
            id="vm-test",
            cpu=1,
            ram_mib=512,
            guest_ip="10.50.0.2",
            gateway="10.50.0.1",
            tap="uvm-test",
            mac="02:fc:00:00:00:01",
            socket="/tmp/firecracker.sock",
            config="/tmp/config.json",
            log="/tmp/firecracker.log",
        )

        config = build_config(settings, vm, Path("/images/vmlinux"), Path("/vm/rootfs.ext4"))

        self.assertIn("ip=10.50.0.2::10.50.0.1:255.255.0.0", config["boot-source"]["boot_args"])