"""Private IPv4 service filtering, not host-port forwarding or publishing.

The lifecycle owner must serialize mutations, keep a live VM's identity immutable,
install the policy before starting its guest, and quarantine before teardown.
Only bridge/TAP traffic is covered; do not offload it around the bridge hooks.
Before making a TAP usable, attach it to an STP-disabled software bridge and set
the TAP's bridge_slave group_fwd_mask bit 0x4000. Otherwise Linux can deliver LLDP
(and STP) before *any* bridge hook and learn spoofed MACs. Network/lifecycle owns
these link settings; this manager only verifies them when the TAP exists.
No br_netfilter sysctls, NAT, persistent applied flags, or live conntrack flushes
are used. Native bridge conntrack (Linux >= 5.3) and nft JSON support are required.
"""

from __future__ import annotations

import hashlib
import json
import re
import uuid
from ipaddress import IPv4Address, IPv4Network
from types import SimpleNamespace
from typing import Any

from .config import Settings
from .domain import PortRule, VmRecord
from .errors import UvmError, ValidationError
from .system import CommandResult, CommandRunner


_COMMAND_TIMEOUT_S = 5.0
_PRIVATE_NETWORKS = tuple(IPv4Network(value) for value in (
    "10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16",
))


def _match(left: dict[str, Any], right: Any, op: str = "==") -> dict[str, Any]:
    return {"match": {"op": op, "left": left, "right": right}}


def _payload(protocol: str, field: str) -> dict[str, Any]:
    return {"payload": {"protocol": protocol, "field": field}}


def _ct(key: str, direction: str | None = None) -> dict[str, Any]:
    # nft 1.0.9 parser_json.c/ct.c use the literal keys "ip saddr"/"ip daddr".
    # Its CT parser ignores the "family" property advertised by the JSON manual.
    return {"ct": {"key": key, **({"dir": direction} if direction else {})}}


class FirewallManager:
    """Own one table per (resolved settings.base, VM id) in the current netns.

    ``ports=None`` is legacy/unmanaged, not an empty allowlist. Normal apply
    refuses it; quarantine works even without a port policy or valid IP/MAC.
    Errors leave existing rules in place (or the committed replacement on a
    readback/timeout failure). Callers must not start/release a VM on an error.
    """

    def __init__(self, settings: Settings, runner: CommandRunner) -> None:
        self._settings = settings
        self._runner = runner

    def apply(self, vm: VmRecord, *, quarantine: bool = False) -> None:
        """Atomically install/replace only this VM's owned table; retain CT state."""
        expected = self._policy(vm, quarantine=quarantine)
        if not self._bridge_paths_safe(vm):
            raise UvmError(
                "firewall needs an STP-disabled software bridge and TAP bridge_slave "
                "group_fwd_mask bit 0x4000 (LLDP forwarding) before enforcement"
            )
        table = expected[0]["table"]
        current = self._read_table(table["name"])
        commands = []
        if current is not None:
            commands.append(self._delete_owned(current, table))
        # Exclusive create plus deletion by handle also protects against a table
        # being replaced by a foreign owner between inspection and commit.
        commands.append({"create": expected[0]})
        commands.extend({"add": item} for item in expected[1:])
        self._transaction(commands)
        installed = self._read_table(table["name"])
        if installed is None or _normalize(installed) != _normalize(expected):
            raise UvmError("firewall readback differs from the requested policy; keep VM isolated")

    def remove(self, vm: VmRecord) -> None:
        """Remove an owned table only after the TAP is gone; absence is harmless."""
        self._require_offline(vm)
        table = self._table_identity(vm)
        current = self._read_table(table["name"])
        if current is not None:
            self._transaction([self._delete_owned(current, table)])

    def status(self, vm: VmRecord) -> str:
        """Inspect actual hooks, ordered rules, identity and complete port set.

        Missing legacy tables are unmanaged. Foreign or altered tables are
        drift; failed inspection/invalid metadata is unknown. A legacy VM can
        report quarantined. This is a ruleset snapshot, not a reachability test.
        """
        try:
            table = self._table_identity(vm)
            current = self._read_table(table["name"])
            if current is None:
                return "unmanaged" if getattr(vm, "ports", None) is None else "missing"
            if _normalize(current) == _normalize(self._policy(vm, quarantine=True)):
                return "quarantined" if self._bridge_paths_safe(vm) else "drift"
            if getattr(vm, "ports", None) is not None:
                if _normalize(current) == _normalize(self._policy(vm, quarantine=False)):
                    return "enforced" if self._bridge_paths_safe(vm) else "drift"
            return "drift"
        except (UvmError, ValueError, TypeError, KeyError):
            return "unknown"

    def check_available(self) -> None:
        """Check cleanup tooling, kernel and a full non-committing nft policy."""
        self._run(("conntrack", "--version"))
        release = self._run(("uname", "-r")).stdout.strip()
        version = re.match(r"^(\d+)\.(\d+)", release)
        if version is None or tuple(map(int, version.groups())) < (5, 3):
            raise UvmError("private ports require Linux >= 5.3 with native bridge conntrack")
        probe = SimpleNamespace(
            id=f"probe-{uuid.uuid4().hex}", tap=f"uvm{uuid.uuid4().hex[:11]}",
            guest_ip="10.0.0.2", gateway="10.0.0.1",
            network="10.0.0.0/24", mac="02:00:00:00:00:02",
            ports=(PortRule(443),),
        )
        policy = self._policy(probe, quarantine=False)  # type: ignore[arg-type]
        commands = [{"create": policy[0]}, *({"add": item} for item in policy[1:])]
        try:
            self._transaction(commands, check_only=True)
        except UvmError as error:
            raise UvmError(f"nft bridge/conntrack policy is unavailable: {error}") from error

    def cleanup_connections(self, vm: VmRecord) -> None:
        """Offline, IP-scoped CT deletion before address reuse, never for close.

        All four tuple addresses are checked, including NAT reply tuples. A
        failed deletion is tolerated only if a successful list proves absence.
        The caller must retain the address reservation until this succeeds.
        """
        self._require_offline(vm)
        address = self._guest_ip(vm)
        for field in ("--orig-src", "--orig-dst", "--reply-src", "--reply-dst"):
            selector = ("--family", "ipv4", field, address)
            self._run(("conntrack", "--delete", *selector), check=False)
            remaining = self._run(("conntrack", "--list", *selector))
            if remaining.stdout.strip():
                raise UvmError(f"conntrack entries remain for {address}; do not reuse its IP")

    def _run(self, command: tuple[str, ...], **kwargs: Any) -> CommandResult:
        return self._runner.run(
            command, capture=True, timeout=_COMMAND_TIMEOUT_S, **kwargs,
        )

    def _transaction(self, commands: list[dict[str, Any]], *, check_only: bool = False) -> None:
        args = ("nft", "--json", *(("--check",) if check_only else ()), "-f", "-")
        self._run(args, input_text=json.dumps({"nftables": commands}))

    def _table_identity(self, vm: VmRecord) -> dict[str, Any]:
        if not isinstance(vm.id, str) or not vm.id:
            raise ValidationError("firewall requires a nonempty VM ID")
        identity = json.dumps([str(self._settings.base.resolve()), vm.id]).encode("utf-8")
        digest = hashlib.sha256(identity).hexdigest()
        return {"family": "bridge", "name": f"uvm_fw_{digest[:24]}",
                "comment": f"uvm-private-ports:v1:{digest}"}

    def _tap(self, vm: VmRecord) -> str:
        if not isinstance(vm.tap, str) or re.fullmatch(r"[A-Za-z0-9_][A-Za-z0-9_.-]{0,14}", vm.tap) is None:
            raise ValidationError("firewall requires an exact, safe TAP name (1-15 characters)")
        if vm.tap in {self._settings.bridge, getattr(vm, "bridge", None)}:
            raise ValidationError("firewall TAP must not be the bridge itself")
        return vm.tap

    def _guest_ip(self, vm: VmRecord) -> str:
        try:
            address = IPv4Address(vm.guest_ip)
            network = IPv4Network(getattr(vm, "network", None) or self._settings.network)
        except (ValueError, TypeError) as error:
            raise ValidationError("firewall requires a valid guest IPv4 address/network") from error
        gateway = str(getattr(vm, "gateway", None) or self._settings.gateway)
        if (not any(address in private for private in _PRIVATE_NETWORKS)
                or address not in network
                or address in {network.network_address, network.broadcast_address}
                or str(address) == gateway):
            raise ValidationError("firewall requires a usable private guest IP inside its network")
        return str(address)

    def _links(self) -> list[dict[str, Any]]:
        try:
            links = json.loads(self._run(("ip", "-j", "-d", "link", "show")).stdout)
            if not isinstance(links, list) or any(
                not isinstance(link, dict) or not isinstance(link.get("ifname"), str)
                for link in links
            ):
                raise ValueError("invalid link listing")
        except (ValueError, TypeError) as error:
            raise UvmError("cannot inspect network interfaces; retaining firewall/conntrack state") from error
        return links

    def _bridge_paths_safe(self, vm: VmRecord) -> bool:
        links = self._links()
        name = self._tap(vm)
        taps = [link for link in links if link["ifname"] == name]
        if not taps:
            return True  # Policy may precede TAP creation; main must prepare it before bringing it up.
        if len(taps) != 1:
            return False
        tap = taps[0]
        bridges = [link for link in links if tap.get("master") in (link["ifname"], link.get("ifindex"))]
        if len(bridges) != 1:
            return False
        try:
            port = tap["linkinfo"]
            bridge = bridges[0]["linkinfo"]
            mask = int(str(port["info_slave_data"]["group_fwd_mask"]), 0)
            return (port["info_slave_kind"] == "bridge" and bridge["info_kind"] == "bridge"
                    and bridge["info_data"]["stp_state"] == 0
                    and 0 <= mask <= 65535 and bool(mask & 0x4000))
        except (KeyError, TypeError, ValueError):
            return False

    def _require_offline(self, vm: VmRecord) -> None:
        tap = self._tap(vm)
        if any(link["ifname"] == tap for link in self._links()):
            raise UvmError(f"TAP {tap} still exists; quarantine it and remove the TAP first")

    def _read_table(self, name: str) -> list[dict[str, Any]] | None:
        # A successful table inventory, not a localized 'not found' error,
        # distinguishes missing policy from missing privileges/tools.
        def objects(command: tuple[str, ...]) -> list[dict[str, Any]]:
            try:
                data = json.loads(self._run(command).stdout)["nftables"]
                if not isinstance(data, list) or any(
                    not isinstance(item, dict) or len(item) != 1
                    or not isinstance(next(iter(item.values())), dict) for item in data
                ):
                    raise ValueError("invalid nft object list")
                return [item for item in data if "metainfo" not in item]
            except (ValueError, KeyError, TypeError) as error:
                raise UvmError("could not inspect nft bridge tables") from error

        tables = objects(("nft", "--json", "list", "tables", "bridge"))
        if any("table" not in item or item["table"].get("family") != "bridge"
               or not isinstance(item["table"].get("name"), str) for item in tables):
            raise UvmError("invalid nft bridge table inventory")
        matches = [item for item in tables if item["table"].get("name") == name]
        if not matches:
            return None
        if len(matches) != 1:
            raise UvmError("ambiguous nft bridge table inventory")
        data = objects(("nft", "--json", "--numeric-protocol", "list", "table", "bridge", name))
        headers = [item["table"] for item in data if "table" in item]
        if len(headers) != 1 or headers[0].get("name") != name or headers[0].get("family") != "bridge":
            raise UvmError("invalid nft bridge table snapshot")
        return data

    @staticmethod
    def _delete_owned(current: list[dict[str, Any]], table: dict[str, Any]) -> dict[str, Any]:
        header = next(item["table"] for item in current if "table" in item)
        if any(header.get(key) != value for key, value in table.items()):
            raise UvmError(f"refusing to touch foreign nft table {table['name']}")
        handle = header.get("handle")
        if type(handle) is not int or handle <= 0:
            raise UvmError("owned nft table has no valid handle; refusing unsafe deletion")
        return {"delete": {"table": {"family": "bridge", "handle": handle}}}

    def _policy(self, vm: VmRecord, *, quarantine: bool) -> list[dict[str, Any]]:
        ports = getattr(vm, "ports", None)
        if ports is None and not quarantine:
            raise ValidationError("legacy VM has no managed port policy; only quarantine is allowed")
        table = self._table_identity(vm)
        tap = self._tap(vm)
        scope = {"family": "bridge", "table": table["name"]}
        policy: list[dict[str, Any]] = [{"table": table}]

        def chain(name: str, hook: str | None = None, priority: int = 0) -> None:
            definition = {**scope, "name": name}
            if hook is not None:
                definition.update(type="filter", hook=hook, prio=priority, policy="accept")
            policy.append({"chain": definition})

        def rule(name: str, *expressions: dict[str, Any]) -> None:
            policy.append({"rule": {**scope, "chain": name, "expr": list(expressions)}})

        for hook in ("prerouting", "input", "forward", "output"):
            chain(hook, hook, -310 if hook == "prerouting" else 0)
        if not quarantine:
            for name in ("identity", "from_vm", "to_vm"):
                chain(name)
        for hook, interface, target in (
            ("prerouting", "iifname", "identity"),
            ("input", "iifname", "from_vm"),
            ("forward", "iifname", "from_vm"),
            ("forward", "oifname", "to_vm"),
            ("output", "oifname", "to_vm"),
        ):
            rule(hook, _match({"meta": {"key": interface}}, tap),
                 {"drop": None} if quarantine else {"jump": {"target": target}})
        if quarantine:
            return policy

        address = self._guest_ip(vm)
        mac = vm.mac.lower() if isinstance(vm.mac, str) else ""
        if (re.fullmatch(r"(?:[0-9a-f]{2}:){5}[0-9a-f]{2}", mac) is None
                or int(mac[:2], 16) & 1 or mac == "00:00:00:00:00:00"):
            raise ValidationError("firewall requires a nonzero unicast guest MAC")
        if not isinstance(ports, tuple):
            raise ValidationError("managed ports must be a tuple of TCP/HTTP/HTTPS port rules")
        allowed = set()
        for port_rule in ports:
            port = getattr(port_rule, "port", None)
            if (type(port) is not int or not 1 <= port <= 65535
                    or getattr(port_rule, "protocol", None) not in ("TCP", "HTTP", "HTTPS")):
                raise ValidationError("private ports must be TCP/HTTP/HTTPS integers from 1 to 65535")
            allowed.add(port)
        # json_parse_set_expr() in nft 1.0.9 rejects [], even for an empty named
        # set. Omit elem, as set_print_json() does when listing an empty set.
        policy.insert(1, {"set": {**scope, "name": "tcp_ports", "type": "inet_service",
                                 "flags": ["constant"],
                                 **({"elem": sorted(allowed)} if allowed else {})}})

        drop, success = {"drop": None}, {"return": None}
        ipv4 = _match(_payload("ether", "type"), "ip")
        arp = _match(_payload("ether", "type"), "arp")
        arp_header = [arp, _match(_payload("arp", "htype"), 1),
                      _match(_payload("arp", "ptype"), "ip"),
                      _match(_payload("arp", "hlen"), 6), _match(_payload("arp", "plen"), 4)]
        # Before conntrack (-200) and FDB learning: no MAC/IP/ARP spoofing,
        # IPv6, VLAN encapsulation or other EtherTypes from a managed port.
        rule("identity", _match(_payload("ether", "saddr"), mac, "!="), drop)
        rule("identity", ipv4, _match(_payload("ip", "saddr"), address), success)
        rule("identity", *arp_header,
             _match(_payload("arp", "operation"), {"set": ["request", "reply"]}),
             _match(_payload("arp", "saddr ether"), mac),
             _match(_payload("arp", "saddr ip"), address), success)
        # RFC 5227 address probes may use zero, but only probe this VM's IP.
        rule("identity", *arp_header, _match(_payload("arp", "operation"), "request"),
             _match(_payload("arp", "saddr ether"), mac),
             _match(_payload("arp", "saddr ip"), "0.0.0.0"),
             _match(_payload("arp", "daddr ip"), address), success)
        rule("identity", drop)

        # Some link-local bridge frames take input without prerouting. Recheck
        # identity there too, before its local-finish callback can learn a MAC.
        rule("from_vm", {"jump": {"target": "identity"}})
        rule("from_vm", arp, success)
        rule("to_vm", *arp_header, _match(_payload("arp", "daddr ip"), address),
             _match(_payload("arp", "operation"), "request"), success)
        rule("to_vm", *arp_header, _match(_payload("arp", "daddr ip"), address),
             _match(_payload("arp", "daddr ether"), mac),
             _match(_payload("arp", "operation"), "reply"), success)

        for name, outbound in (("from_vm", True), ("to_vm", False)):
            rule(name, _match(_payload("ether", "type"), "ip", "!="), drop)
            if not outbound:
                rule(name, ipv4, _match(_payload("ip", "daddr"), address, "!="), drop)
                rule(name, _match(_payload("ether", "daddr"), mac, "!="), drop)
            # Loose TCP tracking must not reclassify a server's midstream data
            # as a newly guest-initiated connection after tracking loss.
            rule(name, ipv4, _match(_payload("ip", "protocol"), "tcp"),
                 _match(_ct("state"), "new", "in"),
                 _match({"&": [_payload("tcp", "flags"), 23]}, 2, "!="), drop)

            for client in (True, False):
                original = outbound == client
                flow = [_match(_ct("direction"), "original" if original else "reply"),
                        _match(_ct("ip saddr" if client else "ip daddr", "original"), address)]
                states = ["new", "established"] if original else ["established"]
                state = _match(_ct("state"), states[0] if len(states) == 1 else states, "in")
                for protocol in (("tcp", "udp", "icmp") if client else ("tcp", "icmp")):
                    parent = [*flow, _match(_ct("protocol"), protocol)]
                    if not client and protocol == "tcp":
                        parent.append(_match(_ct("proto-dst", "original"), "@tcp_ports"))
                    packet = [ipv4, _match(_payload("ip", "protocol"), protocol)]
                    if protocol == "icmp":
                        packet.append(_match(_payload("icmp", "type"),
                                             "echo-request" if original else "echo-reply"))
                    rule(name, *packet, *parent, state, success)
                    # Related is restricted to essential ICMP errors and the
                    # same authorized parent tuple, never TCP/UDP helper flows.
                    rule(name, ipv4, _match(_payload("ip", "protocol"), "icmp"),
                         _match(_payload("icmp", "type"), {"set": [
                             "destination-unreachable", "time-exceeded", "parameter-problem",
                         ]}), *parent, _match(_ct("state"), "related", "in"), success)
            rule(name, drop)
        return policy


def _normalize(objects: list[dict[str, Any]]) -> str:
    """Canonical nft JSON, retaining every effective object, rule and verdict.

    Ignore only handles, empty set presentation and nft's redundant protocol
    dependencies. Rule order is significant; pure conjunctions and sets are not.
    Unknown expressions/properties remain different rather than being ignored.
    """
    def canonical(value: Any) -> Any:
        if isinstance(value, list):
            return [canonical(item) for item in value]
        if not isinstance(value, dict):
            return value
        value = {key: canonical(item) for key, item in value.items()}
        if (set(value) == {"|"} and isinstance(value["|"], list)
                and all(type(item) is int for item in value["|"])):
            result = 0
            for item in value["|"]:
                result |= item
            return result
        if set(value) == {"set"} and isinstance(value["set"], list):
            value["set"] = sorted(value["set"], key=lambda item: json.dumps(item, sort_keys=True))
        return value

    def bitmask(value: Any, names: dict[str, int]) -> Any:
        if isinstance(value, str):
            return names.get(value, value)
        if isinstance(value, dict) and set(value) == {"|"}:
            value = value["|"]
        if isinstance(value, list):
            bits = [bitmask(item, names) for item in value]
            if all(type(item) is int for item in bits):
                result = 0
                for item in bits:
                    result |= item
                return result
        return value

    definitions, rules = [], {}
    for item in objects:
        kind, raw = next(iter(item.items()))
        if kind == "metainfo":
            continue
        obj = {key: canonical(value) for key, value in raw.items() if key != "handle"}
        if kind == "set":
            elements = obj.get("elem", [])
            if not isinstance(elements, list):
                elements = [elements]
            obj["elem"] = sorted(elements, key=lambda value: json.dumps(value, sort_keys=True))
        if obj.get("flags") == []:
            del obj["flags"]
        elif isinstance(obj.get("flags"), str):
            obj["flags"] = [obj["flags"]]
        if kind != "rule":
            definitions.append({kind: obj})
            continue

        expressions = obj.get("expr", [])
        if not isinstance(expressions, list) or not isinstance(obj.get("chain"), str):
            raise UvmError("invalid nft rule snapshot")
        tcp_flags = {"fin": 1, "syn": 2, "rst": 4, "psh": 8, "ack": 16,
                     "urg": 32, "ecn": 64, "cwr": 128}
        for expression in expressions:
            if not isinstance(expression, dict) or not isinstance(expression.get("match", {}), dict):
                continue
            match = expression.get("match", {})
            left = match.get("left")
            if match and not {"left", "right", "op"}.issubset(match):
                raise UvmError("invalid nft match snapshot")
            right = match.get("right")
            if left in (_payload("ip", "protocol"), _ct("protocol")) and type(right) is int:
                match["right"] = {1: "icmp", 6: "tcp", 17: "udp"}.get(right, right)
            if match.get("op") == "in" and (
                (isinstance(right, str) and right.startswith("@"))
                or (isinstance(right, dict) and set(right) == {"set"})
            ):
                match["op"] = "=="
            if left == _ct("state"):
                match["right"] = bitmask(match["right"], {
                    "invalid": 1, "established": 2, "related": 4, "new": 8, "untracked": 64,
                })
            if (isinstance(left, dict) and isinstance(left.get("&"), list)
                    and len(left["&"]) == 2 and left["&"][0] == _payload("tcp", "flags")):
                left["&"][1] = bitmask(left["&"][1], tcp_flags)
                match["right"] = bitmask(match["right"], tcp_flags)

        # nft's replayable JSON omits dependencies implied by named payloads.
        encoded = json.dumps(expressions, sort_keys=True)
        has_tcp = '"protocol": "tcp"' in encoded
        has_icmp = '"protocol": "icmp"' in encoded
        expressions = [expr for expr in expressions if not (
            (has_tcp and expr == _match(_payload("ip", "protocol"), "tcp"))
            or (has_icmp and expr == _match(_payload("ip", "protocol"), "icmp"))
        )]
        encoded = json.dumps(expressions, sort_keys=True)
        expressions = [expr for expr in expressions if not (
            ('"protocol": "arp"' in encoded and expr == _match(_payload("ether", "type"), "arp"))
            or (('"protocol": "ip"' in encoded or has_icmp)
                and expr == _match(_payload("ether", "type"), "ip"))
        )]
        count = len(expressions)
        if expressions and (not isinstance(expressions[-1], dict) or set(expressions[-1]) != {"match"}):
            count -= 1
        if all(isinstance(expr, dict) and set(expr) == {"match"} for expr in expressions[:count]):
            expressions = sorted(expressions[:count], key=lambda expr: json.dumps(expr, sort_keys=True)) + expressions[count:]
        obj["expr"] = expressions
        rules.setdefault(obj.get("chain", ""), []).append(obj)
    definitions.sort(key=lambda item: json.dumps(item, sort_keys=True))
    return json.dumps([definitions, rules], sort_keys=True, separators=(",", ":"))