"""Firecracker launch-plan generation for pyro VMs."""

from __future__ import annotations

import ipaddress
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Protocol

DEFAULT_GUEST_CID_OFFSET = 10_000
DEFAULT_VSOCK_PORT = 5005


@dataclass(frozen=True)
class FirecrackerLaunchPlan:
    """On-disk config artifacts needed to boot a guest."""

    api_socket_path: Path
    config_path: Path
    guest_network_path: Path
    guest_exec_path: Path
    guest_cid: int
    vsock_port: int
    config: dict[str, Any]


class VmInstanceLike(Protocol):
    vm_id: str
    vcpu_count: int
    mem_mib: int
    workdir: Path
    metadata: dict[str, str]
    network: Any


def _guest_cid(vm_id: str) -> int:
    return DEFAULT_GUEST_CID_OFFSET + int(vm_id[:4], 16)


def build_launch_plan(instance: VmInstanceLike) -> FirecrackerLaunchPlan:
    guest_cid = _guest_cid(instance.vm_id)
    vsock_port = DEFAULT_VSOCK_PORT
    api_socket_path = instance.workdir / "firecracker.sock"
    config_path = instance.workdir / "firecracker-config.json"
    guest_network_path = instance.workdir / "guest-network.json"
    guest_exec_path = instance.workdir / "guest-exec.json"

    boot_args = [
        "console=ttyS0",
        "reboot=k",
        "panic=1",
        "acpi=off",
        "pci=off",
        "root=/dev/vda",
        "rootfstype=ext4",
        "rw",
        "init=/opt/pyro/bin/pyro-init",
    ]
    if instance.network is not None:
        network = ipaddress.ip_network(instance.network.subnet_cidr, strict=False)
        prefixlen = network.prefixlen
        boot_args.extend(
            [
                f"pyro.guest_ip={instance.network.guest_ip}",
                f"pyro.gateway_ip={instance.network.gateway_ip}",
                f"pyro.netmask={prefixlen}",
                f"pyro.dns={','.join(instance.network.dns_servers)}",
            ]
        )

    config: dict[str, Any] = {
        "boot-source": {
            "kernel_image_path": instance.metadata["kernel_image"],
            "boot_args": " ".join(boot_args),
        },
        "drives": [
            {
                "drive_id": "rootfs",
                "path_on_host": instance.metadata["rootfs_image"],
                "is_root_device": True,
                "is_read_only": False,
            }
        ],
        "machine-config": {
            "vcpu_count": instance.vcpu_count,
            "mem_size_mib": instance.mem_mib,
            "smt": False,
        },
        "network-interfaces": [],
        "vsock": {
            "guest_cid": guest_cid,
            "uds_path": str(instance.workdir / "vsock.sock"),
        },
    }

    if instance.network is not None:
        config["network-interfaces"] = [
            {
                "iface_id": "eth0",
                "guest_mac": instance.network.mac_address,
                "host_dev_name": instance.network.tap_name,
            }
        ]

    guest_network = {
        "guest_ip": instance.network.guest_ip if instance.network is not None else None,
        "gateway_ip": instance.network.gateway_ip if instance.network is not None else None,
        "subnet_cidr": instance.network.subnet_cidr if instance.network is not None else None,
        "dns_servers": list(instance.network.dns_servers) if instance.network is not None else [],
        "tap_name": instance.network.tap_name if instance.network is not None else None,
    }
    guest_exec = {
        "transport": "vsock",
        "guest_cid": guest_cid,
        "port": vsock_port,
    }

    config_path.write_text(json.dumps(config, indent=2, sort_keys=True), encoding="utf-8")
    guest_network_path.write_text(
        json.dumps(guest_network, indent=2, sort_keys=True), encoding="utf-8"
    )
    guest_exec_path.write_text(json.dumps(guest_exec, indent=2, sort_keys=True), encoding="utf-8")

    return FirecrackerLaunchPlan(
        api_socket_path=api_socket_path,
        config_path=config_path,
        guest_network_path=guest_network_path,
        guest_exec_path=guest_exec_path,
        guest_cid=guest_cid,
        vsock_port=vsock_port,
        config=config,
    )
