"""Abstract nftables ruleset model.

Provides Pydantic models that render to nftables scripts loadable with `nft -f`.
Each container/chain gets its own named table for clean isolation and atomic teardown.
"""

from __future__ import annotations

import subprocess
from enum import Enum

from pydantic import BaseModel, Field


class NftAction(str, Enum):
    """nftables rule action."""

    ACCEPT = "accept"
    DROP = "drop"
    REJECT = "reject"
    MASQUERADE = "masquerade"
    DNAT = "dnat"
    REDIRECT = "redirect"
    RETURN = "return"
    LOG = "log"


class NftHook(str, Enum):
    """nftables chain hook point."""

    INPUT = "input"
    OUTPUT = "output"
    FORWARD = "forward"
    PREROUTING = "prerouting"
    POSTROUTING = "postrouting"


class NftChainType(str, Enum):
    """nftables chain type."""

    FILTER = "filter"
    NAT = "nat"
    ROUTE = "route"


class NftFamily(str, Enum):
    """nftables address family."""

    IP = "ip"
    IP6 = "ip6"
    INET = "inet"


class NftRule(BaseModel):
    """A single nftables rule.

    Rules are rendered in priority order within their chain.
    """

    # Match criteria (all optional — omitted means "match any")
    protocol: str | None = None  # "tcp", "udp", "icmp"
    src_addr: str | None = None  # CIDR or IP
    dst_addr: str | None = None
    src_port: int | None = None
    dst_port: int | None = None
    iif: str | None = None  # incoming interface
    oif: str | None = None  # outgoing interface
    ct_state: str | None = None  # "established,related"

    # Action
    action: NftAction
    dnat_target: str | None = None  # "ip:port" for DNAT
    redirect_port: int | None = None  # port for REDIRECT
    log_prefix: str | None = None  # prefix for LOG

    # Metadata
    comment: str | None = None
    priority: int = 0  # ordering within chain (lower = first)

    def render(self) -> str:
        """Render this rule to nftables syntax."""
        parts: list[str] = []

        if self.iif:
            parts.append(f'iif "{self.iif}"')
        if self.oif:
            parts.append(f'oif "{self.oif}"')
        if self.ct_state:
            parts.append(f"ct state {self.ct_state}")
        if self.protocol:
            parts.append(self.protocol)
        if self.src_addr:
            parts.append(f"ip saddr {self.src_addr}")
        if self.dst_addr:
            parts.append(f"ip daddr {self.dst_addr}")
        if self.src_port:
            parts.append(f"sport {self.src_port}")
        if self.dst_port:
            parts.append(f"dport {self.dst_port}")

        # Action
        if self.action == NftAction.DNAT and self.dnat_target:
            parts.append(f"dnat to {self.dnat_target}")
        elif self.action == NftAction.REDIRECT and self.redirect_port:
            parts.append(f"redirect to :{self.redirect_port}")
        elif self.action == NftAction.LOG and self.log_prefix:
            parts.append(f'log prefix "{self.log_prefix}"')
        else:
            parts.append(self.action.value)

        if self.comment:
            parts.append(f'comment "{self.comment}"')

        return " ".join(parts)


class NftChain(BaseModel):
    """An nftables chain within a table."""

    name: str
    chain_type: NftChainType
    hook: NftHook
    priority: int = 0
    policy: str = "accept"  # default chain policy
    rules: list[NftRule] = Field(default_factory=list)

    def render(self, indent: str = "    ") -> str:
        """Render this chain to nftables syntax."""
        lines = [f"chain {self.name} {{"]
        lines.append(
            f"{indent}type {self.chain_type.value} hook {self.hook.value} "
            f"priority {self.priority}; policy {self.policy};"
        )

        # Sort rules by priority, then render
        sorted_rules = sorted(self.rules, key=lambda r: r.priority)
        if sorted_rules:
            lines.append("")
            for rule in sorted_rules:
                rendered = rule.render()
                if rule.comment:
                    lines.append(f"{indent}{rendered}")
                else:
                    lines.append(f"{indent}{rendered}")

        lines.append("}")
        return "\n".join(lines)


class NftRuleset(BaseModel):
    """Complete nftables ruleset for a container or chain.

    Each ruleset maps to one named table. Tables are isolated from each other
    and can be loaded/deleted atomically.

    Usage:
        ruleset = NftRuleset(table_name="gaol_research")
        ruleset.add_chain(chain)
        script = ruleset.render()  # nftables script
        ruleset.apply()            # atomic load via `nft -f`
        ruleset.destroy()          # delete table instantly
    """

    table_name: str
    family: NftFamily = NftFamily.IP
    chains: list[NftChain] = Field(default_factory=list)

    def add_chain(self, chain: NftChain) -> None:
        """Add a chain to this ruleset."""
        self.chains.append(chain)

    def render(self) -> str:
        """Render to a complete nftables script loadable with `nft -f`."""
        lines = [
            "# Auto-generated by gaol — do not edit manually",
            f"# Table: {self.family.value} {self.table_name}",
            "",
            f"table {self.family.value} {self.table_name} {{",
        ]

        for chain in self.chains:
            rendered = chain.render(indent="        ")
            # Indent the chain block inside the table
            for line in rendered.splitlines():
                lines.append(f"    {line}" if line.strip() else "")

        lines.append("}")
        lines.append("")  # trailing newline
        return "\n".join(lines)

    def apply(self, *, use_sudo: bool = True) -> tuple[int, str, str]:
        """Atomically load this ruleset via `nft -f`.

        First deletes the table if it exists, then loads the new ruleset.
        This ensures clean state without leftover rules.

        Returns:
            (exit_code, stdout, stderr)
        """
        script = self.render()

        # Delete existing table first (ignore error if it doesn't exist)
        delete_cmd = f"table {self.family.value} {self.table_name}"
        cmd_prefix = ["sudo", "-n"] if use_sudo else []
        subprocess.run(
            [*cmd_prefix, "nft", "delete", *delete_cmd.split()],
            capture_output=True,
        )

        # Load new ruleset atomically
        result = subprocess.run(
            [*cmd_prefix, "nft", "-f", "-"],
            input=script,
            capture_output=True,
            text=True,
        )
        return result.returncode, result.stdout, result.stderr

    def destroy(self, *, use_sudo: bool = True) -> tuple[int, str, str]:
        """Delete this table (removes all chains and rules instantly).

        Returns:
            (exit_code, stdout, stderr)
        """
        cmd_prefix = ["sudo", "-n"] if use_sudo else []
        result = subprocess.run(
            [*cmd_prefix, "nft", "delete", "table", self.family.value, self.table_name],
            capture_output=True,
            text=True,
        )
        return result.returncode, result.stdout, result.stderr


# ---------------------------------------------------------------------------
# Archetype rule generators
# ---------------------------------------------------------------------------


def kill_switch_rules(
    gateway_ip: str,
    allowed_ips: list[str] | None = None,
) -> list[NftRule]:
    """Generate kill-switch rules (OUTPUT DROP with exceptions).

    Args:
        gateway_ip: The gateway IP to allow (veth peer on host side).
        allowed_ips: Additional IPs to allow (e.g., proxy, DNS).

    Returns:
        List of rules for an OUTPUT chain with DROP policy.
    """
    rules = [
        NftRule(oif="lo", action=NftAction.ACCEPT, priority=10, comment="allow loopback"),
        NftRule(
            dst_addr=gateway_ip,
            action=NftAction.ACCEPT,
            priority=20,
            comment="allow gateway",
        ),
        NftRule(
            ct_state="established,related",
            action=NftAction.ACCEPT,
            priority=30,
            comment="allow established",
        ),
    ]
    for i, ip in enumerate(allowed_ips or []):
        rules.append(
            NftRule(
                dst_addr=ip,
                action=NftAction.ACCEPT,
                priority=40 + i,
                comment=f"allow {ip}",
            )
        )
    return rules


def chain_forwarding_ruleset(
    table_name: str,
    container_iface: str,
    container_subnet: str,
) -> NftRuleset:
    """Generate a forwarding + masquerade ruleset for a chain hop.

    Args:
        table_name: nftables table name (e.g., "gaol_research").
        container_iface: Host-side veth interface name (e.g., "ve-research").
        container_subnet: Container subnet CIDR (e.g., "10.200.1.0/30").

    Returns:
        Complete ruleset with forward and postrouting chains.
    """
    ruleset = NftRuleset(table_name=table_name)

    # Forward chain: allow container traffic out, established back
    forward = NftChain(
        name="forward",
        chain_type=NftChainType.FILTER,
        hook=NftHook.FORWARD,
        priority=0,
        policy="drop",
        rules=[
            NftRule(
                iif=container_iface,
                action=NftAction.ACCEPT,
                priority=10,
                comment="container outbound",
            ),
            NftRule(
                oif=container_iface,
                ct_state="established,related",
                action=NftAction.ACCEPT,
                priority=20,
                comment="established inbound",
            ),
        ],
    )
    ruleset.add_chain(forward)

    # NAT chain: masquerade container traffic
    postrouting = NftChain(
        name="postrouting",
        chain_type=NftChainType.NAT,
        hook=NftHook.POSTROUTING,
        priority=100,
        rules=[
            NftRule(
                src_addr=container_subnet,
                action=NftAction.MASQUERADE,
                priority=10,
                comment="NAT container traffic",
            ),
        ],
    )
    ruleset.add_chain(postrouting)

    return ruleset


def kill_switch_ruleset(
    table_name: str,
    gateway_ip: str,
    allowed_ips: list[str] | None = None,
) -> NftRuleset:
    """Generate a complete kill-switch ruleset for a container.

    Args:
        table_name: nftables table name (e.g., "gaol_research_ks").
        gateway_ip: Gateway IP to allow.
        allowed_ips: Additional allowed IPs.

    Returns:
        Complete ruleset with OUTPUT DROP policy and exceptions.
    """
    ruleset = NftRuleset(table_name=table_name)

    output = NftChain(
        name="output",
        chain_type=NftChainType.FILTER,
        hook=NftHook.OUTPUT,
        priority=0,
        policy="drop",
        rules=kill_switch_rules(gateway_ip, allowed_ips),
    )
    ruleset.add_chain(output)

    return ruleset
