"""
WireGuard Manager for Vacuum Wall SSL Proxy Firewall.

Generates wg-quick configurations, manages peers, and controls
the WireGuard tunnel interface.
"""

import json
import logging
import os
import subprocess
from datetime import UTC, datetime
from pathlib import Path

from jinja2 import Environment, FileSystemLoader

logger = logging.getLogger(__name__)

PROJECT_DIR = Path(__file__).resolve().parent.parent
CONFIG_PATH = str(PROJECT_DIR / "config" / "wireguard" / "config.json")
WG_CONF_PATH = "/etc/wireguard/wg0.conf"
WG_QUICK_BIN = "wg-quick"
WG_BIN = "wg"

ENV = Environment(
    loader=FileSystemLoader(str(PROJECT_DIR / "system")),
    autoescape=False,
    lstrip_blocks=True,
    trim_blocks=True,
)


# --- Helpers ---


def _run(cmd: list[str], check: bool = True) -> subprocess.CompletedProcess:
    """Run a command via sudo and return the completed process."""
    return subprocess.run(
        ["sudo", *cmd],
        capture_output=True,
        text=True,
        check=check,
    )


def _ensure_dir(path: str) -> None:
    """Create parent directories for *path* if they don't exist."""
    Path(path).parent.mkdir(parents=True, exist_ok=True)


def _default_config() -> dict:
    """Return the skeleton config with no keys and no peers."""
    return {
        "interface": {
            "name": "wg0",
            "listen_port": 51820,
            "private_key": "",
            "public_key": "",
            "addresses": ["10.137.0.1/24"],
            "post_up": None,
            "post_down": None,
        },
        "peers": {},
    }


# --- Core config persistence ---


def get_config() -> dict:
    """Load the current WireGuard configuration from the JSON store."""
    try:
        with open(CONFIG_PATH) as f:
            cfg = json.load(f)
        defaults = _default_config()
        cfg.setdefault("interface", defaults["interface"])
        cfg["interface"].setdefault("name", defaults["interface"]["name"])
        cfg["interface"].setdefault("listen_port", defaults["interface"]["listen_port"])
        cfg["interface"].setdefault("private_key", defaults["interface"]["private_key"])
        cfg["interface"].setdefault("public_key", defaults["interface"]["public_key"])
        cfg["interface"].setdefault("addresses", defaults["interface"]["addresses"])
        cfg["interface"].setdefault("post_up", defaults["interface"]["post_up"])
        cfg["interface"].setdefault("post_down", defaults["interface"]["post_down"])
        cfg.setdefault("peers", {})
        return cfg
    except (FileNotFoundError, json.JSONDecodeError):
        return _default_config()


def save_config(cfg: dict) -> None:
    """Persist *cfg* to the JSON store atomically."""
    _ensure_dir(CONFIG_PATH)
    tmp = CONFIG_PATH + ".tmp"
    with open(tmp, "w") as f:
        json.dump(cfg, f, indent=4)
        f.write("\n")
    os.replace(tmp, CONFIG_PATH)


# --- Key generation ---


def generate_keypair() -> tuple[str, str]:
    """Generate a WireGuard private/public key pair using ``wg`` CLI."""
    res = _run([WG_BIN, "genkey"])
    private_key = res.stdout.strip()
    res2 = _run([WG_BIN, "pubkey"], input=private_key)
    public_key = res2.stdout.strip()
    return private_key, public_key


# --- wg0.conf generation ---


def generate_conf(cfg: dict) -> str:
    """Render a valid wg-quick config file from *cfg* using Jinja2."""
    tmpl = ENV.get_template("wireguard.conf")
    return tmpl.render(
        timestamp=datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"),
        interface=cfg["interface"],
        peers=cfg.get("peers", {}),
    )


# --- Apply / down ---


def apply() -> None:
    """Write the current config to disk and bring the tunnel up with wg-quick."""
    cfg = get_config()
    conf_text = generate_conf(cfg)
    save_config(cfg)

    local_dir = PROJECT_DIR / "data" / "wireguard"
    local_dir.mkdir(parents=True, exist_ok=True)
    local_tmp = local_dir / "wg0.conf.tmp"
    with open(local_tmp, "w") as f:
        f.write(conf_text)
    os.chmod(local_tmp, 0o600)
    _run(["cp", "--", str(local_tmp), WG_CONF_PATH])
    _run(["chown", "root:root", WG_CONF_PATH], check=False)
    local_tmp.unlink(missing_ok=True)

    _run([WG_QUICK_BIN, "up", cfg["interface"]["name"]])
    logger.info("WireGuard tunnel '%s' brought up", cfg["interface"]["name"])


def down() -> None:
    """Bring the WireGuard tunnel interface down."""
    cfg = get_config()
    name = cfg["interface"]["name"]
    _run([WG_QUICK_BIN, "down", name])
    logger.info("WireGuard tunnel '%s' brought down", name)


# --- Status ---


def status() -> dict:
    """Query the live tunnel state via ``wg show``."""
    cfg = get_config()
    name = cfg["interface"]["name"]
    result = {
        "up": False,
        "interface": {},
        "peers": [],
    }

    try:
        proc = _run([WG_BIN, "show", name], check=False)
        if proc.returncode != 0:
            return result

        raw = proc.stdout.strip()
    except Exception:
        return result

    current_peer = None
    peers: list[dict] = []

    for line in raw.splitlines():
        line = line.strip()
        if not line:
            continue

        if line.startswith("interface:"):
            result["up"] = True
            result["interface"] = {}
            current_peer = None
            continue

        if line.startswith("public key:"):
            result["interface"]["public_key"] = line.split(":", 1)[1].strip()
            continue

        if line.startswith("listening port:"):
            val = line.split(":", 1)[1].strip()
            result["interface"]["listen_port"] = int(val)
            continue

        if line.startswith("fwmark:"):
            result["interface"]["fwmark"] = line.split(":", 1)[1].strip()
            continue

        if line.startswith("peer:"):
            cur_key = line.split(":", 1)[1].strip()
            current_peer = {
                "public_key": cur_key,
                "endpoint": None,
                "allowed_ips": [],
                "latest_handshake": None,
                "transfer_received": 0,
                "transfer_sent": 0,
                "persistent_keepalive": None,
            }
            peers.append(current_peer)
            continue

        if current_peer is None:
            continue

        if line.startswith("endpoint:"):
            current_peer["endpoint"] = line.split(":", 1)[1].strip()
            continue

        if line.startswith("allowed ips:"):
            vals = line.split(":", 1)[1].strip().split(", ")
            current_peer["allowed_ips"] = vals
            continue

        if line.startswith("latest handshake:"):
            current_peer["latest_handshake"] = line.split(":", 1)[1].strip()
            continue

        if line.startswith("transfer:"):
            rest = line.split(":", 1)[1].strip()
            parts = rest.split(", ")
            if parts:
                current_peer["transfer_received"] = parts[0].strip()
            if len(parts) > 1:
                current_peer["transfer_sent"] = parts[1].strip()
            continue

        if line.startswith("persistent-keepalive:"):
            val = line.split(":", 1)[1].strip()
            try:
                current_peer["persistent_keepalive"] = int(val)
            except ValueError:
                current_peer["persistent_keepalive"] = None

    result["peers"] = peers
    return result


# --- Peer management ---


def add_peer(
    name: str,
    endpoint: str | None = None,
    allowed_ips: list[str] | None = None,
    persistent_keepalive: int | None = None,
    preshared_key: str | None = None,
) -> dict:
    """Add (or update) a peer in the configuration."""
    cfg = get_config()
    peers = cfg.setdefault("peers", {})
    allowed_ips = allowed_ips or []

    if name in peers:
        peer = peers[name]
        peer["endpoint"] = endpoint
        peer["allowed_ips"] = allowed_ips
        peer["persistent_keepalive"] = persistent_keepalive
        if preshared_key is not None:
            peer["preshared_key"] = preshared_key
        logger.info("WireGuard peer '%s' updated", name)
    else:
        priv, pub = generate_keypair()
        peer = {
            "public_key": pub,
            "private_key": priv,
            "endpoint": endpoint,
            "allowed_ips": allowed_ips,
            "persistent_keepalive": persistent_keepalive,
            "preshared_key": preshared_key,
        }
        peers[name] = peer
        logger.info("WireGuard peer '%s' added (pubkey=%s...)", name, pub[:16])

    save_config(cfg)
    peer_out = dict(peer)
    return peer_out


def remove_peer(name: str) -> None:
    """Remove a peer from the configuration by name."""
    cfg = get_config()
    cfg.setdefault("peers", {}).pop(name, None)
    save_config(cfg)
    logger.info("WireGuard peer '%s' removed", name)


def get_peers() -> list[dict]:
    """List all configured peers (from the JSON store, *not* live)."""
    cfg = get_config()
    peers = []
    for name, info in cfg.get("peers", {}).items():
        entry = dict(info)
        entry["name"] = name
        entry.pop("private_key", None)
        peers.append(entry)
    return peers


def get_peer_status() -> list[dict]:
    """Return live peer status from ``wg show``."""
    st = status()
    return st.get("peers", [])


# --- Client config generation ---


def generate_client_conf(
    peer_name: str,
    server_endpoint: str,
    server_pubkey: str,
) -> str:
    """Build a client-side wg-quick config snippet for *peer_name*."""
    cfg = get_config()
    iface = cfg["interface"]
    peer = cfg["peers"].get(peer_name)
    if peer is None:
        raise KeyError(f"Peer '{peer_name}' not found in configuration")

    client_priv = peer.get("private_key", "")
    if not client_priv:
        raise ValueError(
            f"Peer '{peer_name}' has no private key — cannot generate client config."
        )

    sorted_peers = sorted(cfg.get("peers", {}).keys())
    peer_index = sorted_peers.index(peer_name) + 2
    srv_addr = iface["addresses"][0] if iface["addresses"] else "10.137.0.1/24"
    addr_part, prefix = srv_addr.rsplit("/", 1)
    prefix_base = addr_part.rsplit(".", 1)[0]
    client_addr = f"{prefix_base}.{peer_index}/{prefix}"

    tmpl = ENV.get_template("wireguard-client.conf")
    conf = tmpl.render(
        timestamp=datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"),
        peer_name=peer_name,
        client_priv=client_priv,
        client_addr=client_addr,
        server_pubkey=server_pubkey,
        server_endpoint=server_endpoint,
        allowed_ips=peer.get("allowed_ips", ["0.0.0.0/0"]),
        preshared_key=peer.get("preshared_key"),
        persistent_keepalive=peer.get("persistent_keepalive"),
    )
    logger.info("Client config generated for peer '%s'", peer_name)
    return conf


# --- Interface-level setters ---


def set_listen_port(port: int) -> None:
    """Update the server listen port in the stored configuration."""
    if not (1 <= port <= 65535):
        raise ValueError("Listen port must be in range 1..65535")
    cfg = get_config()
    cfg["interface"]["listen_port"] = port
    save_config(cfg)
    logger.info("WireGuard listen port set to %d", port)


def set_post_up(cmd: str | None) -> None:
    """Set (or clear) the PostUp hook command."""
    cfg = get_config()
    cfg["interface"]["post_up"] = cmd
    save_config(cfg)


def set_post_down(cmd: str | None) -> None:
    """Set (or clear) the PostDown hook command."""
    cfg = get_config()
    cfg["interface"]["post_down"] = cmd
    save_config(cfg)


# --- Initialise ---


def initialize() -> dict:
    """Perform first-time WireGuard setup."""
    cfg = get_config()

    if cfg["interface"].get("private_key"):
        return cfg

    priv, pub = generate_keypair()
    cfg["interface"]["private_key"] = priv
    cfg["interface"]["public_key"] = pub
    save_config(cfg)
    logger.info("WireGuard initialised (pubkey=%s...)", pub[:16])
    return cfg


# --- Utility: parse wg show into structured peer map ---


def _parse_wg_show(output: str) -> dict:
    """Internal parser for ``wg show`` multiline output."""
    peers: dict = {}
    current = None

    for line in output.splitlines():
        line = line.strip()
        if line.startswith("peer:"):
            key = line.split(":", 1)[1].strip()
            current = {"_key": key}
            peers[key] = current
            continue

        if current is None:
            continue

        if line.startswith("endpoint:"):
            val = line.split(":", 1)[1].strip()
            current["endpoint"] = val
        elif line.startswith("allowed ips:"):
            current["allowed_ips"] = line.split(":", 1)[1].strip()
        elif line.startswith("latest handshake:"):
            current["latest_handshake"] = line.split(":", 1)[1].strip()
        elif line.startswith("transfer:"):
            current["transfer_raw"] = line.split(":", 1)[1].strip()
        elif line.startswith("persistent-keepalive:"):
            current["persistent_keepalive"] = line.split(":", 1)[1].strip()

    return peers