"""Network egress control — enforce allowed-domain policies in containers."""

from __future__ import annotations

import socket
import subprocess
from pathlib import Path

from nestdx.config.schema import NetworkPolicy


class NetworkPolicyError(Exception):
    """Network policy setup or teardown failed."""


def resolve_domains(domains: list[str]) -> dict[str, list[str]]:
    """Resolve domain names to IP addresses.

    Returns a mapping of domain → list of IPs.

    Raises:
        NetworkPolicyError: If any domain fails to resolve.
    """
    resolved: dict[str, list[str]] = {}
    for domain in domains:
        try:
            infos = socket.getaddrinfo(domain, None, socket.AF_UNSPEC, socket.SOCK_STREAM)
            ips = sorted({info[4][0] for info in infos})
            if not ips:
                raise NetworkPolicyError(f"DNS resolved no IPs for '{domain}'")
            resolved[domain] = ips
        except socket.gaierror as exc:
            raise NetworkPolicyError(
                f"DNS resolution failed for '{domain}': {exc}"
            ) from exc
    return resolved


class NetworkPolicyEnforcer:
    """Applies network egress restrictions to a container.

    Strategy:
    1. Container starts with --network=none (fully isolated).
    2. A custom bridge network is created with the allowed IPs.
    3. The container is connected to the bridge network.
    4. iptables rules inside the container restrict outbound to allowed IPs only.

    When block_all_other_egress is False, no restrictions are applied.
    """

    def __init__(
        self,
        policy: NetworkPolicy,
        runtime: str,
        project_root: Path | None = None,
    ):
        self.policy = policy
        self.runtime = runtime
        self.project_root = project_root
        self._network_name: str | None = None
        self._resolved_ips: dict[str, list[str]] = {}

    @property
    def should_enforce(self) -> bool:
        """True if network egress restrictions should be applied."""
        return self.policy.block_all_other_egress and len(self.policy.allowed_domains) > 0

    def resolve(self) -> dict[str, list[str]]:
        """Resolve allowed domains to IPs. Caches for session duration."""
        if not self._resolved_ips and self.policy.allowed_domains:
            self._resolved_ips = resolve_domains(self.policy.allowed_domains)
        return self._resolved_ips

    def build_run_args(self) -> list[str]:
        """Return extra `docker run` flags for network isolation.

        When enforcement is active, starts with --network=none.
        When not enforcing, returns an empty list (default bridge network).
        """
        if not self.should_enforce:
            return []
        return ["--network=none"]

    def setup_network(self, container_id: str, session_id: str) -> None:
        """Create a custom bridge network and connect the container.

        After connecting, adds iptables rules inside the container to allow
        only the resolved IPs (plus DNS at 127.0.0.11 for Docker's embedded DNS).
        """
        if not self.should_enforce:
            return

        resolved = self.resolve()
        self._network_name = f"nestdx-net-{session_id}"

        # Create custom bridge network
        result = subprocess.run(
            [self.runtime, "network", "create", self._network_name],
            capture_output=True,
            timeout=30,
        )
        if result.returncode != 0:
            stderr = result.stderr.decode(errors="replace")
            raise NetworkPolicyError(
                f"Failed to create network '{self._network_name}': {stderr}"
            )

        # Connect container to the custom network
        result = subprocess.run(
            [self.runtime, "network", "connect", self._network_name, container_id],
            capture_output=True,
            timeout=30,
        )
        if result.returncode != 0:
            stderr = result.stderr.decode(errors="replace")
            raise NetworkPolicyError(
                f"Failed to connect container to network: {stderr}"
            )

        # Disconnect from the "none" network is implicit (already --network=none).
        # Now set iptables rules inside the container to restrict egress.
        allowed_ips = []
        for ips in resolved.values():
            allowed_ips.extend(ips)

        # Build iptables script: default DROP on OUTPUT, allow only resolved IPs + loopback + DNS
        rules = [
            # Allow loopback
            "iptables -A OUTPUT -o lo -j ACCEPT",
            # Allow established connections (responses to allowed requests)
            "iptables -A OUTPUT -m state --state ESTABLISHED,RELATED -j ACCEPT",
            # Allow Docker embedded DNS (127.0.0.11)
            "iptables -A OUTPUT -d 127.0.0.11 -p udp --dport 53 -j ACCEPT",
            "iptables -A OUTPUT -d 127.0.0.11 -p tcp --dport 53 -j ACCEPT",
        ]

        for ip in sorted(set(allowed_ips)):
            rules.append(f"iptables -A OUTPUT -d {ip} -j ACCEPT")

        # Default drop everything else
        rules.append("iptables -A OUTPUT -j DROP")

        script = " && ".join(rules)

        result = subprocess.run(
            [self.runtime, "exec", container_id, "sh", "-c", script],
            capture_output=True,
            timeout=30,
        )
        if result.returncode != 0:
            stderr = result.stderr.decode(errors="replace")
            raise NetworkPolicyError(
                f"Failed to apply iptables rules: {stderr}"
            )

    def teardown_network(self, container_id: str) -> None:
        """Disconnect container from custom network and remove the network."""
        if self._network_name is None:
            return

        # Disconnect (ignore errors — container may already be stopped)
        subprocess.run(
            [self.runtime, "network", "disconnect", self._network_name, container_id],
            capture_output=True,
            timeout=30,
        )

        # Remove network
        subprocess.run(
            [self.runtime, "network", "rm", self._network_name],
            capture_output=True,
            timeout=30,
        )

        self._network_name = None

    def validate_egress(self, domain: str) -> bool:
        """Check if a domain would be reachable under this policy.

        For dry-run validation: resolves domain and checks if it's in allowed list.
        Returns True if allowed, False if would be blocked.
        """
        if not self.policy.block_all_other_egress:
            return True
        return domain in self.policy.allowed_domains
