#!/usr/bin/env python3
"""Wallarm Node Manager — deploy, edit, remove, status, tunnel."""

import json, os, sys, time, subprocess, shutil, hashlib
from pathlib import Path

VERSION = "1.0.0"
BASE = "/opt/fw"
APP = f"{BASE}/app"
STATE = f"{APP}/state.json"
CONF = f"{APP}/fw.conf"
AIO_URL = "https://storage.googleapis.com/meganode_storage/6.12/wallarm-6.12.5.x86_64-glibc.sh"
AIO_PATH = f"{BASE}/wallarm-aio.sh"
DEPLOY_TO = "/opt/wallarm"

# ─── Helpers ───────────────────────────────────────────────────────────

def run(cmd, timeout=120, check=False):
    return subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout)

def run_ok(cmd):
    return subprocess.run(cmd, shell=True, capture_output=True).returncode == 0

def hash_port(s):
    return abs(hash(s)) % 500

def is_root():
    return os.geteuid() == 0

# ─── Preflight ─────────────────────────────────────────────────────────

def preflight():
    checks = []
    ok = True
    def check(name, passed, detail=""):
        nonlocal ok
        if not passed: ok = False
        m = "✅" if passed else "❌"
        print(f"  {m} {name} — {detail}")
        checks.append((name, passed))

    check("root", is_root())
    check("systemd", run_ok("systemctl --version"))
    arch = run("uname -m").stdout.strip()
    check("arch", arch in ("x86_64", "aarch64"), arch)
    check("curl/wget", run_ok("which curl") or run_ok("which wget"), "downloader")
    for c in ("systemctl", "sed", "mkdir", "rm"):
        check(f"cmd:{c}", run_ok(f"which {c}"))
    check("installer", run_ok(f"curl -fsSL -o /dev/null {AIO_URL} 2>/dev/null"), "meganode")
    for cloud, host in (("US", "us1.api.wallarm.com"), ("EU", "api.wallarm.com")):
        reachable = run_ok(f"curl -sL --connect-timeout 5 -o /dev/null https://{host} 2>/dev/null")
        m = "✅" if reachable else "⚠️ "
        print(f"  {m} cloud:{cloud} — {host}{' (unreachable, will retry on deploy)' if not reachable else ''}")
    return ok

# ─── Config ────────────────────────────────────────────────────────────

def load_config():
    if not os.path.exists(CONF):
        return {}
    with open(CONF) as f:
        return json.load(f)

def load_state():
    if not os.path.exists(STATE):
        return {"nodes": []}
    with open(STATE) as f:
        return json.load(f)

def save_state(s):
    with open(STATE, "w") as f:
        json.dump(s, f, indent=2)

def is_deployed(name):
    s = load_state()
    return any(n["name"] == name for n in s.get("nodes", []))

# ─── Node Operations ───────────────────────────────────────────────────

def install_node(name, token, cloud, port, upstream_ip, upstream_port, labels, mode):
    instance = f"{BASE}/{name}/wallarm"
    api_host = "us1.api.wallarm.com" if cloud == "US" else "api.wallarm.com"
    address = f"0.0.0.0:{port}"
    upstream = f"{upstream_ip}:{upstream_port}"

    # Download AIO
    if not os.path.exists(AIO_PATH):
        print(f"[{name}] Downloading installer...")
        dl = "curl -fsSL" if run_ok("which curl") else "wget -q"
        run(f"{dl} -o {AIO_PATH} {AIO_URL}")
        os.chmod(AIO_PATH, 0o755)

    # Prepare /opt/wallarm
    if os.path.islink(DEPLOY_TO) or os.path.isfile(DEPLOY_TO):
        os.remove(DEPLOY_TO)
    shutil.rmtree(DEPLOY_TO, ignore_errors=True)
    os.makedirs(DEPLOY_TO, exist_ok=True)

    # Install NGINX
    nginx_bin = f"{DEPLOY_TO}/nginx/sbin/nginx"
    if not os.path.exists(nginx_bin):
        os.makedirs(os.path.dirname(nginx_bin), exist_ok=True)
        sys_nginx = shutil.which("nginx")
        if sys_nginx:
            shutil.copy(sys_nginx, nginx_bin)
        else:
            run("apt-get install -y -qq nginx 2>/dev/null || yum install -y -q nginx 2>/dev/null", timeout=60)
            if os.path.exists("/usr/sbin/nginx"):
                shutil.copy("/usr/sbin/nginx", nginx_bin)

    # Extract AIO
    print(f"[{name}] Extracting...")
    run(f"bash {AIO_PATH} --noexec --keep --target {DEPLOY_TO} --noprogress --accept 2>&1", timeout=120)

    # Nginx config
    nginx_dir = f"{DEPLOY_TO}/nginx"
    os.makedirs(f"{nginx_dir}/conf", exist_ok=True)
    with open(f"{nginx_dir}/conf/nginx.conf", "w") as f:
        f.write(f"""load_module {DEPLOY_TO}/modules/nginx_v1.26.3_s0ff5dffff/ngx_http_wallarm_module.so;

worker_processes auto;
pid {nginx_dir}/nginx.pid;
error_log {nginx_dir}/error.log;
events {{ worker_connections 10240; }}
http {{
    access_log {nginx_dir}/access.log;
    wallarm_mode {mode};
    server {{
        listen {port};
        server_name _;
        location /wallarm-status {{ wallarm_status on; allow 127.0.0.0/8; deny all; }}
        location / {{ proxy_pass http://{upstream}; proxy_set_header Host $host; proxy_set_header X-Real-IP $remote_addr; }}
    }}
}}
""")
    run(f"{nginx_bin} -c {nginx_dir}/conf/nginx.conf", timeout=5)

    # Pick module + register
    run(f"bash {DEPLOY_TO}/pick-module.sh 2>&1", timeout=10)
    print(f"[{name}] Registering node...")
    r = run(f"bash -c 'source {DEPLOY_TO}/env.list 2>/dev/null; {DEPLOY_TO}/register-node job:register -token {token} -host {api_host}'",
            timeout=120)
    # Check for success indicators
    registered = "node instance registered" in r.stdout or "init done" in r.stdout
    access_denied = "access denied" in r.stdout + r.stderr
    invalid_token = "invalid token" in (r.stdout + r.stderr).lower()

    if access_denied or invalid_token:
        raise RuntimeError("Token rejected by Wallarm. Check token permissions and cloud region.")
    if not registered:
        raise RuntimeError(f"Registration failed:\n{r.stdout[-500:]}\n{r.stderr[-500:]}")

    # Kill deploy nginx + move to instance
    run(f"pkill -9 -f '{DEPLOY_TO}/nginx' 2>/dev/null; true", timeout=5)
    print(f"[{name}] Installing...")
    shutil.rmtree(instance, ignore_errors=True)
    os.makedirs(os.path.dirname(instance), exist_ok=True)
    try:
        os.rename(DEPLOY_TO, instance)
    except OSError:
        run(f"cp -a {DEPLOY_TO} {instance}", timeout=30)
        shutil.rmtree(DEPLOY_TO, ignore_errors=True)

    # Patch paths + ports
    offset = hash_port(instance)
    for pattern in ("*.sh", "*.list", "*.conf", "*.yaml", "*.yml"):
        for f in Path(instance).rglob(pattern):
            content = f.read_text()
            content = content.replace("/opt/wallarm", instance)
            for port_num in (3313, 6388, 9001, 8088, 9667, 6060, 5005):
                content = content.replace(f":{port_num}", f":{port_num + offset}")
            f.write_text(content)

    # Add -cfg flag to wcli so it reads its own node.yaml
    for conf in Path(f"{instance}/etc").glob("*.conf"):
        c = conf.read_text()
        if "wcli run" in c and "-cfg" not in c:
            conf.write_text(c.replace("wcli run", f"wcli run -cfg {instance}/etc/wallarm/node.yaml"))

    # Systemd
    tmpl = f"""[Unit]
Description=Wallarm Node - %i
After=network.target
[Service]
Type=simple
WorkingDirectory={BASE}/%i/wallarm
EnvironmentFile=-{BASE}/%i/wallarm/env.list
ExecStartPre=/bin/ln -sf {BASE}/%i/wallarm /opt/wallarm
ExecStartPre=-{BASE}/%i/wallarm/nginx/sbin/nginx -c {BASE}/%i/wallarm/nginx/conf/nginx.conf
ExecStartPre=/bin/sleep 1
ExecStart={BASE}/%i/wallarm/usr/bin/python3.10 {BASE}/%i/wallarm/usr/bin/supervisord -c {BASE}/%i/wallarm/etc/supervisord.conf
ExecStop={BASE}/%i/wallarm/usr/bin/python3.10 {BASE}/%i/wallarm/usr/bin/supervisord -c {BASE}/%i/wallarm/etc/supervisord.conf shutdown
Restart=on-failure
RestartSec=5
User=root
[Install]
WantedBy=multi-user.target
"""
    with open("/etc/systemd/system/wallarm-node@.service", "w") as f:
        f.write(tmpl)
    run("systemctl daemon-reload")
    run(f"systemctl enable wallarm-node@{name}")
    run(f"systemctl start wallarm-node@{name}")
    print(f"✅ {name} deployed. systemctl status wallarm-node@{name}")
    return True

# ─── Remove ────────────────────────────────────────────────────────────

def remove_node(name):
    run(f"systemctl stop wallarm-node@{name} 2>/dev/null")
    run(f"systemctl disable wallarm-node@{name} 2>/dev/null")
    shutil.rmtree(f"{BASE}/{name}", ignore_errors=True)
    s = load_state()
    s["nodes"] = [n for n in s["nodes"] if n["name"] != name]
    save_state(s)
    print(f"✅ {name} removed.")

# ─── Status ────────────────────────────────────────────────────────────

def show_status():
    s = load_state()
    if not s.get("nodes"):
        print("No nodes deployed.")
        return
    print("\n─── Node Status ───")
    for n in s["nodes"]:
        r = run(f"systemctl is-active wallarm-node@{n['name']}")
        status = "●" if r.stdout.strip() == "active" else "○"
        print(f"  {status} {n['name']} — {r.stdout.strip()} — {n.get('address','')}")

# ─── Tunnel ────────────────────────────────────────────────────────────

def start_tunnel():
    host = input("Jumphost URL [ssh.sechpoint.app:443]: ") or "ssh.sechpoint.app:443"
    user = input("Username [wallarm-tunnel]: ") or "wallarm-tunnel"
    pw = input("Password or SSH key path: ")
    if not pw:
        print("No credentials.")
        return
    print(f"\nShare: ssh -p 9042 {user}@{host}\nCtrl+C to close.\n")
    # Import tunnel module on demand
    try:
        from tunnel import start_tunnel_ssh
        start_tunnel_ssh(host, user, pw)
    except ImportError:
        print("Tunnel module not available. Install paramiko: pip install paramiko")

# ─── Deploy Menu ───────────────────────────────────────────────────────

def deploy_menu():
    cfg = load_config()
    nodes = cfg.get("nodes", {})
    if not nodes:
        print("No nodes in /opt/fw/fw.conf")
        return

    deployed = {n["name"] for n in load_state().get("nodes", [])}
    available = []
    print("\n─── Deploy Node ───")
    for name, nc in nodes.items():
        if name in deployed:
            print(f"  ✓  {name} (port {nc.get('port','?')}) [deployed]")
        else:
            available.append(name)
            print(f"  [{len(available)}] {name} (port {nc.get('port','?')})")

    if not available:
        print("All nodes deployed.")
        return

    c = input("\nSelect number (Enter=all): ").strip()
    selected = available if not c else [available[int(c)-1]] if c.isdigit() and 1 <= int(c) <= len(available) else []

    if not selected:
        print("Invalid.")
        return

    for name in selected:
        nc = nodes[name]
        print(f"\nDeploying {name}...")
        mode = nc.get("mode", "monitoring")
        print("\nTraffic mode: [1] monitoring [2] safe_blocking [3] block [4] off")
        m = input(f"Choose [{['','1','2','3','4'][['monitoring','safe_blocking','block','off'].index(mode)] if mode in ['monitoring','safe_blocking','block','off'] else '1'}]: ").strip()
        modes = {"2":"safe_blocking","3":"block","4":"off"}
        mode = modes.get(m, mode)

        try:
            install_node(name, nc["token"], nc.get("cloud","EU"),
                        str(nc.get("port","8081")), nc.get("upstream_ip","127.0.0.1"),
                        str(nc.get("upstream_port","80")), nc.get("labels",f"group={name}"), mode)
            s = load_state()
            s.setdefault("nodes", []).append({
                "name": name, "type": "native", "address": f"0.0.0.0:{nc.get('port','8081')}",
                "upstream_ip": nc.get("upstream_ip",""), "upstream_port": int(nc.get("upstream_port",80)),
                "status": "running", "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ")
            })
            save_state(s)
            print(f"✅ {name} deployed.")
        except Exception as e:
            print(f"❌ {name} failed: {e}")

# ─── Edit Menu ─────────────────────────────────────────────────────────

def edit_menu():
    s = load_state()
    if not s.get("nodes"):
        print("No nodes.")
        return
    print("\n─── Edit Node ───")
    for i, n in enumerate(s["nodes"]):
        r = run(f"systemctl is-active wallarm-node@{n['name']}")
        st = "●" if r.stdout.strip() == "active" else "○"
        print(f"  [{i+1}] {st} {n['name']}  {n.get('address','')} → {n.get('upstream_ip','')}:{n.get('upstream_port','')}")
    c = input("\nSelect: ").strip()
    if not c.isdigit(): return
    n = s["nodes"][int(c)-1]

    port = input(f"Port [{n['address'].split(':')[-1]}]: ").strip()
    ip = input(f"Upstream IP [{n.get('upstream_ip','')}]: ").strip()
    up = input(f"Upstream port [{n.get('upstream_port','')}]: ").strip()
    if not any([port, ip, up]):
        print("No changes.")
        return
    if port: n["address"] = f"0.0.0.0:{port}"
    if ip: n["upstream_ip"] = ip
    if up: n["upstream_port"] = int(up)
    save_state(s)

    # Update nginx config
    nginx_conf = f"{BASE}/{n['name']}/wallarm/nginx/conf/nginx.conf"
    if os.path.exists(nginx_conf):
        c = Path(nginx_conf).read_text()
        if port: c = c.replace(f"listen {n['address'].split(':')[-1]};", f"listen {port};")
        if ip or up:
            old_up = f"{n.get('upstream_ip','')}:{n.get('upstream_port','')}"
            new_up = f"{n.get('upstream_ip','')}:{n.get('upstream_port','')}"
            c = c.replace(f"http://{old_up}", f"http://{new_up}")
        Path(nginx_conf).write_text(c)
    print(f"✅ Updated. Restart: systemctl restart wallarm-node@{n['name']}")

# ─── Remove Menu ───────────────────────────────────────────────────────

def remove_menu():
    s = load_state()
    if not s.get("nodes"):
        print("No nodes.")
        return
    print("\n─── Remove Node ───")
    for i, n in enumerate(s["nodes"]):
        r = run(f"systemctl is-active wallarm-node@{n['name']}")
        st = "●" if r.stdout.strip() == "active" else "○"
        print(f"  [{i+1}] {st} {n['name']} {n.get('address','')}")
    c = input("\nSelect: ").strip()
    if not c.isdigit(): return
    n = s["nodes"][int(c)-1]
    remove_node(n["name"])

# ─── Main ──────────────────────────────────────────────────────────────

def main():
    if len(sys.argv) > 1 and sys.argv[1] in ("--version", "-v"):
        print(f"deploy.py v{VERSION}")
        return
    if len(sys.argv) > 1 and sys.argv[1] in ("--help", "-h", "help"):
        print(f"deploy.py v{VERSION} — Wallarm Node Manager")
        print("  deploy.py              Interactive menu")
        print("  deploy.py --deploy-all  Deploy all nodes from fw.conf")
        print("  deploy.py --version     Show version")
        return
    if len(sys.argv) > 1 and sys.argv[1] == "--deploy-all":
        deploy_all()
        return

    print(f"═══ Wallarm Deployment Manager ═══")
    print(f"   v{VERSION}\n")
    if not preflight():
        print("\n❌ Preflight failed.")
        sys.exit(1)
    print("✅ Preflight passed.\n")

    while True:
        print("\n─── Menu ───")
        print("  [1] Deploy  [2] Edit  [3] Status  [4] Remove  [5] Tunnel  [q] Quit")
        c = input("\nChoice: ").strip()
        if c == "1": deploy_menu()
        elif c == "2": edit_menu()
        elif c == "3": show_status()
        elif c == "4": remove_menu()
        elif c == "5": start_tunnel()
        elif c.lower() == "q": break
        # Clean up stray symlinks
        if os.path.islink("/opt/wallarm"):
            os.remove("/opt/wallarm")

def deploy_all():
    print("Deploying all from fw.conf...")
    cfg = load_config()
    if not cfg.get("nodes"):
        print("No nodes found in /opt/fw/fw.conf")
        return
    for name, nc in cfg.get("nodes", {}).items():
        if is_deployed(name):
            print(f"✓ {name} already deployed")
            continue
        print(f"\nDeploying {name}...")
        install_node(name, nc["token"], nc.get("cloud","EU"),
                    str(nc.get("port","8081")), nc.get("upstream_ip","127.0.0.1"),
                    str(nc.get("upstream_port","80")), nc.get("labels",f"group={name}"),
                    nc.get("mode","monitoring"))
    # Clean up
    if os.path.islink("/opt/wallarm"):
        os.remove("/opt/wallarm")
    print("Done.")

if __name__ == "__main__":
    main()