"""
firewall.py - firewalld manager for Vacuum Wall SSL proxy firewall appliance.

Wraps firewall-cmd CLI via sudo, manages zones, rules, masquerade/NAT,
and port-forwarding. All mutations are --permanent followed by --reload.

A JSON snapshot of all rules is persisted at DATA_DIR/rules.json so the
Flask UI can inspect or restore previous configurations.
"""

import json
import os
import subprocess
from contextlib import suppress
from datetime import UTC
from pathlib import Path
from typing import Any

PROJECT_DIR = Path(__file__).resolve().parent.parent
DATA_DIR: str = str(PROJECT_DIR / "data" / "firewall")
RULES_FILE: str = os.path.join(DATA_DIR, "rules.json")
CONFIG_DIR = PROJECT_DIR / "config" / "firewall"
CONFIG_FILE = CONFIG_DIR / "config.json"

DEFAULT_CONFIG: dict[str, Any] = {"zones": {}}


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _run(cmd: list[str], check: bool = True) -> str:
    """Run a command via subprocess and return its stdout.

    Callers must include ``"sudo"`` as the first argument when the
    command requires elevated privileges.

    Raises:
        RuntimeError: When ``check=True`` and the process exits non-zero.
    """
    result = subprocess.run(cmd, capture_output=True, text=True, check=check)
    return result.stdout.strip()


def _reload() -> None:
    """Reload firewalld so permanent changes take effect immediately."""
    _run(["sudo", "firewall-cmd", "--reload"])


def _ensure_data_dir() -> None:
    """Create the data directory tree if it does not exist."""
    os.makedirs(DATA_DIR, exist_ok=True)
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)


# ---------------------------------------------------------------------------
# Read-only queries
# ---------------------------------------------------------------------------


def get_available_zones() -> list[str]:
    """Return the list of all built-in (available) firewalld zone names."""
    output = _run(["sudo", "firewall-cmd", "--get-zones"])
    return output.split()


def get_active_zones() -> dict[str, list[str]]:
    """Return a dict mapping active zone names to their assigned interfaces.

    Example return value::

        {
            "public": ["eth0"],
            "internal": ["eth1"],
        }
    """
    output = _run(["sudo", "firewall-cmd", "--get-active-zones"])
    zones: dict[str, list[str]] = {}
    current_zone: str | None = None
    for raw_line in output.splitlines():
        stripped = raw_line.strip()
        if not stripped:
            continue
        # Indented lines belong to the current zone section.
        if raw_line.startswith(" "):
            current_ifaces = (
                zones[current_zone]
                if current_zone
                else zones.get(list(zones.keys())[-1], [])
            )
            for piece in stripped.split():
                if current_zone and piece not in current_ifaces:
                    current_ifaces.append(piece)
        else:
            current_zone = stripped
            zones[current_zone] = []
    return zones


def get_zone_info(zone: str) -> dict[str, Any]:
    """Return detailed information for *zone*.

    Keys in the returned dict include:
    ``name``, ``target``, ``interfaces``, ``sources``, ``services``,
    ``ports``, ``protocols``, ``forward-ports``, ``masquerade``,
    ``rich-rules``, ``ics``, ``icmp-blocks``, ``module``.
    """
    output = _run(["sudo", "firewall-cmd", f"--zone={zone}", "--list-all"])
    info: dict[str, Any] = {"name": zone}
    for line in output.splitlines():
        line = line.strip()
        if not line or ":" not in line:
            continue
        key, _, value = line.partition(":")
        key = key.strip()
        value = value.strip()

        if not value:
            # Lines like  "interfaces:  " or "masquerade: " when disabled
            if key in ("masquerade", "ics"):
                info[key] = False
            else:
                info[key] = []
        else:
            if key in (
                "interfaces",
                "sources",
                "services",
                "ports",
                "protocols",
                "icmp-blocks",
                "module",
            ):
                info[key] = value.split()
            elif key == "forward-ports":
                info[key] = _parse_forward_ports(value)
            elif key in ("masquerade", "ics"):
                info[key] = value.lower() == "yes"
            elif key == "rich-rules":
                # rich-rules can span multiple lines; we'll parse below.
                info[key] = [value] if value else []
            else:
                info[key] = value

    # rich-rules may already have been set; if not, default to empty.
    info.setdefault("rich-rules", [])
    info.setdefault("interfaces", [])
    info.setdefault("sources", [])
    info.setdefault("services", [])
    info.setdefault("ports", [])
    info.setdefault("protocols", [])
    info.setdefault("forward-ports", [])
    info.setdefault("masquerade", False)
    info.setdefault("ics", False)
    info.setdefault("icmp-blocks", [])
    info.setdefault("module", [])
    info.setdefault("target", "default")
    return info


def get_services() -> list[str]:
    """Return the list of available service names known to firewalld."""
    output = _run(["sudo", "firewall-cmd", "--get-services"])
    return output.split()


def get_icmp_blocks() -> list[str]:
    """Return the list of available ICMP block names."""
    output = _run(["sudo", "firewall-cmd", "--get-icmptypes"])
    return output.split()


def get_interfaces() -> list[str]:
    """Return the list of network interfaces visible via iproute2."""
    output = _run(["ip", "-o", "link", "show"])
    ifaces: list[str] = []
    for line in output.splitlines():
        if line:
            # Format: "NUM: NAME: <FLAGS> ..."
            parts = line.split()
            if len(parts) >= 2:
                name = parts[1].rstrip(":")
                ifaces.append(name)
    return ifaces


def get_rich_rules(zone: str) -> list[str]:
    """Return the rich rules defined for *zone* as a list of raw strings."""
    output = _run(["sudo", "firewall-cmd", f"--zone={zone}", "--list-rich-rules"])
    output = output.strip()
    if not output:
        return []
    rules: list[str] = []
    current: list[str] = []
    for line in output.splitlines():
        raw = line.rstrip()
        if not raw.endswith(";"):
            current.append(raw)
        else:
            current.append(raw)
            rules.append(" ".join(current))
            current = []
    if current:
        rules.append(" ".join(current))
    return rules


# ---------------------------------------------------------------------------
# Zone CRUD
# ---------------------------------------------------------------------------


def create_zone(zone: str, target: str = "default") -> None:
    """Create a new permanent zone in firewalld.

    Args:
        zone: Name of the zone to create.

    Raises:
        RuntimeError: If the zone already exists or creation fails.
    """
    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            f"--set-target={target}",
            "--permanent",
        ]
    )
    _reload()


def delete_zone(zone: str) -> None:
    """Delete an existing zone.

    Raises:
        RuntimeError: If the zone does not exist or the deletion fails.
    """
    _run(["sudo", "firewall-cmd", f"--zone={zone}", "--delete", "--permanent"])
    _reload()


# ---------------------------------------------------------------------------
# Interface assignment
# ---------------------------------------------------------------------------


def set_zone_interfaces(zone: str, interfaces: list[str]) -> None:
    """Assign *interfaces* to *zone*, replacing any existing assignments.

    Existing interfaces on the zone are removed first so only the
    provided list remains.
    """
    # Remove current permanent interfaces for this zone.
    try:
        current = get_zone_info(zone).get("interfaces", [])
    except Exception:
        current = []
    for iface in current:
        _run(
            [
                "sudo",
                "firewall-cmd",
                f"--zone={zone}",
                "--remove-interface=" + iface,
                "--permanent",
            ],
            check=False,
        )

    # Add the desired set.
    for iface in interfaces:
        _run(
            [
                "sudo",
                "firewall-cmd",
                f"--zone={zone}",
                "--add-interface=" + iface,
                "--permanent",
            ]
        )
    _reload()


def add_zone_interface(zone: str, iface: str) -> None:
    """Add a single interface to *zone*."""
    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            "--add-interface=" + iface,
            "--permanent",
        ]
    )
    _reload()


def remove_zone_interface(zone: str, iface: str) -> None:
    """Remove a single interface from *zone*."""
    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            "--remove-interface=" + iface,
            "--permanent",
        ]
    )
    _reload()


# ---------------------------------------------------------------------------
# Service management
# ---------------------------------------------------------------------------


def set_zone_services(zone: str, services: list[str]) -> None:
    """Set services for *zone*, replacing any previously allowed services."""
    # Remove all current services.
    current = get_zone_info(zone).get("services", [])
    for svc in current:
        _run(
            [
                "sudo",
                "firewall-cmd",
                f"--zone={zone}",
                f"--remove-service={svc}",
                "--permanent",
            ],
            check=False,
        )

    for svc in services:
        _run(
            [
                "sudo",
                "firewall-cmd",
                f"--zone={zone}",
                f"--add-service={svc}",
                "--permanent",
            ]
        )
    _reload()


def add_zone_service(zone: str, service: str) -> None:
    """Add a single service to *zone*."""
    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            f"--add-service={service}",
            "--permanent",
        ]
    )
    _reload()


def remove_zone_service(zone: str, service: str) -> None:
    """Remove a single service from *zone*."""
    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            f"--remove-service={service}",
            "--permanent",
        ]
    )
    _reload()


# ---------------------------------------------------------------------------
# Rich rules
# ---------------------------------------------------------------------------


def add_rich_rule(zone: str, rule: str) -> None:
    """Add a rich rule to *zone*.

    The *rule* argument should be a fully-formed rich-rule expression,
    e.g. ``rule family="ipv4" port protocol="tcp" port="443" accept``.
    """
    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            "--add-rich-rule=" + rule,
            "--permanent",
        ]
    )
    _reload()


def remove_rich_rule(zone: str, rule: str) -> None:
    """Remove a rich rule from *zone*.

    The rule string must match exactly what was added.
    """
    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            "--remove-rich-rule=" + rule,
            "--permanent",
        ]
    )
    _reload()


# ---------------------------------------------------------------------------
# Masquerade (NAT)
# ---------------------------------------------------------------------------


def set_masquerade(zone: str, enable: bool) -> None:
    """Enable or disable masquerade (source-NAT) on *zone*."""
    action = "--add-masquerade" if enable else "--remove-masquerade"
    _run(["sudo", "firewall-cmd", f"--zone={zone}", action, "--permanent"])
    _reload()


# ---------------------------------------------------------------------------
# Port forwarding
# ---------------------------------------------------------------------------


def add_forward_port(
    zone: str,
    port: int,
    protocol: str,
    toaddr: str | None = None,
    toport: int | None = None,
) -> None:
    """Add a port forwarding rule to *zone*.

    Forward traffic arriving on ``port/protocol`` to
    ``toaddr:toport`` (or just ``toport`` when *toaddr* is omitted).
    """
    fwd = f"port={port}/proto={protocol}"
    if toaddr and toport:
        fwd += f"/toaddr={toaddr}/toport={toport}"
    elif toport:
        fwd += f"/toport={toport}"
    else:
        fwd += f"/toaddr={toaddr}" if toaddr else ""

    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            f"--add-forward-port={fwd}",
            "--permanent",
        ]
    )
    _reload()


def remove_forward_port(
    zone: str,
    port: int,
    protocol: str,
    toaddr: str | None = None,
    toport: int | None = None,
) -> None:
    """Remove a previously added port-forwarding rule from *zone*.

    All parameters must match the original rule exactly.
    """
    fwd = f"port={port}/proto={protocol}"
    if toaddr and toport:
        fwd += f"/toaddr={toaddr}/toport={toport}"
    elif toport:
        fwd += f"/toport={toport}"
    else:
        fwd += f"/toaddr={toaddr}" if toaddr else ""

    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            f"--remove-forward-port={fwd}",
            "--permanent",
        ]
    )
    _reload()


# ---------------------------------------------------------------------------
# Helpers for parsing forward-port lines
# ---------------------------------------------------------------------------


def _parse_forward_port(raw: str) -> dict[str, Any]:
    """Parse a single forward-port specifier into a structured dict.

    Input: ``port=443/proto=tcp/toaddr=192.168.1.5/toport=8080``
    """
    result: dict[str, Any] = {}
    for piece in raw.split("/"):
        if "=" not in piece:
            continue
        key, _, val = piece.partition("=")
        if key == "port":
            result["port"] = int(val)
        elif key == "proto":
            result["proto"] = val
        elif key == "toaddr":
            result["toaddr"] = val
        elif key == "toport":
            result["toport"] = int(val)
    return result


def _parse_forward_ports(value: str) -> list[dict[str, Any]]:
    """Parse the 'forward-ports' line into a list of structured dicts."""
    if not value:
        return []
    return [_parse_forward_port(raw) for raw in value.split()]


# ---------------------------------------------------------------------------
# State snapshot / backup helpers
# ---------------------------------------------------------------------------


def get_state() -> dict[str, Any]:
    """Return the complete current state of firewalld as a Python dict.

    The dict contains all zones with their per-zone configuration, all
    rich rules, masquerade settings, forward-port rules, and the set of
    active interfaces.
    """
    zones: dict[str, dict[str, Any]] = {}
    for name in get_available_zones():
        try:
            zones[name] = get_zone_info(name)
        except Exception:
            continue

    return {
        "active_zones": get_active_zones(),
        "interfaces": get_interfaces(),
        "available_services": get_services(),
        "zones": zones,
        "rich_rules": {name: get_rich_rules(name) for name in zones},
        "timestamp": _now_iso(),
    }


def _now_iso() -> str:
    """Return the current UTC time as an ISO-8601 string."""
    from datetime import datetime

    return datetime.now(UTC).isoformat()


def save_backup() -> str:
    """Capture the full state and write it to RULES_FILE on disk.

    Returns:
        Absolute path to the written file.
    """
    _ensure_data_dir()
    state = get_state()
    with open(RULES_FILE, "w") as fh:
        json.dump(state, fh, indent=2, default=str)
    return RULES_FILE


def load_backup() -> dict[str, Any]:
    """Read the JSON backup file and return the state dict.

    Use :func:`restore_backup` to actually apply the loaded state.

    Raises:
        FileNotFoundError: When no backup file exists at RULES_FILE.
        json.JSONDecodeError: When the file is not valid JSON.

    Returns:
        The loaded state dict.
    """
    with open(RULES_FILE) as fh:
        state: dict[str, Any] = json.load(fh)
    return state


def restore_backup(state: dict[str, Any]) -> None:
    """Apply the zone configuration described in *state*.

    Walks every zone in *state*["zones"] and re-creates services,
    interfaces, forward ports, masquerade, and rich rules.

    This is a *merge*: zones not present in the snapshot are **not**
    touched.
    """
    zones_cfg = state.get("zones", {})
    for zone_name, zinfo in zones_cfg.items():
        # Ensure the zone exists.
        if zone_name not in get_available_zones():
            target = zinfo.get("target", "default")
            create_zone(zone_name, target)

        # Services
        services = zinfo.get("services", [])
        set_zone_services(zone_name, services)

        # Interfaces
        interfaces = zinfo.get("interfaces", [])
        set_zone_interfaces(zone_name, interfaces)

        # Masquerade
        if zinfo.get("masquerade"):
            set_masquerade(zone_name, True)

        # Forward ports  (stored as dicts, or raw strings from old backups)
        for fp in zinfo.get("forward-ports", []):
            if isinstance(fp, str):
                fp_str = fp
            else:
                parts = [f"port={fp['port']}", f"proto={fp['proto']}"]
                if "toaddr" in fp:
                    parts.append(f"toaddr={fp['toaddr']}")
                if "toport" in fp:
                    parts.append(f"toport={fp['toport']}")
                fp_str = "/".join(parts)
            _run(
                [
                    "sudo",
                    "firewall-cmd",
                    f"--zone={zone_name}",
                    f"--add-forward-port={fp_str}",
                    "--permanent",
                ],
                check=False,
            )

        # Rich rules
        for rule in zinfo.get("rich-rules", []):
            _run(
                [
                    "sudo",
                    "firewall-cmd",
                    f"--zone={zone_name}",
                    f"--add-rich-rule={rule}",
                    "--permanent",
                ],
                check=False,
            )

    _reload()


# ---------------------------------------------------------------------------
# Declarative config management (config/firewall/config.json)
# ---------------------------------------------------------------------------


def _ensure_config_file() -> None:
    """Create config directory and file if they do not exist."""
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    if not CONFIG_FILE.exists():
        with open(CONFIG_FILE, "w") as fh:
            json.dump(DEFAULT_CONFIG, fh, indent=2)
            fh.write("\n")


def config_get() -> dict[str, Any]:
    """Return the declarative config from ``config/firewall/config.json``."""
    _ensure_config_file()
    with open(CONFIG_FILE) as fh:
        return json.load(fh)


def config_set(cfg: dict[str, Any]) -> None:
    """Write *cfg* to ``config/firewall/config.json`` (atomic replace)."""
    _ensure_config_file()
    tmp = CONFIG_FILE.with_name(CONFIG_FILE.name + ".tmp")
    with open(tmp, "w") as fh:
        json.dump(cfg, fh, indent=2)
        fh.write("\n")
    os.replace(tmp, CONFIG_FILE)


def _normalize_target(target: str) -> str:
    """Map between config JSON target names and firewalld target values."""
    if target == "ACCEPT":
        return "ACCEPT"
    if target == "DROP":
        return "DROP"
    if target == "REJECT":
        return "REJECT"
    return "default"


def _live_target_to_config(target: str) -> str:
    """Map firewalld target value back to config JSON canonical form."""
    if target == "ACCEPT":
        return "ACCEPT"
    if target == "DROP":
        return "DROP"
    if target == "REJECT":
        return "REJECT"
    return "DEFAULT"


def config_pending() -> dict[str, Any]:
    """Compare declarative config against live firewalld state, return diff.

    Returns a dict with ``pending`` (list of change dicts), ``needs_apply``
    (bool), and ``live_zones`` (dict of zones not yet in config).
    """
    cfg = config_get()
    live_state = get_state()
    cfg_zones = cfg.get("zones", {})
    live_zones = live_state.get("zones", {})

    changes: list[dict[str, Any]] = []
    unknown_live: dict[str, Any] = {}

    for zone_name, zone_cfg in cfg_zones.items():
        live_zone = live_zones.get(zone_name, {})
        if not zone_cfg.get("interfaces"):
            continue

        cfg_ifaces = set(zone_cfg.get("interfaces", []))
        live_ifaces = set(live_zone.get("interfaces", []))
        if cfg_ifaces != live_ifaces:
            changes.append(
                {
                    "zone": zone_name,
                    "type": "interfaces",
                    "config": sorted(cfg_ifaces),
                    "live": sorted(live_ifaces),
                }
            )

        cfg_services = set(zone_cfg.get("services", []))
        live_services = set(live_zone.get("services", []))
        if cfg_services != live_services:
            changes.append(
                {
                    "zone": zone_name,
                    "type": "services",
                    "config": sorted(cfg_services),
                    "live": sorted(live_services),
                }
            )

        cfg_target = _normalize_target(zone_cfg.get("target", "DEFAULT"))
        live_target = live_zone.get("target", "default")
        if cfg_target != live_target:
            changes.append(
                {
                    "zone": zone_name,
                    "type": "target",
                    "config": cfg_target,
                    "live": live_target,
                }
            )

        cfg_mq = zone_cfg.get("masquerade", False)
        live_mq = live_zone.get("masquerade", False)
        if cfg_mq != live_mq:
            changes.append(
                {
                    "zone": zone_name,
                    "type": "masquerade",
                    "config": cfg_mq,
                    "live": live_mq,
                }
            )

    for zone_name in live_zones:
        if zone_name not in cfg_zones:
            unknown_live[zone_name] = {
                "interfaces": live_zones[zone_name].get("interfaces", []),
            }

    return {
        "pending": changes,
        "needs_apply": len(changes) > 0,
        "unmanaged_zones": unknown_live,
    }


def config_apply() -> dict[str, Any]:
    """Apply the declarative config to live firewalld.

    Takes a snapshot via ``save_backup()`` first, then reconciles each zone
    in the config (create/update, interfaces, services, masquerade), reloads,
    and takes another snapshot.

    Returns a dict with ``applied_zones`` and a ``backup`` path.
    """
    cfg = config_get()
    cfg_zones = cfg.get("zones", {})

    save_backup()

    available = get_available_zones()
    applied: list[str] = []
    for zone_name, zone_cfg in cfg_zones.items():
        need_create = zone_name not in available

        if need_create:
            target = _normalize_target(zone_cfg.get("target", "DEFAULT"))
            create_zone(zone_name, target)
        else:
            desired_target = _normalize_target(zone_cfg.get("target", "DEFAULT"))
            if desired_target != "default":
                with suppress(RuntimeError):
                    _run(
                        [
                            "sudo",
                            "firewall-cmd",
                            f"--zone={zone_name}",
                            f"--set-target={desired_target}",
                            "--permanent",
                        ],
                        check=False,
                    )

        set_zone_services(zone_name, zone_cfg.get("services", []))
        set_zone_interfaces(zone_name, zone_cfg.get("interfaces", []))

        mq = zone_cfg.get("masquerade", False)
        if mq is not None:
            set_masquerade(zone_name, mq)

        applied.append(zone_name)

    _reload()
    backup_path = save_backup()

    return {
        "applied_zones": applied,
        "backup": backup_path,
    }


__all__ = [
    "CONFIG_DIR",
    "CONFIG_FILE",
    "DATA_DIR",
    "DEFAULT_CONFIG",
    "RULES_FILE",
    "_reload",
    "_run",
    "add_forward_port",
    "add_rich_rule",
    "add_zone_interface",
    "add_zone_service",
    "config_apply",
    "config_get",
    "config_pending",
    "config_set",
    "create_zone",
    "delete_zone",
    "get_active_zones",
    "get_available_zones",
    "get_icmp_blocks",
    "get_interfaces",
    "get_rich_rules",
    "get_services",
    "get_state",
    "get_zone_info",
    "load_backup",
    "remove_forward_port",
    "remove_rich_rule",
    "remove_zone_interface",
    "remove_zone_service",
    "restore_backup",
    "save_backup",
    "set_masquerade",
    "set_zone_interfaces",
    "set_zone_services",
]