#!/usr/bin/env python3
"""Pi-local wlan-star WG identity from /etc/ethergold (Apply SoT — not PC PS1)."""

from __future__ import annotations

import os
import re
import subprocess
from pathlib import Path
from typing import Any

# Position 1 as spoke cannot use overlay .1 (reserved for hub topology).
SPOKE_ID_POS1_REMAP = 101
WG_PORT_DEFAULT = 51820
# Master-service return: spokes must re-handshake within ~5s (was 25).
DEFAULT_WG_KEEPALIVE_S = 5
WG_CONF = Path("/etc/wireguard/wg0.conf")
HUB_STATE_DIR = Path("/etc/ethergold/hub")
SPOKE_KEYS_DIR = Path("/etc/ethergold/spoke")


def resolve_wg_keepalive_s() -> int:
    """WireGuard PersistentKeepalive seconds (EDGE_WG_KEEPALIVE_S, default 5)."""
    raw = (os.environ.get("EDGE_WG_KEEPALIVE_S") or "").strip()
    if raw:
        try:
            return max(1, min(25, int(raw)))
        except ValueError:
            pass
    return DEFAULT_WG_KEEPALIVE_S


def spoke_id_from_bank_position(bank_position: int) -> int:
    """Map seat Position Number → EDGE_SPOKE_ID (hub overlay .1 stays reserved)."""
    pos = int(bank_position)
    if pos == 1:
        return SPOKE_ID_POS1_REMAP
    if 2 <= pos <= 254:
        return pos
    raise ValueError(f"bank_position {pos} out of range for spoke id")


def tip_pog_edge_root() -> Path | None:
    """Locate tip pog-edge/ next to this module or under ~/EtherGold/dev."""
    here = Path(__file__).resolve().parent
    for cand in (
        here / "pog-edge",
        Path(os.environ.get("HOME") or "/home/topbox") / "EtherGold" / "dev" / "pog-edge",
        Path("/home/topbox/EtherGold/dev/pog-edge"),
    ):
        if (cand / "provision_hub_spoke.sh").is_file():
            return cand
    return None


def _iface_ipv4(iface: str) -> str:
    iface = (iface or "").strip()
    if not iface:
        return ""
    try:
        proc = subprocess.run(
            ["ip", "-4", "-o", "addr", "show", "dev", iface],
            capture_output=True,
            text=True,
            timeout=5,
            check=False,
        )
    except (OSError, subprocess.TimeoutExpired):
        return ""
    # "2: wlan0    inet 10.1.10.110/24 ..."
    for line in (proc.stdout or "").splitlines():
        m = re.search(r"\binet\s+(\d+\.\d+\.\d+\.\d+)", line)
        if m:
            return m.group(1)
    return ""


def _read_text(path: Path) -> str:
    try:
        return path.read_text(encoding="utf-8", errors="replace").strip()
    except OSError:
        return ""


def _write_wg_conf(text: str) -> None:
    """Install ``/etc/wireguard/wg0.conf`` (root-owned dir — use sudo when needed)."""
    import tempfile

    payload = text if text.endswith("\n") else text + "\n"
    fd, tmp_path = tempfile.mkstemp(prefix="eg-wg0-", suffix=".conf")
    try:
        with os.fdopen(fd, "w", encoding="utf-8", newline="\n") as fh:
            fh.write(payload)
        os.chmod(tmp_path, 0o600)
        need_sudo = hasattr(os, "geteuid") and os.geteuid() != 0
        if need_sudo:
            subprocess.run(
                ["sudo", "-n", "mkdir", "-p", str(WG_CONF.parent)],
                check=True,
                capture_output=True,
                text=True,
            )
            if _path_exists_sudo(WG_CONF):
                subprocess.run(
                    [
                        "sudo",
                        "-n",
                        "cp",
                        "-a",
                        str(WG_CONF),
                        str(WG_CONF.with_name(WG_CONF.name + ".bak.apply")),
                    ],
                    check=False,
                    capture_output=True,
                    text=True,
                )
            proc = subprocess.run(
                ["sudo", "-n", "install", "-m", "600", tmp_path, str(WG_CONF)],
                check=False,
                capture_output=True,
                text=True,
            )
            if proc.returncode != 0:
                raise OSError(
                    proc.returncode,
                    (proc.stderr or proc.stdout or "sudo install wg0.conf failed").strip(),
                )
            return
        WG_CONF.parent.mkdir(parents=True, exist_ok=True)
        if WG_CONF.is_file():
            bak = WG_CONF.with_name(WG_CONF.name + ".bak.apply")
            try:
                bak.write_bytes(WG_CONF.read_bytes())
            except OSError:
                pass
        dest = WG_CONF.with_suffix(".conf.tmp")
        dest.write_text(payload, encoding="utf-8", newline="\n")
        os.chmod(dest, 0o600)
        dest.replace(WG_CONF)
    finally:
        try:
            os.unlink(tmp_path)
        except OSError:
            pass


def _path_exists_sudo(path: Path) -> bool:
    try:
        proc = subprocess.run(
            ["sudo", "-n", "test", "-e", str(path)],
            check=False,
            capture_output=True,
        )
        return proc.returncode == 0
    except OSError:
        return False


def _netadmin(cmd: list[str], *, check: bool = True) -> subprocess.CompletedProcess[str]:
    real = list(cmd)
    if real and real[0] in ("wg", "ip") and hasattr(os, "geteuid") and os.geteuid() != 0:
        real = ["sudo", "-n", *real]
    return subprocess.run(real, check=check, capture_output=True, text=True)


def _apply_spoke_hub_peer_live(
    *,
    hub_wg_pub: str,
    hub_overlay: str,
    endpoint: str,
) -> None:
    """Live-add hub Peer+Endpoint so enroll works without restarting tunnel_up."""
    _netadmin(
        [
            "wg",
            "set",
            "wg0",
            "peer",
            hub_wg_pub.strip(),
            "allowed-ips",
            f"{hub_overlay}/32",
            "endpoint",
            endpoint,
            "persistent-keepalive",
            str(resolve_wg_keepalive_s()),
        ]
    )


def ensure_live_wg_keepalive(*, wg_if: str = "wg0") -> dict[str, Any]:
    """Bump running peers to resolve_wg_keepalive_s() without rewriting conf.

    Field seats may still have PersistentKeepalive = 25 on disk from older tips;
    Master power-return must not wait that full interval.
    """
    ka = resolve_wg_keepalive_s()
    try:
        proc = _netadmin(["wg", "show", wg_if, "peers"], check=False)
    except OSError as exc:
        return {"ok": False, "error": str(exc), "keepalive_s": ka}
    pubs = [ln.strip() for ln in (proc.stdout or "").splitlines() if ln.strip()]
    applied: list[str] = []
    for pub in pubs:
        try:
            _netadmin(
                [
                    "wg",
                    "set",
                    wg_if,
                    "peer",
                    pub,
                    "persistent-keepalive",
                    str(ka),
                ],
                check=False,
            )
            applied.append(pub[:16])
        except OSError:
            continue
    return {"ok": True, "keepalive_s": ka, "peers": len(applied), "pubs": applied}


def _ensure_keypair(keys_dir: Path) -> tuple[str, str]:
    keys_dir.mkdir(parents=True, exist_ok=True)
    try:
        os.chmod(keys_dir, 0o700)
    except OSError:
        pass
    priv_path = keys_dir / "wg_private.key"
    pub_path = keys_dir / "wg_public.key"
    try:
        need_priv = (not priv_path.is_file()) or priv_path.stat().st_size == 0
    except OSError:
        need_priv = True
    if need_priv:
        proc = subprocess.run(
            ["wg", "genkey"],
            capture_output=True,
            text=True,
            timeout=10,
            check=False,
        )
        if proc.returncode != 0:
            raise RuntimeError(f"wg genkey failed: {(proc.stderr or '').strip()}")
        priv = (proc.stdout or "").strip()
        priv_path.write_text(priv + "\n", encoding="utf-8")
        os.chmod(priv_path, 0o600)
    priv = _read_text(priv_path)
    try:
        need_pub = (not pub_path.is_file()) or pub_path.stat().st_size == 0
    except OSError:
        need_pub = True
    if need_pub:
        proc = subprocess.run(
            ["wg", "pubkey"],
            input=priv + "\n",
            capture_output=True,
            text=True,
            timeout=10,
            check=False,
        )
        if proc.returncode != 0:
            raise RuntimeError(f"wg pubkey failed: {(proc.stderr or '').strip()}")
        pub_path.write_text((proc.stdout or "").strip() + "\n", encoding="utf-8")
        os.chmod(pub_path, 0o644)
    _chown_topbox(keys_dir, priv_path, pub_path)
    return priv, _read_text(pub_path)


def _chown_topbox(*paths: Path) -> None:
    """Apply often runs as root (HTTP); bridge enroll runs as topbox."""
    try:
        import pwd

        ent = pwd.getpwnam("topbox")
        uid, gid = int(ent.pw_uid), int(ent.pw_gid)
    except Exception:
        return
    for p in paths:
        try:
            os.chown(p, uid, gid)
        except OSError:
            pass



def _run_provision_hub(edge_root: Path, env: dict[str, str]) -> dict[str, Any]:
    script = edge_root / "provision_hub_spoke.sh"
    merged = os.environ.copy()
    merged.update(env)
    try:
        proc = subprocess.run(
            ["bash", str(script), "hub"],
            cwd=str(edge_root),
            capture_output=True,
            text=True,
            timeout=60,
            check=False,
            env=merged,
        )
    except (OSError, subprocess.TimeoutExpired) as exc:
        return {"ok": False, "error": str(exc), "step": "provision_hub"}
    out = {
        "ok": proc.returncode == 0,
        "exit": proc.returncode,
        "stdout": (proc.stdout or "")[-2000:],
        "stderr": (proc.stderr or "")[-1000:],
        "step": "provision_hub",
    }
    if proc.returncode != 0:
        out["error"] = (proc.stderr or proc.stdout or "provision hub failed").strip()
    return out


def _patch_edge_env(etc: Path, updates: dict[str, str]) -> None:
    path = etc / "edge.env"
    lines: list[str] = []
    if path.is_file():
        lines = path.read_text(encoding="utf-8", errors="replace").splitlines()
    keys = set(updates)
    new_lines: list[str] = []
    for line in lines:
        if not line or line.lstrip().startswith("#") or "=" not in line:
            new_lines.append(line)
            continue
        k = line.split("=", 1)[0].strip()
        if k in keys:
            new_lines.append(f"{k}={updates[k]}")
            keys.discard(k)
        else:
            new_lines.append(line)
    for k in sorted(keys):
        new_lines.append(f"{k}={updates[k]}")
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".env.tmp")
    tmp.write_text("\n".join(new_lines) + "\n", encoding="utf-8", newline="\n")
    tmp.replace(path)


def write_spoke_wg_conf(
    *,
    bank: int,
    spoke_id: int,
    hub_wg_pub: str,
    hub_endpoint: str,
    priv: str | None = None,
) -> dict[str, Any]:
    """Install full spoke wg0.conf (Interface + Peer)."""
    if not hub_wg_pub or not hub_endpoint:
        return {"ok": False, "error": "hub_wg_pub and hub_endpoint required"}
    if priv is None:
        priv, _pub = _ensure_keypair(SPOKE_KEYS_DIR)
    hub_overlay = f"10.99.{int(bank)}.1"
    overlay = f"10.99.{int(bank)}.{int(spoke_id)}"
    gretap = f"eg-s{int(spoke_id):02d}"
    endpoint = hub_endpoint.strip()
    if ":" not in endpoint:
        endpoint = f"{endpoint}:{WG_PORT_DEFAULT}"
    text = (
        f"# EtherGold wlan-star SPOKE {spoke_id} — bank {bank} overlay {overlay}\n"
        f"# gretap {gretap} rides this tunnel toward hub {hub_overlay}.\n"
        f"# Written by Apply / enroll (Pi SoT — not PC provision_wlan_star.ps1).\n"
        f"[Interface]\n"
        f"PrivateKey = {priv}\n"
        f"ListenPort = {WG_PORT_DEFAULT}\n"
        f"Table = off\n"
        f"\n"
        f"[Peer]\n"
        f"PublicKey = {hub_wg_pub.strip()}\n"
        f"AllowedIPs = {hub_overlay}/32\n"
        f"Endpoint = {endpoint}\n"
        f"PersistentKeepalive = {resolve_wg_keepalive_s()}\n"
    )
    try:
        _write_wg_conf(text)
        _apply_spoke_hub_peer_live(
            hub_wg_pub=hub_wg_pub,
            hub_overlay=hub_overlay,
            endpoint=endpoint,
        )
    except OSError as exc:
        return {"ok": False, "error": str(exc)}
    except subprocess.CalledProcessError as exc:
        err = (exc.stderr or exc.stdout or str(exc)).strip()
        return {"ok": False, "error": err or str(exc)}
    return {
        "ok": True,
        "spoke_id": int(spoke_id),
        "overlay": overlay,
        "gretap": gretap,
        "hub_endpoint": endpoint,
    }


def write_spoke_wg_interface_only(*, bank: int, spoke_id: int) -> dict[str, Any]:
    """Interface-only spoke conf — clears stale Endpoint poison before enroll."""
    priv, pub = _ensure_keypair(SPOKE_KEYS_DIR)
    overlay = f"10.99.{int(bank)}.{int(spoke_id)}"
    gretap = f"eg-s{int(spoke_id):02d}"
    text = (
        f"# EtherGold wlan-star SPOKE {spoke_id} — bank {bank} overlay {overlay}\n"
        f"# gretap {gretap} — Peer added after enroll (no stale Endpoint).\n"
        f"# Written by Apply (Pi SoT).\n"
        f"[Interface]\n"
        f"PrivateKey = {priv}\n"
        f"ListenPort = {WG_PORT_DEFAULT}\n"
        f"Table = off\n"
    )
    try:
        _write_wg_conf(text)
    except OSError as exc:
        return {"ok": False, "error": str(exc), "spoke_pub": pub}
    return {
        "ok": True,
        "spoke_id": int(spoke_id),
        "overlay": overlay,
        "gretap": gretap,
        "spoke_pub": pub,
        "peer": "pending_enroll",
    }


def apply_wlan_star_local(
    cfg: Any,
    *,
    etc: Path | None = None,
) -> dict[str, Any]:
    """Rewrite local WG identity from FieldConfig after edge.env is on disk."""
    from field_config import etc_dir  # local import avoids cycle at module load

    e = etc_dir(etc)
    if getattr(cfg, "profile", "") != "atlas_wifi_star":
        return {"ok": True, "skipped": "not atlas_wifi_star"}

    role = (getattr(cfg, "node_role", "") or "").strip()
    bank = int(getattr(cfg, "bank", 0) or 0)
    uplink = (getattr(cfg, "uplink", "") or "wlan0").strip()
    hub_ip = (getattr(cfg, "hub_uplink_ip", "") or "").strip()
    edge_env = {
        "EDGE_MODE": "1",
        "EDGE_BANK": str(bank),
        "EDGE_UPLINK": uplink,
        "EDGE_TOPOLOGY": "hub" if role == "master_hub" else "spoke",
        "EDGE_ROLE": "hub" if role == "master_hub" else "spoke",
        "EDGE_BANK_WLAN": "1",
    }

    if role == "master_hub":
        edge_root = tip_pog_edge_root()
        if edge_root is None:
            return {"ok": False, "error": "provision_hub_spoke.sh not found under tip pog-edge"}
        # Force hub conf even if keys already match a stale SPOKE file:
        # provision script skips rewrite when PrivateKey matches — remove SPOKE header.
        if WG_CONF.is_file():
            head = _read_text(WG_CONF)[:200]
            if "SPOKE" in head or "HUB" not in head:
                try:
                    WG_CONF.unlink()
                except OSError:
                    pass
        prov = _run_provision_hub(edge_root, edge_env)
        if not prov.get("ok"):
            # Fallback: write hub Interface ourselves if script needs root.
            try:
                priv, pub = _ensure_keypair(HUB_STATE_DIR)
                overlay = f"10.99.{bank}.1"
                text = (
                    f"# EtherGold wlan-star HUB — bank {bank} overlay {overlay}\n"
                    f"# Spoke peers added live by enrollment / peers.json.\n"
                    f"# Written by Apply (Pi SoT).\n"
                    f"[Interface]\n"
                    f"PrivateKey = {priv}\n"
                    f"ListenPort = {WG_PORT_DEFAULT}\n"
                    f"Table = off\n"
                )
                _write_wg_conf(text)
                HUB_STATE_DIR.mkdir(parents=True, exist_ok=True)
                peers = HUB_STATE_DIR / "peers.json"
                if not peers.is_file():
                    peers.write_text("{}\n", encoding="utf-8")
                pub = _read_text(HUB_STATE_DIR / "wg_public.key") or pub
                uplink_ip = _iface_ipv4(uplink)
                updates = {}
                if uplink_ip:
                    updates["EDGE_UPLINK_IP"] = uplink_ip
                if pub:
                    updates["EDGE_HUB_WG_PUB"] = pub
                if updates:
                    _patch_edge_env(e, updates)
                return {
                    "ok": True,
                    "role": "hub",
                    "method": "python_fallback",
                    "provision": prov,
                    "uplink_ip": uplink_ip or None,
                    "hub_wg_pub": pub or None,
                }
            except Exception as exc:
                return {"ok": False, "error": str(exc), "provision": prov, "role": "hub"}

        pub = _read_text(HUB_STATE_DIR / "wg_public.key")
        for line in (prov.get("stdout") or "").splitlines():
            if line.startswith("HUB_PUB="):
                pub = line.split("=", 1)[-1].strip() or pub
        uplink_ip = _iface_ipv4(uplink)
        updates = {}
        if uplink_ip:
            updates["EDGE_UPLINK_IP"] = uplink_ip
        if pub:
            updates["EDGE_HUB_WG_PUB"] = pub
        if updates:
            _patch_edge_env(e, updates)
        return {
            "ok": True,
            "role": "hub",
            "method": "provision_hub_spoke",
            "provision": prov,
            "uplink_ip": uplink_ip or None,
            "hub_wg_pub": pub or None,
        }

    if role != "slave_spoke":
        return {"ok": False, "error": f"unexpected node_role {role!r}"}

    try:
        spoke_id = spoke_id_from_bank_position(int(getattr(cfg, "bank_position", 0) or 0))
    except ValueError as exc:
        return {"ok": False, "error": str(exc), "role": "spoke"}

    hub_pub = (
        (getattr(cfg, "hub_wg_pub", "") or "").strip()
        or os.environ.get("EDGE_HUB_WG_PUB", "").strip()
    )
    # Prefer operator Hub IP from Apply; else leave Peer for enroll.
    if hub_ip and hub_pub:
        result = write_spoke_wg_conf(
            bank=bank,
            spoke_id=spoke_id,
            hub_wg_pub=hub_pub,
            hub_endpoint=hub_ip,
        )
        result["role"] = "spoke"
        result["spoke_id"] = spoke_id
        if hub_ip:
            _patch_edge_env(e, {"EDGE_HUB_IP": hub_ip, "EDGE_SPOKE_ID": str(spoke_id)})
        else:
            _patch_edge_env(e, {"EDGE_SPOKE_ID": str(spoke_id)})
        return result

    result = write_spoke_wg_interface_only(bank=bank, spoke_id=spoke_id)
    result["role"] = "spoke"
    updates = {"EDGE_SPOKE_ID": str(spoke_id)}
    if hub_ip:
        updates["EDGE_HUB_IP"] = hub_ip
    _patch_edge_env(e, updates)
    return result