"""Firewall rule management for port forwarding."""

from __future__ import annotations

import shutil
import subprocess
from enum import Enum
from typing import NamedTuple

from gaol.networking.models import Protocol


class FirewallBackend(str, Enum):
    """Firewall backend type."""

    IPTABLES = "iptables"
    NFTABLES = "nftables"
    AUTO = "auto"


class FirewallError(Exception):
    """Firewall operation error."""

    pass


class ForwardingRule(NamedTuple):
    """A port forwarding rule."""

    host_port: int
    guest_ip: str
    guest_port: int
    protocol: Protocol
    host_ip: str = "0.0.0.0"
    comment: str = ""


# Chain/table names for gaol rules
GAOL_CHAIN = "GAOL"
GAOL_NFT_TABLE = "gaol"


def detect_firewall_backend() -> FirewallBackend:
    """Detect the available firewall backend.

    Returns:
        FirewallBackend that is available

    Raises:
        FirewallError: If no firewall backend is available
    """
    # Prefer nftables if available and active
    if shutil.which("nft"):
        result = subprocess.run(
            ["nft", "list", "tables"],
            capture_output=True,
        )
        if result.returncode == 0:
            return FirewallBackend.NFTABLES

    if shutil.which("iptables"):
        return FirewallBackend.IPTABLES

    raise FirewallError("No firewall backend available (need iptables or nftables)")


def _run_cmd(
    cmd: list[str], use_sudo: bool = True
) -> subprocess.CompletedProcess[str]:
    """Run a command with optional sudo."""
    if use_sudo:
        cmd = ["sudo"] + cmd
    return subprocess.run(cmd, capture_output=True, text=True)


class FirewallManager:
    """Manages firewall rules for port forwarding."""

    def __init__(
        self,
        backend: FirewallBackend = FirewallBackend.AUTO,
        use_sudo: bool = True,
    ) -> None:
        self.use_sudo = use_sudo
        self._setup_done = False

        if backend == FirewallBackend.AUTO:
            self.backend = detect_firewall_backend()
        else:
            self.backend = backend

    def setup(self) -> None:
        """Set up firewall chains/tables for gaol.

        Creates the necessary chains/tables if they don't exist.
        """
        if self._setup_done:
            return

        if self.backend == FirewallBackend.NFTABLES:
            self._setup_nftables()
        else:
            self._setup_iptables()

        self._setup_done = True

    def _setup_nftables(self) -> None:
        """Set up nftables table and chains."""
        # Create table if not exists
        _run_cmd(["nft", "add", "table", "ip", GAOL_NFT_TABLE], self.use_sudo)

        # Create prerouting chain for DNAT
        _run_cmd(
            [
                "nft",
                "add",
                "chain",
                "ip",
                GAOL_NFT_TABLE,
                "prerouting",
                "{ type nat hook prerouting priority -100 ; }",
            ],
            self.use_sudo,
        )

        # Create postrouting chain for SNAT/masquerade
        _run_cmd(
            [
                "nft",
                "add",
                "chain",
                "ip",
                GAOL_NFT_TABLE,
                "postrouting",
                "{ type nat hook postrouting priority 100 ; }",
            ],
            self.use_sudo,
        )

        # Create filter chain for forwarding
        _run_cmd(
            [
                "nft",
                "add",
                "chain",
                "ip",
                GAOL_NFT_TABLE,
                "forward",
                "{ type filter hook forward priority 0 ; }",
            ],
            self.use_sudo,
        )

    def _setup_iptables(self) -> None:
        """Set up iptables chains."""
        # Create custom chain in nat table
        result = _run_cmd(
            ["iptables", "-t", "nat", "-N", GAOL_CHAIN], self.use_sudo
        )

        if result.returncode == 0:
            # Chain created, add jump rules
            _run_cmd(
                ["iptables", "-t", "nat", "-A", "PREROUTING", "-j", GAOL_CHAIN],
                self.use_sudo,
            )
            _run_cmd(
                ["iptables", "-t", "nat", "-A", "OUTPUT", "-j", GAOL_CHAIN],
                self.use_sudo,
            )

        # Create custom chain in filter table for forwarding
        _run_cmd(["iptables", "-N", f"{GAOL_CHAIN}_FWD"], self.use_sudo)
        _run_cmd(
            ["iptables", "-A", "FORWARD", "-j", f"{GAOL_CHAIN}_FWD"],
            self.use_sudo,
        )

    def add_forwarding_rule(self, rule: ForwardingRule) -> None:
        """Add a port forwarding rule.

        Args:
            rule: Forwarding rule to add

        Raises:
            FirewallError: If rule creation fails
        """
        self.setup()

        if self.backend == FirewallBackend.NFTABLES:
            self._add_nft_rule(rule)
        else:
            self._add_iptables_rule(rule)

    def _add_nft_rule(self, rule: ForwardingRule) -> None:
        """Add forwarding rule using nftables."""
        proto = rule.protocol.value
        if rule.protocol == Protocol.BOTH:
            # Add both TCP and UDP rules
            for p in [Protocol.TCP, Protocol.UDP]:
                self._add_nft_rule(
                    ForwardingRule(
                        rule.host_port,
                        rule.guest_ip,
                        rule.guest_port,
                        p,
                        rule.host_ip,
                        rule.comment,
                    )
                )
            return

        # DNAT rule (prerouting)
        dnat_rule = [
            "nft",
            "add",
            "rule",
            "ip",
            GAOL_NFT_TABLE,
            "prerouting",
            proto,
            "dport",
            str(rule.host_port),
            "dnat",
            "to",
            f"{rule.guest_ip}:{rule.guest_port}",
        ]
        if rule.comment:
            dnat_rule.extend(["comment", f'"{rule.comment}"'])

        result = _run_cmd(dnat_rule, self.use_sudo)
        if result.returncode != 0:
            raise FirewallError(f"Failed to add DNAT rule: {result.stderr}")

        # Forward allow rule
        fwd_rule = [
            "nft",
            "add",
            "rule",
            "ip",
            GAOL_NFT_TABLE,
            "forward",
            "ip",
            "daddr",
            rule.guest_ip,
            proto,
            "dport",
            str(rule.guest_port),
            "accept",
        ]
        result = _run_cmd(fwd_rule, self.use_sudo)
        if result.returncode != 0:
            raise FirewallError(f"Failed to add forward rule: {result.stderr}")

    def _add_iptables_rule(self, rule: ForwardingRule) -> None:
        """Add forwarding rule using iptables."""
        proto = rule.protocol.value
        if rule.protocol == Protocol.BOTH:
            for p in [Protocol.TCP, Protocol.UDP]:
                self._add_iptables_rule(
                    ForwardingRule(
                        rule.host_port,
                        rule.guest_ip,
                        rule.guest_port,
                        p,
                        rule.host_ip,
                        rule.comment,
                    )
                )
            return

        # DNAT rule
        dnat_cmd = [
            "iptables",
            "-t",
            "nat",
            "-A",
            GAOL_CHAIN,
            "-p",
            proto,
            "--dport",
            str(rule.host_port),
            "-j",
            "DNAT",
            "--to-destination",
            f"{rule.guest_ip}:{rule.guest_port}",
        ]
        if rule.host_ip != "0.0.0.0":
            dnat_cmd.extend(["-d", rule.host_ip])
        if rule.comment:
            dnat_cmd.extend(["-m", "comment", "--comment", rule.comment])

        result = _run_cmd(dnat_cmd, self.use_sudo)
        if result.returncode != 0:
            raise FirewallError(f"Failed to add DNAT rule: {result.stderr}")

        # Forward allow rule
        fwd_cmd = [
            "iptables",
            "-A",
            f"{GAOL_CHAIN}_FWD",
            "-d",
            rule.guest_ip,
            "-p",
            proto,
            "--dport",
            str(rule.guest_port),
            "-j",
            "ACCEPT",
        ]
        result = _run_cmd(fwd_cmd, self.use_sudo)
        if result.returncode != 0:
            raise FirewallError(f"Failed to add forward rule: {result.stderr}")

    def remove_forwarding_rule(self, rule: ForwardingRule) -> None:
        """Remove a port forwarding rule.

        Args:
            rule: Forwarding rule to remove
        """
        if self.backend == FirewallBackend.NFTABLES:
            self._remove_nft_rule(rule)
        else:
            self._remove_iptables_rule(rule)

    def _remove_nft_rule(self, rule: ForwardingRule) -> None:
        """Remove nftables rule by pattern matching."""
        proto = rule.protocol.value
        if rule.protocol == Protocol.BOTH:
            for p in [Protocol.TCP, Protocol.UDP]:
                self._remove_nft_rule(
                    ForwardingRule(
                        rule.host_port,
                        rule.guest_ip,
                        rule.guest_port,
                        p,
                        rule.host_ip,
                        rule.comment,
                    )
                )
            return

        # List rules and find handle for prerouting
        result = _run_cmd(
            [
                "nft",
                "-a",
                "list",
                "chain",
                "ip",
                GAOL_NFT_TABLE,
                "prerouting",
            ],
            self.use_sudo,
        )

        if result.returncode == 0:
            for line in result.stdout.splitlines():
                if (
                    f"dport {rule.host_port}" in line
                    and proto in line
                    and "handle" in line
                ):
                    # Extract handle number
                    handle = line.split("handle")[-1].strip()
                    _run_cmd(
                        [
                            "nft",
                            "delete",
                            "rule",
                            "ip",
                            GAOL_NFT_TABLE,
                            "prerouting",
                            "handle",
                            handle,
                        ],
                        self.use_sudo,
                    )

        # Remove forward rule
        result = _run_cmd(
            ["nft", "-a", "list", "chain", "ip", GAOL_NFT_TABLE, "forward"],
            self.use_sudo,
        )

        if result.returncode == 0:
            for line in result.stdout.splitlines():
                if (
                    f"dport {rule.guest_port}" in line
                    and rule.guest_ip in line
                    and "handle" in line
                ):
                    handle = line.split("handle")[-1].strip()
                    _run_cmd(
                        [
                            "nft",
                            "delete",
                            "rule",
                            "ip",
                            GAOL_NFT_TABLE,
                            "forward",
                            "handle",
                            handle,
                        ],
                        self.use_sudo,
                    )

    def _remove_iptables_rule(self, rule: ForwardingRule) -> None:
        """Remove iptables rule."""
        proto = rule.protocol.value
        if rule.protocol == Protocol.BOTH:
            for p in [Protocol.TCP, Protocol.UDP]:
                self._remove_iptables_rule(
                    ForwardingRule(
                        rule.host_port,
                        rule.guest_ip,
                        rule.guest_port,
                        p,
                        rule.host_ip,
                        rule.comment,
                    )
                )
            return

        # Remove DNAT rule
        _run_cmd(
            [
                "iptables",
                "-t",
                "nat",
                "-D",
                GAOL_CHAIN,
                "-p",
                proto,
                "--dport",
                str(rule.host_port),
                "-j",
                "DNAT",
                "--to-destination",
                f"{rule.guest_ip}:{rule.guest_port}",
            ],
            self.use_sudo,
        )

        # Remove forward rule
        _run_cmd(
            [
                "iptables",
                "-D",
                f"{GAOL_CHAIN}_FWD",
                "-d",
                rule.guest_ip,
                "-p",
                proto,
                "--dport",
                str(rule.guest_port),
                "-j",
                "ACCEPT",
            ],
            self.use_sudo,
        )

    def clear_rules_for_vm(self, vm_ip: str) -> None:
        """Remove all forwarding rules for a VM.

        Args:
            vm_ip: IP address of the VM
        """
        if self.backend == FirewallBackend.NFTABLES:
            # List and remove all rules targeting this IP
            for chain in ["prerouting", "forward"]:
                result = _run_cmd(
                    ["nft", "-a", "list", "chain", "ip", GAOL_NFT_TABLE, chain],
                    self.use_sudo,
                )
                if result.returncode == 0:
                    for line in result.stdout.splitlines():
                        if vm_ip in line and "handle" in line:
                            handle = line.split("handle")[-1].strip()
                            _run_cmd(
                                [
                                    "nft",
                                    "delete",
                                    "rule",
                                    "ip",
                                    GAOL_NFT_TABLE,
                                    chain,
                                    "handle",
                                    handle,
                                ],
                                self.use_sudo,
                            )
        else:
            # For iptables, list rules with line numbers and delete
            # List DNAT rules
            result = _run_cmd(
                ["iptables", "-t", "nat", "-L", GAOL_CHAIN, "-n", "--line-numbers"],
                self.use_sudo,
            )
            if result.returncode == 0:
                # Parse in reverse to avoid line number shifting
                lines_to_delete = []
                for line in result.stdout.splitlines():
                    if vm_ip in line:
                        parts = line.split()
                        if parts and parts[0].isdigit():
                            lines_to_delete.append(int(parts[0]))

                for line_num in reversed(lines_to_delete):
                    _run_cmd(
                        [
                            "iptables",
                            "-t",
                            "nat",
                            "-D",
                            GAOL_CHAIN,
                            str(line_num),
                        ],
                        self.use_sudo,
                    )

            # List forward rules
            result = _run_cmd(
                ["iptables", "-L", f"{GAOL_CHAIN}_FWD", "-n", "--line-numbers"],
                self.use_sudo,
            )
            if result.returncode == 0:
                lines_to_delete = []
                for line in result.stdout.splitlines():
                    if vm_ip in line:
                        parts = line.split()
                        if parts and parts[0].isdigit():
                            lines_to_delete.append(int(parts[0]))

                for line_num in reversed(lines_to_delete):
                    _run_cmd(
                        [
                            "iptables",
                            "-D",
                            f"{GAOL_CHAIN}_FWD",
                            str(line_num),
                        ],
                        self.use_sudo,
                    )

    def list_rules(self) -> list[ForwardingRule]:
        """List all active forwarding rules.

        Returns:
            List of ForwardingRule objects
        """
        rules: list[ForwardingRule] = []

        if self.backend == FirewallBackend.NFTABLES:
            result = _run_cmd(
                ["nft", "list", "chain", "ip", GAOL_NFT_TABLE, "prerouting"],
                self.use_sudo,
            )
            # Parse nftables output (simplified parsing)
            if result.returncode == 0:
                for line in result.stdout.splitlines():
                    if "dnat to" in line:
                        # Try to parse: tcp dport 8080 dnat to 192.168.1.2:80
                        parts = line.strip().split()
                        try:
                            proto_idx = (
                                parts.index("tcp")
                                if "tcp" in parts
                                else parts.index("udp")
                            )
                            proto = Protocol(parts[proto_idx])
                            dport_idx = parts.index("dport") + 1
                            host_port = int(parts[dport_idx])
                            dnat_idx = parts.index("to") + 1
                            dest = parts[dnat_idx]
                            guest_ip, guest_port = dest.split(":")
                            rules.append(
                                ForwardingRule(
                                    host_port=host_port,
                                    guest_ip=guest_ip,
                                    guest_port=int(guest_port),
                                    protocol=proto,
                                )
                            )
                        except (ValueError, IndexError):
                            continue
        else:
            result = _run_cmd(
                ["iptables", "-t", "nat", "-L", GAOL_CHAIN, "-n"],
                self.use_sudo,
            )
            if result.returncode == 0:
                for line in result.stdout.splitlines():
                    if "DNAT" in line:
                        parts = line.split()
                        try:
                            proto = Protocol(parts[1].lower())
                            # Find dpt: and to:
                            for part in parts:
                                if part.startswith("dpt:"):
                                    host_port = int(part.split(":")[1])
                                if part.startswith("to:"):
                                    dest = part.split(":")[1:]
                                    guest_ip = dest[0]
                                    guest_port = int(dest[1]) if len(dest) > 1 else host_port
                            rules.append(
                                ForwardingRule(
                                    host_port=host_port,
                                    guest_ip=guest_ip,
                                    guest_port=guest_port,
                                    protocol=proto,
                                )
                            )
                        except (ValueError, IndexError):
                            continue

        return rules


# Convenience functions


def add_port_forward(
    host_port: int,
    guest_ip: str,
    guest_port: int,
    protocol: Protocol = Protocol.TCP,
    comment: str = "",
) -> None:
    """Add a port forwarding rule.

    Convenience function that creates a FirewallManager and adds the rule.
    """
    manager = FirewallManager()
    manager.add_forwarding_rule(
        ForwardingRule(
            host_port=host_port,
            guest_ip=guest_ip,
            guest_port=guest_port,
            protocol=protocol,
            comment=comment,
        )
    )


def remove_port_forward(
    host_port: int,
    guest_ip: str,
    guest_port: int,
    protocol: Protocol = Protocol.TCP,
) -> None:
    """Remove a port forwarding rule.

    Convenience function that creates a FirewallManager and removes the rule.
    """
    manager = FirewallManager()
    manager.remove_forwarding_rule(
        ForwardingRule(
            host_port=host_port,
            guest_ip=guest_ip,
            guest_port=guest_port,
            protocol=protocol,
        )
    )
