"""Host-side network allocation and diagnostics for Firecracker VMs."""

from __future__ import annotations

import os
import shutil
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

CommandRunner = Callable[[list[str]], subprocess.CompletedProcess[str]]

DEFAULT_NETWORK_PREFIX = "172.29."
DEFAULT_GATEWAY_SUFFIX = 1
DEFAULT_GUEST_SUFFIX = 2


@dataclass(frozen=True)
class NetworkConfig:
    vm_id: str
    tap_name: str
    guest_ip: str
    gateway_ip: str
    subnet_cidr: str
    mac_address: str
    dns_servers: tuple[str, ...] = ("1.1.1.1", "8.8.8.8")


@dataclass(frozen=True)
class NetworkDiagnostics:
    tun_available: bool
    ip_binary: str | None
    nft_binary: str | None
    iptables_binary: str | None
    ip_forward_enabled: bool
    sudo_non_interactive: bool


class TapNetworkManager:
    """Allocate per-VM TAP/NAT configuration on the host."""

    def __init__(
        self,
        *,
        enabled: bool | None = None,
        runner: CommandRunner | None = None,
        use_sudo: bool | None = None,
    ) -> None:
        if enabled is None:
            self._enabled = os.environ.get("PYRO_VM_ENABLE_NETWORK") == "1"
        else:
            self._enabled = enabled
        self._use_sudo = self._detect_sudo() if use_sudo is None else use_sudo
        self._runner = runner or self._run

    @staticmethod
    def diagnostics() -> NetworkDiagnostics:
        ip_forward = False
        try:
            ip_forward = (
                Path("/proc/sys/net/ipv4/ip_forward").read_text(encoding="utf-8").strip()
                == "1"
            )
        except OSError:
            ip_forward = False
        return NetworkDiagnostics(
            tun_available=Path("/dev/net/tun").exists(),
            ip_binary=shutil.which("ip"),
            nft_binary=shutil.which("nft"),
            iptables_binary=shutil.which("iptables"),
            ip_forward_enabled=ip_forward,
            sudo_non_interactive=TapNetworkManager._detect_sudo(),
        )

    @property
    def enabled(self) -> bool:
        return self._enabled

    def allocate(self, vm_id: str) -> NetworkConfig:
        octet = int(vm_id[:2], 16)
        third_octet = 1 + (octet % 200)
        subnet_cidr = f"{DEFAULT_NETWORK_PREFIX}{third_octet}.0/24"
        gateway_ip = f"{DEFAULT_NETWORK_PREFIX}{third_octet}.{DEFAULT_GATEWAY_SUFFIX}"
        guest_ip = f"{DEFAULT_NETWORK_PREFIX}{third_octet}.{DEFAULT_GUEST_SUFFIX}"
        tap_name = f"pyro{vm_id[:8]}"
        mac_address = f"06:00:{vm_id[0:2]}:{vm_id[2:4]}:{vm_id[4:6]}:{vm_id[6:8]}"
        config = NetworkConfig(
            vm_id=vm_id,
            tap_name=tap_name,
            guest_ip=guest_ip,
            gateway_ip=gateway_ip,
            subnet_cidr=subnet_cidr,
            mac_address=mac_address,
        )
        if self._enabled:
            try:
                self._ensure_host_network(config)
            except Exception:
                table_name = self._nft_table_name(config.vm_id)
                self._run_ignore(["nft", "delete", "table", "ip", table_name])
                self._run_ignore(
                    [
                        "iptables",
                        "-t",
                        "nat",
                        "-D",
                        "POSTROUTING",
                        "-s",
                        config.subnet_cidr,
                        "-j",
                        "MASQUERADE",
                    ]
                )
                self._run_ignore(
                    ["iptables", "-D", "FORWARD", "-i", config.tap_name, "-j", "ACCEPT"]
                )
                self._run_ignore(
                    [
                        "iptables",
                        "-D",
                        "FORWARD",
                        "-o",
                        config.tap_name,
                        "-m",
                        "conntrack",
                        "--ctstate",
                        "RELATED,ESTABLISHED",
                        "-j",
                        "ACCEPT",
                    ]
                )
                self._run_ignore(["ip", "link", "del", config.tap_name])
                raise
        return config

    def cleanup(self, config: NetworkConfig) -> None:
        if not self._enabled:
            return
        table_name = self._nft_table_name(config.vm_id)
        self._run_ignore(["nft", "delete", "table", "ip", table_name])
        self._run_ignore(
            [
                "iptables",
                "-t",
                "nat",
                "-D",
                "POSTROUTING",
                "-s",
                config.subnet_cidr,
                "-j",
                "MASQUERADE",
            ]
        )
        self._run_ignore(["iptables", "-D", "FORWARD", "-i", config.tap_name, "-j", "ACCEPT"])
        self._run_ignore(
            [
                "iptables",
                "-D",
                "FORWARD",
                "-o",
                config.tap_name,
                "-m",
                "conntrack",
                "--ctstate",
                "RELATED,ESTABLISHED",
                "-j",
                "ACCEPT",
            ]
        )
        self._run_ignore(["ip", "link", "del", config.tap_name])

    def to_metadata(self, config: NetworkConfig) -> dict[str, str]:
        return {
            "network_enabled": "true" if self._enabled else "false",
            "tap_name": config.tap_name,
            "guest_ip": config.guest_ip,
            "gateway_ip": config.gateway_ip,
            "subnet_cidr": config.subnet_cidr,
            "mac_address": config.mac_address,
            "dns_servers": ",".join(config.dns_servers),
        }

    def network_info(self, config: NetworkConfig) -> dict[str, object]:
        return {
            "network_enabled": self._enabled,
            "tap_name": config.tap_name,
            "guest_ip": config.guest_ip,
            "gateway_ip": config.gateway_ip,
            "subnet_cidr": config.subnet_cidr,
            "mac_address": config.mac_address,
            "dns_servers": list(config.dns_servers),
            "outbound_connectivity_expected": self._enabled,
        }

    def _ensure_host_network(self, config: NetworkConfig) -> None:
        diagnostics = self.diagnostics()
        if not diagnostics.tun_available:
            raise RuntimeError("/dev/net/tun is not available on this host")
        if diagnostics.ip_binary is None:
            raise RuntimeError("`ip` binary is required for TAP setup")
        if diagnostics.nft_binary is None and diagnostics.iptables_binary is None:
            raise RuntimeError("`nft` or `iptables` is required for outbound NAT setup")
        if not diagnostics.ip_forward_enabled:
            raise RuntimeError("IPv4 forwarding is disabled on this host")

        self._runner(
            [
                "ip",
                "tuntap",
                "add",
                "dev",
                config.tap_name,
                "mode",
                "tap",
                "user",
                str(os.getuid()),
            ]
        )
        self._runner(["ip", "addr", "add", f"{config.gateway_ip}/24", "dev", config.tap_name])
        self._runner(["ip", "link", "set", config.tap_name, "up"])

        if diagnostics.nft_binary is not None:
            self._ensure_nft_network(config)
            return
        self._ensure_iptables_network(config)

    def _ensure_nft_network(self, config: NetworkConfig) -> None:
        table_name = self._nft_table_name(config.vm_id)
        self._runner(["nft", "add", "table", "ip", table_name])
        self._runner(
            [
                "nft",
                "add",
                "chain",
                "ip",
                table_name,
                "postrouting",
                "{",
                "type",
                "nat",
                "hook",
                "postrouting",
                "priority",
                "srcnat",
                ";",
                "}",
            ]
        )
        self._runner(
            [
                "nft",
                "add",
                "chain",
                "ip",
                table_name,
                "forward",
                "{",
                "type",
                "filter",
                "hook",
                "forward",
                "priority",
                "filter",
                ";",
                "policy",
                "accept",
                ";",
                "}",
            ]
        )
        self._runner(
            [
                "nft",
                "add",
                "rule",
                "ip",
                table_name,
                "postrouting",
                "ip",
                "saddr",
                config.subnet_cidr,
                "masquerade",
            ]
        )
        self._runner(
            [
                "nft",
                "add",
                "rule",
                "ip",
                table_name,
                "forward",
                "iifname",
                config.tap_name,
                "accept",
            ]
        )
        self._runner(
            [
                "nft",
                "add",
                "rule",
                "ip",
                table_name,
                "forward",
                "oifname",
                config.tap_name,
                "ct",
                "state",
                "related,established",
                "accept",
            ]
        )

    def _ensure_iptables_network(self, config: NetworkConfig) -> None:
        self._runner(
            [
                "iptables",
                "-t",
                "nat",
                "-A",
                "POSTROUTING",
                "-s",
                config.subnet_cidr,
                "-j",
                "MASQUERADE",
            ]
        )
        self._runner(["iptables", "-A", "FORWARD", "-i", config.tap_name, "-j", "ACCEPT"])
        self._runner(
            [
                "iptables",
                "-A",
                "FORWARD",
                "-o",
                config.tap_name,
                "-m",
                "conntrack",
                "--ctstate",
                "RELATED,ESTABLISHED",
                "-j",
                "ACCEPT",
            ]
        )

    def _run_ignore(self, command: list[str]) -> None:
        try:
            self._runner(command)
        except RuntimeError:
            return

    @staticmethod
    def _nft_table_name(vm_id: str) -> str:
        return f"pyro_{vm_id[:12]}"

    @staticmethod
    def _detect_sudo() -> bool:
        if os.geteuid() == 0:
            return False
        if shutil.which("sudo") is None:
            return False
        completed = subprocess.run(
            ["sudo", "-n", "true"],
            text=True,
            capture_output=True,
            check=False,
        )
        return completed.returncode == 0

    def _run(self, command: list[str]) -> subprocess.CompletedProcess[str]:
        effective_command = command
        if self._use_sudo:
            effective_command = ["sudo", "-n", *command]
        completed = subprocess.run(
            effective_command,
            text=True,
            capture_output=True,
            check=False,
        )
        if completed.returncode != 0:
            stderr = completed.stderr.strip() or completed.stdout.strip()
            raise RuntimeError(f"command {' '.join(effective_command)!r} failed: {stderr}")
        return completed
