#!/usr/bin/env python3
"""
VayDNS Terminal Management Menu
Run: python3 menu.py  OR  vaydns
"""

import os
import sys
import json
import signal
import subprocess
import time
import threading
import configparser
from pathlib import Path

# ── Try to import rich, fallback to plain output ──────────────────────────────
try:
    from rich.console import Console
    from rich.table import Table
    from rich.panel import Panel
    from rich.prompt import Prompt, Confirm
    from rich.text import Text
    from rich.live import Live
    from rich.columns import Columns
    from rich import box
    HAS_RICH = True
except ImportError:
    HAS_RICH = False

VAYDNS_DIR   = Path("/opt/vaydns")
CONFIG_FILE  = VAYDNS_DIR / "vaydns.conf"
PID_SERVER   = VAYDNS_DIR / "vaydns-server.pid"
PID_CLIENT   = VAYDNS_DIR / "vaydns-client.pid"
LOG_SERVER   = VAYDNS_DIR / "logs/server.log"
LOG_CLIENT   = VAYDNS_DIR / "logs/client.log"
BIN_SERVER   = VAYDNS_DIR / "vaydns-server"
BIN_CLIENT   = VAYDNS_DIR / "vaydns-client"
KEY_DIR      = VAYDNS_DIR / "keys"
SYSTEMD_SRV  = Path("/etc/systemd/system/vaydns-server.service")
SYSTEMD_CLI  = Path("/etc/systemd/system/vaydns-client.service")

(VAYDNS_DIR / "logs").mkdir(parents=True, exist_ok=True)

console = Console() if HAS_RICH else None

# ── Config helpers ─────────────────────────────────────────────────────────────

DEFAULT_CONFIG = {
    "server": {
        "domain": "",
        "listen_port": "5300",
        "upstream": "127.0.0.1:8000",
        "privkey_file": str(KEY_DIR / "server.key"),
        "mtu": "1232",
        "record_type": "txt",
        "idle_timeout": "10s",
        "keepalive": "2s",
        "log_level": "info",
    },
    "client": {
        "domain": "",
        "listen": "127.0.0.1:7000",
        "pubkey_file": str(KEY_DIR / "server.pub"),
        "transport": "doh",
        "doh_url": "https://doh.cloudflare.com/dns-query",
        "dot_addr": "",
        "udp_addr": "8.8.8.8:53",
        "record_type": "txt",
        "idle_timeout": "10s",
        "keepalive": "2s",
        "log_level": "info",
        "utls": "weighted",
    },
    "proxy": {
        "mode": "none",
        "proxy_type": "socks5",
    },
}

def load_config() -> configparser.ConfigParser:
    cfg = configparser.ConfigParser()
    for section, values in DEFAULT_CONFIG.items():
        cfg[section] = values
    if CONFIG_FILE.exists():
        cfg.read(CONFIG_FILE)
    return cfg

def save_config(cfg: configparser.ConfigParser):
    with open(CONFIG_FILE, "w") as f:
        cfg.write(f)

# ── Process helpers ────────────────────────────────────────────────────────────

def read_pid(pid_file: Path) -> int | None:
    try:
        return int(pid_file.read_text().strip())
    except Exception:
        return None

def is_running(pid_file: Path) -> bool:
    pid = read_pid(pid_file)
    if pid is None:
        return False
    try:
        os.kill(pid, 0)
        return True
    except ProcessLookupError:
        return False
    except PermissionError:
        return True

def stop_process(pid_file: Path, name: str):
    pid = read_pid(pid_file)
    if pid is None:
        print_warn(f"{name} is not running (no PID file)")
        return
    try:
        os.kill(pid, signal.SIGTERM)
        time.sleep(1)
        try:
            os.kill(pid, 0)
            os.kill(pid, signal.SIGKILL)
            print_warn(f"{name} force-killed (SIGKILL)")
        except ProcessLookupError:
            print_ok(f"{name} stopped")
    except ProcessLookupError:
        print_warn(f"{name} was not running")
    finally:
        pid_file.unlink(missing_ok=True)

def start_process(cmd: list[str], pid_file: Path, log_file: Path, name: str):
    log_file.parent.mkdir(parents=True, exist_ok=True)
    with open(log_file, "a") as lf:
        proc = subprocess.Popen(
            cmd,
            stdout=lf,
            stderr=lf,
            stdin=subprocess.DEVNULL,
            start_new_session=True,
        )
    pid_file.write_text(str(proc.pid))
    time.sleep(0.8)
    if is_running(pid_file):
        print_ok(f"{name} started (PID {proc.pid})")
    else:
        print_err(f"{name} failed to start — check {log_file}")

# ── Output helpers ─────────────────────────────────────────────────────────────

def print_ok(msg):
    if HAS_RICH:
        console.print(f"  [green][+][/green] {msg}")
    else:
        print(f"  [+] {msg}")

def print_warn(msg):
    if HAS_RICH:
        console.print(f"  [yellow][!][/yellow] {msg}")
    else:
        print(f"  [!] {msg}")

def print_err(msg):
    if HAS_RICH:
        console.print(f"  [red][✗][/red] {msg}")
    else:
        print(f"  [✗] {msg}")

def print_info(msg):
    if HAS_RICH:
        console.print(f"  [cyan]•[/cyan] {msg}")
    else:
        print(f"  • {msg}")

def clear():
    os.system("clear" if os.name != "nt" else "cls")

def pause():
    input("\n  Press ENTER to continue...")

# ── Banner ─────────────────────────────────────────────────────────────────────

BANNER = r"""
 __   __           ___  _  _ ___
 \ \ / /__ _ _  _ |   \| \| / __|
  \ V / _ \ || || || |) | .` \__ \
   \_/\___/\_|| ||_|___/|_|\_|___/
             |__/
"""

def print_banner():
    srv_status = "[green]● RUNNING[/green]" if is_running(PID_SERVER) else "[red]○ STOPPED[/red]"
    cli_status = "[green]● RUNNING[/green]" if is_running(PID_CLIENT) else "[red]○ STOPPED[/red]"
    if HAS_RICH:
        console.print(f"[cyan]{BANNER}[/cyan]")
        console.print(
            Panel(
                f"  Server: {srv_status}    Client: {cli_status}",
                title="[bold]VayDNS Management[/bold]",
                subtitle="[dim]vaydns.orx.ma[/dim]",
                border_style="cyan",
                expand=False,
            )
        )
    else:
        print(BANNER)
        srv = "● RUNNING" if is_running(PID_SERVER) else "○ STOPPED"
        cli = "● RUNNING" if is_running(PID_CLIENT) else "○ STOPPED"
        print(f"  Server: {srv}    Client: {cli}")
        print("  ─────────────────────────────────")

# ── Menu rendering ─────────────────────────────────────────────────────────────

def render_menu(title: str, options: list[tuple[str, str]]):
    if HAS_RICH:
        table = Table(box=box.SIMPLE, show_header=False, padding=(0, 2))
        table.add_column("Key", style="bold cyan", width=6)
        table.add_column("Action", style="white")
        for key, label in options:
            table.add_row(f"[{key}]", label)
        console.print(Panel(table, title=f"[bold]{title}[/bold]", border_style="cyan"))
    else:
        print(f"\n  ── {title} ──")
        for key, label in options:
            print(f"    [{key}] {label}")
        print()

def prompt(text: str, default: str = "") -> str:
    if HAS_RICH:
        return Prompt.ask(f"  [bold cyan]{text}[/bold cyan]", default=default)
    else:
        val = input(f"  {text} [{default}]: ").strip()
        return val if val else default

def confirm(text: str, default: bool = False) -> bool:
    if HAS_RICH:
        return Confirm.ask(f"  [bold yellow]{text}[/bold yellow]", default=default)
    else:
        ans = input(f"  {text} [{'Y/n' if default else 'y/N'}]: ").strip().lower()
        if not ans:
            return default
        return ans.startswith("y")

def get_choice(options: list[str]) -> str:
    valid = {o[0].lower() for o in options}
    valid.add("q")
    while True:
        ch = input("  Choice: ").strip().lower()
        if ch in valid:
            return ch
        print_warn("Invalid choice — try again")

# ── Server commands ────────────────────────────────────────────────────────────

def build_server_cmd(cfg: configparser.ConfigParser) -> list[str]:
    s = cfg["server"]
    cmd = [
        str(BIN_SERVER),
        "-udp", f":{s['listen_port']}",
        "-domain", s["domain"],
        "-upstream", s["upstream"],
        "-privkey-file", s["privkey_file"],
        "-mtu", s["mtu"],
        "-record-type", s["record_type"],
        "-idle-timeout", s["idle_timeout"],
        "-keepalive", s["keepalive"],
        "-log-level", s["log_level"],
    ]
    return cmd

def build_client_cmd(cfg: configparser.ConfigParser) -> list[str]:
    c = cfg["client"]
    cmd = [str(BIN_CLIENT)]
    transport = c["transport"]
    if transport == "doh":
        cmd += ["-doh", c["doh_url"]]
    elif transport == "dot":
        cmd += ["-dot", c["dot_addr"]]
    elif transport == "udp":
        cmd += ["-udp", c["udp_addr"]]
    cmd += [
        "-domain", c["domain"],
        "-listen", c["listen"],
        "-pubkey-file", c["pubkey_file"],
        "-record-type", c["record_type"],
        "-idle-timeout", c["idle_timeout"],
        "-keepalive", c["keepalive"],
        "-log-level", c["log_level"],
    ]
    if transport == "doh" and c.get("utls", "weighted") != "none":
        cmd += ["-utls", c.get("utls", "weighted")]
    return cmd

# ── Sections ───────────────────────────────────────────────────────────────────

def section_status():
    clear()
    print_banner()
    cfg = load_config()
    if HAS_RICH:
        t = Table(title="Process Status", box=box.ROUNDED, border_style="cyan")
        t.add_column("Component", style="bold")
        t.add_column("Status")
        t.add_column("PID")
        t.add_column("Log")
        for name, pid_f, log_f in [
            ("vaydns-server", PID_SERVER, LOG_SERVER),
            ("vaydns-client", PID_CLIENT, LOG_CLIENT),
        ]:
            running = is_running(pid_f)
            pid     = read_pid(pid_f) or "—"
            status  = "[green]● RUNNING[/green]" if running else "[red]○ STOPPED[/red]"
            console.print()
            t.add_row(name, status, str(pid), str(log_f))
        console.print(t)
        console.print()
        t2 = Table(title="Configuration", box=box.ROUNDED, border_style="cyan")
        t2.add_column("Key", style="cyan")
        t2.add_column("Value", style="white")
        s = cfg["server"]
        c = cfg["client"]
        t2.add_row("Server domain",   s.get("domain", "—"))
        t2.add_row("Server port",     s.get("listen_port", "5300"))
        t2.add_row("Server upstream", s.get("upstream", "—"))
        t2.add_row("Client domain",   c.get("domain", "—"))
        t2.add_row("Client listen",   c.get("listen", "127.0.0.1:7000"))
        t2.add_row("Client transport",c.get("transport", "doh"))
        console.print(t2)
    else:
        for name, pid_f, log_f in [
            ("vaydns-server", PID_SERVER, LOG_SERVER),
            ("vaydns-client", PID_CLIENT, LOG_CLIENT),
        ]:
            running = is_running(pid_f)
            pid = read_pid(pid_f) or "—"
            status = "● RUNNING" if running else "○ STOPPED"
            print(f"  {name}: {status}  PID={pid}  Log={log_f}")
    pause()

def section_server_menu():
    cfg = load_config()
    while True:
        clear()
        print_banner()
        opts = [
            ("s", "Start server"),
            ("k", "Stop server"),
            ("r", "Restart server"),
            ("c", "Configure server"),
            ("g", "Generate new keypair"),
            ("l", "View server logs"),
            ("i", "Install as systemd service"),
            ("u", "Uninstall systemd service"),
            ("b", "Back"),
        ]
        render_menu("Server Management", opts)
        ch = get_choice(opts)

        if ch == "s":
            cfg = load_config()
            if not cfg["server"]["domain"]:
                print_warn("Domain not configured. Please configure first.")
                pause()
                continue
            if is_running(PID_SERVER):
                print_warn("Server already running")
            else:
                cmd = build_server_cmd(cfg)
                print_info(f"Command: {' '.join(cmd)}")
                start_process(cmd, PID_SERVER, LOG_SERVER, "vaydns-server")
            pause()

        elif ch == "k":
            stop_process(PID_SERVER, "vaydns-server")
            pause()

        elif ch == "r":
            stop_process(PID_SERVER, "vaydns-server")
            time.sleep(1)
            cfg = load_config()
            if cfg["server"]["domain"]:
                cmd = build_server_cmd(cfg)
                start_process(cmd, PID_SERVER, LOG_SERVER, "vaydns-server")
            pause()

        elif ch == "c":
            configure_server(cfg)
            cfg = load_config()

        elif ch == "g":
            generate_keys()
            pause()

        elif ch == "l":
            view_logs(LOG_SERVER, "Server Logs")

        elif ch == "i":
            install_systemd_server(cfg)
            pause()

        elif ch == "u":
            uninstall_systemd("vaydns-server", SYSTEMD_SRV)
            pause()

        elif ch in ("b", "q"):
            break

def section_client_menu():
    cfg = load_config()
    while True:
        clear()
        print_banner()
        opts = [
            ("s", "Start client"),
            ("k", "Stop client"),
            ("r", "Restart client"),
            ("c", "Configure client"),
            ("l", "View client logs"),
            ("t", "Test connection"),
            ("i", "Install as systemd service"),
            ("u", "Uninstall systemd service"),
            ("b", "Back"),
        ]
        render_menu("Client Management", opts)
        ch = get_choice(opts)

        if ch == "s":
            cfg = load_config()
            if not cfg["client"]["domain"]:
                print_warn("Domain not configured. Please configure first.")
                pause()
                continue
            if is_running(PID_CLIENT):
                print_warn("Client already running")
            else:
                cmd = build_client_cmd(cfg)
                print_info(f"Command: {' '.join(cmd)}")
                start_process(cmd, PID_CLIENT, LOG_CLIENT, "vaydns-client")
            pause()

        elif ch == "k":
            stop_process(PID_CLIENT, "vaydns-client")
            pause()

        elif ch == "r":
            stop_process(PID_CLIENT, "vaydns-client")
            time.sleep(1)
            cfg = load_config()
            if cfg["client"]["domain"]:
                cmd = build_client_cmd(cfg)
                start_process(cmd, PID_CLIENT, LOG_CLIENT, "vaydns-client")
            pause()

        elif ch == "c":
            configure_client(cfg)
            cfg = load_config()

        elif ch == "l":
            view_logs(LOG_CLIENT, "Client Logs")

        elif ch == "t":
            test_connection(cfg)
            pause()

        elif ch == "i":
            install_systemd_client(cfg)
            pause()

        elif ch == "u":
            uninstall_systemd("vaydns-client", SYSTEMD_CLI)
            pause()

        elif ch in ("b", "q"):
            break

def section_tools_menu():
    while True:
        clear()
        print_banner()
        opts = [
            ("k", "Key manager"),
            ("p", "Proxy setup wizard"),
            ("f", "Firewall / iptables helper"),
            ("z", "DNS zone setup guide"),
            ("u", "Update VayDNS binaries"),
            ("b", "Back"),
        ]
        render_menu("Tools", opts)
        ch = get_choice(opts)

        if ch == "k":
            key_manager()
        elif ch == "p":
            proxy_wizard()
        elif ch == "f":
            firewall_helper()
        elif ch == "z":
            dns_zone_guide()
        elif ch == "u":
            update_vaydns()
            pause()
        elif ch in ("b", "q"):
            break

# ── Configure ──────────────────────────────────────────────────────────────────

def configure_server(cfg: configparser.ConfigParser):
    clear()
    if HAS_RICH:
        console.print(Panel("[bold]Configure Server[/bold]", border_style="cyan"))
    else:
        print("\n  ── Configure Server ──")
    s = cfg["server"]
    s["domain"]       = prompt("Tunnel domain (e.g. t.example.com)", s["domain"])
    s["listen_port"]  = prompt("UDP listen port", s["listen_port"])
    s["upstream"]     = prompt("Upstream TCP address (proxy/SSH/etc)", s["upstream"])
    s["privkey_file"] = prompt("Private key file", s["privkey_file"])
    s["mtu"]          = prompt("MTU (1232 recommended)", s["mtu"])
    s["record_type"]  = prompt("DNS record type [txt/null/cname/a/aaaa/mx/ns/srv/caa]", s["record_type"])
    s["idle_timeout"] = prompt("Idle timeout", s["idle_timeout"])
    s["keepalive"]    = prompt("Keepalive interval (< idle-timeout)", s["keepalive"])
    s["log_level"]    = prompt("Log level [debug/info/warning/error]", s["log_level"])
    save_config(cfg)
    print_ok("Server configuration saved")
    pause()

def configure_client(cfg: configparser.ConfigParser):
    clear()
    if HAS_RICH:
        console.print(Panel("[bold]Configure Client[/bold]", border_style="cyan"))
    else:
        print("\n  ── Configure Client ──")
    c = cfg["client"]
    c["domain"]    = prompt("Tunnel domain (must match server)", c["domain"])
    c["listen"]    = prompt("Local listen address", c["listen"])
    c["pubkey_file"] = prompt("Server public key file", c["pubkey_file"])

    print_info("Transport: [1] DoH  [2] DoT  [3] UDP")
    tr = input("  Select transport [1]: ").strip() or "1"
    if tr == "1":
        c["transport"] = "doh"
        c["doh_url"] = prompt("DoH resolver URL", c.get("doh_url", "https://doh.cloudflare.com/dns-query"))
        c["utls"] = prompt("uTLS profile [weighted/Firefox/Chrome/random/none]", c.get("utls", "weighted"))
    elif tr == "2":
        c["transport"] = "dot"
        c["dot_addr"] = prompt("DoT resolver (host:853)", c.get("dot_addr", ""))
    else:
        c["transport"] = "udp"
        c["udp_addr"] = prompt("UDP resolver (host:53)", c.get("udp_addr", "8.8.8.8:53"))

    c["record_type"]  = prompt("DNS record type [txt/null/cname/a/aaaa/mx/ns/srv/caa]", c["record_type"])
    c["idle_timeout"] = prompt("Idle timeout (must match server)", c["idle_timeout"])
    c["keepalive"]    = prompt("Keepalive (must match server, < idle-timeout)", c["keepalive"])
    c["log_level"]    = prompt("Log level [debug/info/warning/error]", c["log_level"])
    save_config(cfg)
    print_ok("Client configuration saved")
    pause()

# ── Key manager ────────────────────────────────────────────────────────────────

def generate_keys():
    KEY_DIR.mkdir(parents=True, exist_ok=True)
    if not BIN_SERVER.exists():
        print_err(f"vaydns-server binary not found at {BIN_SERVER}")
        return
    cmd = [
        str(BIN_SERVER), "-gen-key",
        "-privkey-file", str(KEY_DIR / "server.key"),
        "-pubkey-file",  str(KEY_DIR / "server.pub"),
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode == 0:
        print_ok(f"Keys saved to {KEY_DIR}/")
        if (KEY_DIR / "server.pub").exists():
            pub = (KEY_DIR / "server.pub").read_text().strip()
            if HAS_RICH:
                console.print(f"\n  [bold yellow]Public key:[/bold yellow] [cyan]{pub}[/cyan]")
            else:
                print(f"\n  Public key: {pub}")
    else:
        print_err(f"Key generation failed: {result.stderr}")

def key_manager():
    while True:
        clear()
        opts = [
            ("g", "Generate new keypair"),
            ("s", "Show public key"),
            ("p", "Show private key (⚠ sensitive)"),
            ("c", "Copy pubkey to clipboard"),
            ("b", "Back"),
        ]
        render_menu("Key Manager", opts)
        ch = get_choice(opts)

        if ch == "g":
            if (KEY_DIR / "server.key").exists():
                if not confirm("Overwrite existing keys?", default=False):
                    continue
            generate_keys()
            pause()
        elif ch == "s":
            pub = KEY_DIR / "server.pub"
            if pub.exists():
                if HAS_RICH:
                    console.print(f"\n  [cyan]{pub.read_text().strip()}[/cyan]")
                else:
                    print(f"\n  {pub.read_text().strip()}")
            else:
                print_warn("No public key found. Generate one first.")
            pause()
        elif ch == "p":
            priv = KEY_DIR / "server.key"
            if priv.exists():
                if confirm("⚠  Show private key?", default=False):
                    if HAS_RICH:
                        console.print(f"\n  [red]{priv.read_text().strip()}[/red]")
                    else:
                        print(f"\n  {priv.read_text().strip()}")
            else:
                print_warn("No private key found.")
            pause()
        elif ch == "c":
            pub = KEY_DIR / "server.pub"
            if pub.exists():
                try:
                    subprocess.run(["xclip", "-selection", "clipboard"],
                                   input=pub.read_text(), text=True, check=True)
                    print_ok("Copied to clipboard!")
                except Exception:
                    try:
                        subprocess.run(["xsel", "--clipboard", "--input"],
                                       input=pub.read_text(), text=True, check=True)
                        print_ok("Copied to clipboard!")
                    except Exception:
                        print_warn("Clipboard not available. Key:")
                        print(f"  {pub.read_text().strip()}")
            pause()
        elif ch in ("b", "q"):
            break

# ── Logs ───────────────────────────────────────────────────────────────────────

def view_logs(log_file: Path, title: str):
    clear()
    if HAS_RICH:
        console.print(Panel(f"[bold]{title}[/bold]  [dim]{log_file}[/dim]", border_style="cyan"))
    else:
        print(f"\n  ── {title} ─────────────")
        print(f"  File: {log_file}")
    if not log_file.exists():
        print_warn("Log file not found yet")
        pause()
        return

    try:
        result = subprocess.run(["tail", "-n", "80", str(log_file)],
                                capture_output=True, text=True)
        if HAS_RICH:
            from rich.syntax import Syntax
            console.print(Syntax(result.stdout, "text", theme="monokai",
                                 line_numbers=True, word_wrap=True))
        else:
            print(result.stdout)
    except Exception as e:
        print_err(str(e))

    print_info("Press [f] to follow live, or ENTER to return")
    ch = input("  Choice: ").strip().lower()
    if ch == "f":
        try:
            proc = subprocess.Popen(["tail", "-f", str(log_file)])
            print_info("Following log — press Ctrl+C to stop")
            proc.wait()
        except KeyboardInterrupt:
            proc.terminate()
    return

# ── Test connection ────────────────────────────────────────────────────────────

def test_connection(cfg: configparser.ConfigParser):
    clear()
    listen = cfg["client"].get("listen", "127.0.0.1:7000")
    host, port = listen.rsplit(":", 1)
    print_info(f"Testing TCP connectivity to {listen} ...")
    result = subprocess.run(
        ["bash", "-c", f"echo | timeout 3 nc -z {host} {port} && echo OK || echo FAIL"],
        capture_output=True, text=True,
    )
    if "OK" in result.stdout:
        print_ok(f"Connection to {listen} is OPEN — tunnel is working!")
    else:
        print_warn(f"Cannot connect to {listen} — client may not be running or tunnel is down")

    print_info("Testing via curl (if proxy mode http)...")
    try:
        result2 = subprocess.run(
            ["curl", "-s", "--max-time", "5",
             "--proxy", f"http://{listen}/",
             "https://wtfismyip.com/text"],
            capture_output=True, text=True,
        )
        if result2.returncode == 0 and result2.stdout.strip():
            print_ok(f"Exit IP via tunnel: {result2.stdout.strip()}")
        else:
            print_info("HTTP proxy test inconclusive (may not be in http proxy mode)")
    except Exception:
        pass

# ── Proxy wizard ───────────────────────────────────────────────────────────────

def proxy_wizard():
    clear()
    if HAS_RICH:
        console.print(Panel("[bold]Proxy Setup Wizard[/bold]", border_style="cyan"))
    modes = [
        ("1", "HTTP proxy (ncat)"),
        ("2", "SOCKS5 via SSH (server-side)"),
        ("3", "SOCKS5 via SSH (client-side, private)"),
        ("4", "Tor bridge"),
    ]
    render_menu("Select Proxy Mode", modes)
    ch = input("  Choice [1]: ").strip() or "1"

    upstream = "127.0.0.1:8000"
    cfg = load_config()

    if ch == "1":
        print_ok("HTTP proxy via ncat:")
        print_info("Server side:")
        print(f"    ncat -l -k --proxy-type http 127.0.0.1 8000")
        print(f"    {BIN_SERVER} -udp :{cfg['server']['listen_port']} -privkey-file {cfg['server']['privkey_file']} -domain {cfg['server']['domain'] or 't.example.com'} -upstream {upstream}")
        print_info("Client side:")
        print(f"    {BIN_CLIENT} -doh {cfg['client']['doh_url']} -pubkey-file {cfg['client']['pubkey_file']} -domain {cfg['client']['domain'] or 't.example.com'} -listen {cfg['client']['listen']}")
        print(f"    curl --proxy http://{cfg['client']['listen']}/ https://wtfismyip.com/text")
    elif ch == "2":
        print_ok("SOCKS5 via SSH (server-side):")
        print(f"    ssh -N -D 127.0.0.1:8000 -o NoHostAuthenticationForLocalhost=yes 127.0.0.1")
        print(f"    {BIN_SERVER} -udp :{cfg['server']['listen_port']} -privkey-file {cfg['server']['privkey_file']} -domain {cfg['server']['domain'] or 't.example.com'} -upstream {upstream}")
        print_info("Client side:")
        print(f"    {BIN_CLIENT} -doh {cfg['client']['doh_url']} -pubkey-file {cfg['client']['pubkey_file']} -domain {cfg['client']['domain'] or 't.example.com'} -listen {cfg['client']['listen']}")
        print(f"    curl --proxy socks5h://{cfg['client']['listen']}/ https://wtfismyip.com/text")
    elif ch == "3":
        print_ok("SOCKS5 via SSH (client-side, private):")
        print(f"    {BIN_SERVER} -udp :{cfg['server']['listen_port']} -privkey-file {cfg['server']['privkey_file']} -domain {cfg['server']['domain'] or 't.example.com'} -upstream 127.0.0.1:22")
        print_info("Client side:")
        print(f"    {BIN_CLIENT} -doh {cfg['client']['doh_url']} -pubkey-file {cfg['client']['pubkey_file']} -domain {cfg['client']['domain'] or 't.example.com'} -listen 127.0.0.1:8000")
        print(f"    ssh -N -D 127.0.0.1:7000 -o HostKeyAlias=tunnel-server -p 8000 127.0.0.1")
        print(f"    curl --proxy socks5h://127.0.0.1:7000/ https://wtfismyip.com/text")
    elif ch == "4":
        print_ok("Tor bridge setup:")
        print(f"    {BIN_SERVER} -udp :{cfg['server']['listen_port']} -privkey-file {cfg['server']['privkey_file']} -domain {cfg['server']['domain'] or 't.example.com'} -upstream 127.0.0.1:9001")
        print_info("Client torrc:")
        print(f"    Bridge 127.0.0.1:{cfg['client']['listen'].split(':')[-1]} <FINGERPRINT>")

    pause()

# ── Firewall helper ────────────────────────────────────────────────────────────

def firewall_helper():
    clear()
    cfg = load_config()
    port = cfg["server"].get("listen_port", "5300")
    if HAS_RICH:
        console.print(Panel("[bold]Firewall / iptables Setup[/bold]", border_style="cyan"))
    print_info(f"Redirecting UDP port 53 → {port} (run as root):\n")
    cmds = [
        f"sudo iptables  -I INPUT -p udp --dport {port} -j ACCEPT",
        f"sudo iptables  -t nat -I PREROUTING -i eth0 -p udp --dport 53 -j REDIRECT --to-ports {port}",
        f"sudo ip6tables -I INPUT -p udp --dport {port} -j ACCEPT",
        f"sudo ip6tables -t nat -I PREROUTING -i eth0 -p udp --dport 53 -j REDIRECT --to-ports {port}",
    ]
    for c in cmds:
        print(f"    {c}")
    print()
    if confirm("Apply these rules now?", default=False):
        for c in cmds:
            result = subprocess.run(c, shell=True, capture_output=True)
            if result.returncode == 0:
                print_ok(c)
            else:
                print_err(f"Failed: {c}  — {result.stderr.decode()}")
    pause()

# ── DNS zone guide ─────────────────────────────────────────────────────────────

def dns_zone_guide():
    clear()
    cfg = load_config()
    domain = cfg["server"].get("domain") or "t.example.com"
    base = domain.split(".", 1)[1] if "." in domain else "example.com"
    tns = f"tns.{base}"
    if HAS_RICH:
        console.print(Panel("[bold]DNS Zone Setup[/bold]", border_style="cyan"))
        t = Table(box=box.ROUNDED)
        t.add_column("Type")
        t.add_column("Name")
        t.add_column("Value")
        t.add_row("A",  tns,    "YOUR_SERVER_IP")
        t.add_row("AAAA", tns,  "YOUR_SERVER_IPv6  (optional)")
        t.add_row("NS", domain, tns)
        console.print(t)
    else:
        print(f"\n  DNS Records needed:")
        print(f"    A     {tns}   YOUR_SERVER_IP")
        print(f"    AAAA  {tns}   YOUR_SERVER_IPv6  (optional)")
        print(f"    NS    {domain}  {tns}")
    print_info("Queries for *.{} will be forwarded to your tunnel server".format(domain))
    print_info("Keep tunnel subdomain short to maximize payload space")
    pause()

# ── Systemd ────────────────────────────────────────────────────────────────────

def install_systemd_server(cfg: configparser.ConfigParser):
    cmd = " ".join(build_server_cmd(cfg))
    svc = f"""[Unit]
Description=VayDNS Server
After=network.target

[Service]
ExecStart={cmd}
Restart=on-failure
RestartSec=5
StandardOutput=append:{LOG_SERVER}
StandardError=append:{LOG_SERVER}

[Install]
WantedBy=multi-user.target
"""
    try:
        SYSTEMD_SRV.write_text(svc)
        subprocess.run(["systemctl", "daemon-reload"], check=True)
        subprocess.run(["systemctl", "enable", "vaydns-server"], check=True)
        print_ok("vaydns-server systemd service installed and enabled")
        print_info("Start with: systemctl start vaydns-server")
    except Exception as e:
        print_err(f"Systemd install failed: {e}")

def install_systemd_client(cfg: configparser.ConfigParser):
    cmd = " ".join(build_client_cmd(cfg))
    svc = f"""[Unit]
Description=VayDNS Client
After=network.target

[Service]
ExecStart={cmd}
Restart=on-failure
RestartSec=5
StandardOutput=append:{LOG_CLIENT}
StandardError=append:{LOG_CLIENT}

[Install]
WantedBy=multi-user.target
"""
    try:
        SYSTEMD_CLI.write_text(svc)
        subprocess.run(["systemctl", "daemon-reload"], check=True)
        subprocess.run(["systemctl", "enable", "vaydns-client"], check=True)
        print_ok("vaydns-client systemd service installed and enabled")
        print_info("Start with: systemctl start vaydns-client")
    except Exception as e:
        print_err(f"Systemd install failed: {e}")

def uninstall_systemd(name: str, svc_file: Path):
    try:
        subprocess.run(["systemctl", "stop", name], capture_output=True)
        subprocess.run(["systemctl", "disable", name], capture_output=True)
        svc_file.unlink(missing_ok=True)
        subprocess.run(["systemctl", "daemon-reload"], check=True)
        print_ok(f"{name} service removed")
    except Exception as e:
        print_err(str(e))

# ── Update ─────────────────────────────────────────────────────────────────────

def update_vaydns():
    clear()
    print_info("Pulling latest VayDNS source...")
    src = VAYDNS_DIR / "src"
    if not src.exists():
        print_err("Source directory not found. Run setup.sh first.")
        return
    try:
        subprocess.run(["git", "-C", str(src), "pull"], check=True)
        was_srv = is_running(PID_SERVER)
        was_cli = is_running(PID_CLIENT)
        stop_process(PID_SERVER, "vaydns-server")
        stop_process(PID_CLIENT, "vaydns-client")
        cfg = load_config()
        print_info("Rebuilding binaries...")
        subprocess.run(["go", "build", "-o", str(BIN_SERVER), "./vaydns-server"],
                       cwd=str(src), check=True)
        subprocess.run(["go", "build", "-o", str(BIN_CLIENT), "./vaydns-client"],
                       cwd=str(src), check=True)
        print_ok("Binaries updated!")
        if was_srv:
            start_process(build_server_cmd(cfg), PID_SERVER, LOG_SERVER, "vaydns-server")
        if was_cli:
            start_process(build_client_cmd(cfg), PID_CLIENT, LOG_CLIENT, "vaydns-client")
    except subprocess.CalledProcessError as e:
        print_err(f"Update failed: {e}")

# ── Main menu ──────────────────────────────────────────────────────────────────

def main_menu():
    while True:
        clear()
        print_banner()
        opts = [
            ("1", "Status overview"),
            ("2", "Server management"),
            ("3", "Client management"),
            ("4", "Tools & Utilities"),
            ("q", "Quit"),
        ]
        render_menu("Main Menu", opts)
        ch = get_choice(opts)

        if ch == "1":
            section_status()
        elif ch == "2":
            section_server_menu()
        elif ch == "3":
            section_client_menu()
        elif ch == "4":
            section_tools_menu()
        elif ch == "q":
            if HAS_RICH:
                console.print("\n  [cyan]Goodbye![/cyan]\n")
            else:
                print("\n  Goodbye!\n")
            sys.exit(0)

# ── CLI flags ──────────────────────────────────────────────────────────────────

def cli_dispatch():
    args = sys.argv[1:]
    cfg = load_config()
    if "--server" in args:
        if is_running(PID_SERVER):
            print_warn("Server already running")
        else:
            start_process(build_server_cmd(cfg), PID_SERVER, LOG_SERVER, "vaydns-server")
    elif "--client" in args:
        if is_running(PID_CLIENT):
            print_warn("Client already running")
        else:
            start_process(build_client_cmd(cfg), PID_CLIENT, LOG_CLIENT, "vaydns-client")
    elif "--stop-server" in args:
        stop_process(PID_SERVER, "vaydns-server")
    elif "--stop-client" in args:
        stop_process(PID_CLIENT, "vaydns-client")
    elif "--status" in args:
        for name, pid_f in [("vaydns-server", PID_SERVER), ("vaydns-client", PID_CLIENT)]:
            status = "RUNNING" if is_running(pid_f) else "STOPPED"
            pid    = read_pid(pid_f) or "—"
            print(f"  {name}: {status}  PID={pid}")
    elif "--gen-key" in args:
        generate_keys()
    else:
        main_menu()

if __name__ == "__main__":
    cli_dispatch()