"""Firewall rule management for privacy networking.

Manages iptables/nftables rules for:
- Transparent Tor proxy (redirect all traffic through Tor)
- DNS leak prevention (block unauthorized DNS queries)
- VPN kill switch (block traffic if VPN disconnects)
"""

from __future__ import annotations

import shutil
import subprocess
from enum import Enum

from gaol.privacy.exceptions import (
    FirewallNotAvailableError,
    FirewallRuleError,
    InsufficientPermissionsError,
)


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

    IPTABLES = "iptables"
    NFTABLES = "nftables"


# Chain/table names for gaol privacy rules
PRIVACY_CHAIN = "GAOL_PRIVACY"
PRIVACY_NFT_TABLE = "gaol_privacy"


def _run_cmd(
    cmd: list[str], use_sudo: bool = True, check: bool = False
) -> subprocess.CompletedProcess[str]:
    """Run a command with optional sudo.

    Args:
        cmd: Command to run
        use_sudo: Whether to use sudo
        check: Whether to raise on non-zero exit

    Returns:
        Completed process result
    """
    if use_sudo:
        cmd = ["sudo"] + cmd
    return subprocess.run(cmd, capture_output=True, text=True, check=check)


class PrivacyFirewall:
    """Manages firewall rules for privacy networking."""

    def __init__(self, use_sudo: bool = True) -> None:
        """Initialize privacy firewall manager.

        Args:
            use_sudo: Whether to use sudo for commands
        """
        self.use_sudo = use_sudo
        self._backend: FirewallBackend | None = None
        self._setup_done = False

    @property
    def backend(self) -> FirewallBackend:
        """Get the firewall backend, detecting if needed."""
        if self._backend is None:
            self._backend = self._detect_backend()
        return self._backend

    def _detect_backend(self) -> FirewallBackend:
        """Detect available firewall backend.

        Returns:
            Available FirewallBackend

        Raises:
            FirewallNotAvailableError: If no backend available
        """
        # Prefer nftables if available and functional
        if shutil.which("nft"):
            result = _run_cmd(["nft", "list", "tables"], self.use_sudo)
            if result.returncode == 0:
                return FirewallBackend.NFTABLES

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

        raise FirewallNotAvailableError()

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

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

        try:
            if self.backend == FirewallBackend.NFTABLES:
                self._setup_nftables()
            else:
                self._setup_iptables()
            self._setup_done = True
        except subprocess.CalledProcessError as e:
            if "Permission denied" in str(e.stderr):
                raise InsufficientPermissionsError("firewall setup")
            raise FirewallRuleError("setup", str(e.stderr))

    def _setup_nftables(self) -> None:
        """Set up nftables table and chains for privacy rules."""
        # Create table
        _run_cmd(["nft", "add", "table", "ip", PRIVACY_NFT_TABLE], self.use_sudo)

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

        # Create filter chain for blocking
        _run_cmd(
            [
                "nft",
                "add",
                "chain",
                "ip",
                PRIVACY_NFT_TABLE,
                "filter_output",
                "{ type filter hook output priority 0 ; }",
            ],
            self.use_sudo,
        )

    def _setup_iptables(self) -> None:
        """Set up iptables chains for privacy rules."""
        # Create custom chain in nat table
        result = _run_cmd(
            ["iptables", "-t", "nat", "-N", PRIVACY_CHAIN], self.use_sudo
        )
        if result.returncode == 0:
            # Chain created, add jump from OUTPUT
            _run_cmd(
                ["iptables", "-t", "nat", "-A", "OUTPUT", "-j", PRIVACY_CHAIN],
                self.use_sudo,
            )

        # Create custom chain in filter table
        result = _run_cmd(["iptables", "-N", f"{PRIVACY_CHAIN}_FILTER"], self.use_sudo)
        if result.returncode == 0:
            _run_cmd(
                ["iptables", "-A", "OUTPUT", "-j", f"{PRIVACY_CHAIN}_FILTER"],
                self.use_sudo,
            )

    def setup_transparent_tor(
        self,
        tor_user: str = "debian-tor",
        trans_port: int = 9040,
        dns_port: int = 5353,
        non_tor_net: list[str] | None = None,
    ) -> None:
        """Set up iptables rules for transparent Tor proxy.

        All TCP traffic is redirected to Tor's TransPort.
        All DNS (UDP 53) is redirected to Tor's DNSPort.
        Tor daemon's own traffic is allowed to bypass.

        Args:
            tor_user: Username running the Tor daemon
            trans_port: Tor transparent proxy port
            dns_port: Tor DNS port
            non_tor_net: Networks to exclude from Tor (e.g., local networks)
        """
        self.setup()

        if non_tor_net is None:
            non_tor_net = [
                "127.0.0.0/8",
                "10.0.0.0/8",
                "172.16.0.0/12",
                "192.168.0.0/16",
            ]

        if self.backend == FirewallBackend.NFTABLES:
            self._setup_transparent_tor_nft(tor_user, trans_port, dns_port, non_tor_net)
        else:
            self._setup_transparent_tor_ipt(tor_user, trans_port, dns_port, non_tor_net)

    def _setup_transparent_tor_nft(
        self,
        tor_user: str,
        trans_port: int,
        dns_port: int,
        non_tor_net: list[str],
    ) -> None:
        """Set up transparent Tor using nftables."""
        # Get Tor user UID
        result = _run_cmd(["id", "-u", tor_user], use_sudo=False)
        if result.returncode != 0:
            raise FirewallRuleError("transparent tor", f"User {tor_user} not found")
        tor_uid = result.stdout.strip()

        # Allow Tor daemon's own traffic
        _run_cmd(
            [
                "nft",
                "add",
                "rule",
                "ip",
                PRIVACY_NFT_TABLE,
                "output",
                "meta",
                "skuid",
                tor_uid,
                "return",
            ],
            self.use_sudo,
        )

        # Skip local networks
        for net in non_tor_net:
            _run_cmd(
                [
                    "nft",
                    "add",
                    "rule",
                    "ip",
                    PRIVACY_NFT_TABLE,
                    "output",
                    "ip",
                    "daddr",
                    net,
                    "return",
                ],
                self.use_sudo,
            )

        # Redirect DNS to Tor DNS port
        _run_cmd(
            [
                "nft",
                "add",
                "rule",
                "ip",
                PRIVACY_NFT_TABLE,
                "output",
                "udp",
                "dport",
                "53",
                "redirect",
                "to",
                f":{dns_port}",
            ],
            self.use_sudo,
        )

        # Redirect all TCP to TransPort
        _run_cmd(
            [
                "nft",
                "add",
                "rule",
                "ip",
                PRIVACY_NFT_TABLE,
                "output",
                "tcp",
                "flags",
                "&",
                "syn",
                "==",
                "syn",
                "redirect",
                "to",
                f":{trans_port}",
            ],
            self.use_sudo,
        )

    def _setup_transparent_tor_ipt(
        self,
        tor_user: str,
        trans_port: int,
        dns_port: int,
        non_tor_net: list[str],
    ) -> None:
        """Set up transparent Tor using iptables."""
        # Allow Tor daemon's own traffic
        _run_cmd(
            [
                "iptables",
                "-t",
                "nat",
                "-A",
                PRIVACY_CHAIN,
                "-m",
                "owner",
                "--uid-owner",
                tor_user,
                "-j",
                "RETURN",
            ],
            self.use_sudo,
        )

        # Skip local networks
        for net in non_tor_net:
            _run_cmd(
                [
                    "iptables",
                    "-t",
                    "nat",
                    "-A",
                    PRIVACY_CHAIN,
                    "-d",
                    net,
                    "-j",
                    "RETURN",
                ],
                self.use_sudo,
            )

        # Redirect DNS (UDP 53) to Tor DNS port
        _run_cmd(
            [
                "iptables",
                "-t",
                "nat",
                "-A",
                PRIVACY_CHAIN,
                "-p",
                "udp",
                "--dport",
                "53",
                "-j",
                "REDIRECT",
                "--to-ports",
                str(dns_port),
            ],
            self.use_sudo,
        )

        # Redirect all TCP SYN packets to TransPort
        _run_cmd(
            [
                "iptables",
                "-t",
                "nat",
                "-A",
                PRIVACY_CHAIN,
                "-p",
                "tcp",
                "--syn",
                "-j",
                "REDIRECT",
                "--to-ports",
                str(trans_port),
            ],
            self.use_sudo,
        )

    def setup_dns_leak_prevention(
        self,
        allowed_dns: list[str] | None = None,
        allow_local: bool = False,
    ) -> None:
        """Set up DNS leak prevention rules.

        Blocks all DNS queries (UDP/TCP port 53) except to allowed servers.

        Args:
            allowed_dns: List of allowed DNS server IPs (empty = block all)
            allow_local: Whether to allow local resolvers (127.0.0.1)
        """
        self.setup()

        if allowed_dns is None:
            allowed_dns = []

        if allow_local:
            allowed_dns = ["127.0.0.1", "127.0.0.53"] + allowed_dns

        if self.backend == FirewallBackend.NFTABLES:
            self._setup_dns_leak_nft(allowed_dns)
        else:
            self._setup_dns_leak_ipt(allowed_dns)

    def _setup_dns_leak_nft(self, allowed_dns: list[str]) -> None:
        """Set up DNS leak prevention using nftables."""
        # Allow DNS to specified servers
        for dns in allowed_dns:
            _run_cmd(
                [
                    "nft",
                    "add",
                    "rule",
                    "ip",
                    PRIVACY_NFT_TABLE,
                    "filter_output",
                    "ip",
                    "daddr",
                    dns,
                    "udp",
                    "dport",
                    "53",
                    "accept",
                ],
                self.use_sudo,
            )
            _run_cmd(
                [
                    "nft",
                    "add",
                    "rule",
                    "ip",
                    PRIVACY_NFT_TABLE,
                    "filter_output",
                    "ip",
                    "daddr",
                    dns,
                    "tcp",
                    "dport",
                    "53",
                    "accept",
                ],
                self.use_sudo,
            )

        # Block all other DNS
        _run_cmd(
            [
                "nft",
                "add",
                "rule",
                "ip",
                PRIVACY_NFT_TABLE,
                "filter_output",
                "udp",
                "dport",
                "53",
                "drop",
            ],
            self.use_sudo,
        )
        _run_cmd(
            [
                "nft",
                "add",
                "rule",
                "ip",
                PRIVACY_NFT_TABLE,
                "filter_output",
                "tcp",
                "dport",
                "53",
                "drop",
            ],
            self.use_sudo,
        )

    def _setup_dns_leak_ipt(self, allowed_dns: list[str]) -> None:
        """Set up DNS leak prevention using iptables."""
        # Allow DNS to specified servers
        for dns in allowed_dns:
            _run_cmd(
                [
                    "iptables",
                    "-A",
                    f"{PRIVACY_CHAIN}_FILTER",
                    "-d",
                    dns,
                    "-p",
                    "udp",
                    "--dport",
                    "53",
                    "-j",
                    "ACCEPT",
                ],
                self.use_sudo,
            )
            _run_cmd(
                [
                    "iptables",
                    "-A",
                    f"{PRIVACY_CHAIN}_FILTER",
                    "-d",
                    dns,
                    "-p",
                    "tcp",
                    "--dport",
                    "53",
                    "-j",
                    "ACCEPT",
                ],
                self.use_sudo,
            )

        # Block all other DNS
        _run_cmd(
            [
                "iptables",
                "-A",
                f"{PRIVACY_CHAIN}_FILTER",
                "-p",
                "udp",
                "--dport",
                "53",
                "-j",
                "DROP",
            ],
            self.use_sudo,
        )
        _run_cmd(
            [
                "iptables",
                "-A",
                f"{PRIVACY_CHAIN}_FILTER",
                "-p",
                "tcp",
                "--dport",
                "53",
                "-j",
                "DROP",
            ],
            self.use_sudo,
        )

    def setup_kill_switch(
        self,
        vpn_interface: str,
        allowed_hosts: list[str] | None = None,
    ) -> None:
        """Set up VPN kill switch.

        Blocks all traffic that doesn't go through the VPN interface.
        Only allows traffic to specified hosts (e.g., VPN servers).

        Args:
            vpn_interface: VPN interface name (e.g., tun0, wg0)
            allowed_hosts: Hosts to allow even without VPN (e.g., VPN server IPs)
        """
        self.setup()

        if allowed_hosts is None:
            allowed_hosts = []

        # Always allow loopback
        allowed_hosts = ["127.0.0.1"] + allowed_hosts

        if self.backend == FirewallBackend.NFTABLES:
            self._setup_kill_switch_nft(vpn_interface, allowed_hosts)
        else:
            self._setup_kill_switch_ipt(vpn_interface, allowed_hosts)

    def _setup_kill_switch_nft(
        self, vpn_interface: str, allowed_hosts: list[str]
    ) -> None:
        """Set up VPN kill switch using nftables."""
        # Allow traffic on VPN interface
        _run_cmd(
            [
                "nft",
                "add",
                "rule",
                "ip",
                PRIVACY_NFT_TABLE,
                "filter_output",
                "oif",
                vpn_interface,
                "accept",
            ],
            self.use_sudo,
        )

        # Allow loopback
        _run_cmd(
            [
                "nft",
                "add",
                "rule",
                "ip",
                PRIVACY_NFT_TABLE,
                "filter_output",
                "oif",
                "lo",
                "accept",
            ],
            self.use_sudo,
        )

        # Allow specific hosts (e.g., VPN servers)
        for host in allowed_hosts:
            _run_cmd(
                [
                    "nft",
                    "add",
                    "rule",
                    "ip",
                    PRIVACY_NFT_TABLE,
                    "filter_output",
                    "ip",
                    "daddr",
                    host,
                    "accept",
                ],
                self.use_sudo,
            )

        # Block everything else
        _run_cmd(
            [
                "nft",
                "add",
                "rule",
                "ip",
                PRIVACY_NFT_TABLE,
                "filter_output",
                "drop",
            ],
            self.use_sudo,
        )

    def _setup_kill_switch_ipt(
        self, vpn_interface: str, allowed_hosts: list[str]
    ) -> None:
        """Set up VPN kill switch using iptables."""
        # Allow traffic on VPN interface
        _run_cmd(
            [
                "iptables",
                "-A",
                f"{PRIVACY_CHAIN}_FILTER",
                "-o",
                vpn_interface,
                "-j",
                "ACCEPT",
            ],
            self.use_sudo,
        )

        # Allow loopback
        _run_cmd(
            [
                "iptables",
                "-A",
                f"{PRIVACY_CHAIN}_FILTER",
                "-o",
                "lo",
                "-j",
                "ACCEPT",
            ],
            self.use_sudo,
        )

        # Allow specific hosts
        for host in allowed_hosts:
            _run_cmd(
                [
                    "iptables",
                    "-A",
                    f"{PRIVACY_CHAIN}_FILTER",
                    "-d",
                    host,
                    "-j",
                    "ACCEPT",
                ],
                self.use_sudo,
            )

        # Block everything else
        _run_cmd(
            [
                "iptables",
                "-A",
                f"{PRIVACY_CHAIN}_FILTER",
                "-j",
                "DROP",
            ],
            self.use_sudo,
        )

    def cleanup(self) -> None:
        """Remove all privacy firewall rules.

        Flushes and removes the privacy chains/tables.
        """
        if self._backend is None:
            # Try to detect to clean up
            try:
                _ = self.backend
            except FirewallNotAvailableError:
                return

        if self.backend == FirewallBackend.NFTABLES:
            self._cleanup_nftables()
        else:
            self._cleanup_iptables()

        self._setup_done = False

    def _cleanup_nftables(self) -> None:
        """Remove nftables privacy table."""
        _run_cmd(
            ["nft", "delete", "table", "ip", PRIVACY_NFT_TABLE],
            self.use_sudo,
        )

    def _cleanup_iptables(self) -> None:
        """Remove iptables privacy chains."""
        # Remove jump rules first
        _run_cmd(
            ["iptables", "-t", "nat", "-D", "OUTPUT", "-j", PRIVACY_CHAIN],
            self.use_sudo,
        )
        _run_cmd(
            ["iptables", "-D", "OUTPUT", "-j", f"{PRIVACY_CHAIN}_FILTER"],
            self.use_sudo,
        )

        # Flush and delete chains
        _run_cmd(
            ["iptables", "-t", "nat", "-F", PRIVACY_CHAIN],
            self.use_sudo,
        )
        _run_cmd(
            ["iptables", "-t", "nat", "-X", PRIVACY_CHAIN],
            self.use_sudo,
        )
        _run_cmd(
            ["iptables", "-F", f"{PRIVACY_CHAIN}_FILTER"],
            self.use_sudo,
        )
        _run_cmd(
            ["iptables", "-X", f"{PRIVACY_CHAIN}_FILTER"],
            self.use_sudo,
        )

    def is_available(self) -> bool:
        """Check if firewall backend is available.

        Returns:
            True if a backend is available
        """
        try:
            _ = self.backend
            return True
        except FirewallNotAvailableError:
            return False
