#!/usr/bin/env python3
"""Keep SSH reverse tunnel: VPS 127.0.0.1:8765 -> box serve_dashboard :8765.

Caddy on the Thai VPS serves https://lighter.160.236.128.88.sslip.io → VPS:8765.
Without this -R forward, the public dashboard dies even if local serve is up.
Does NOT touch SOCKS 10808, Agogo 8877, or kill serve_dashboard/guard_dashboard_port.
"""
from __future__ import annotations

import fcntl
import json
import os
import re
import signal
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
DATA = ROOT / "data"
LOCAL_PORT = 8765
REMOTE_BIND = "127.0.0.1"
INTERVAL_S = 20
LOCK = DATA / "dashboard_reverse.guard.lock"
PIDFILE = DATA / "dashboard_reverse_guard.pid"
EVENTS = DATA / "dashboard_reverse_guard_events.jsonl"
LOG = DATA / "dashboard_reverse_guard.log"
URL_FILE = DATA / "dashboard_public_url.txt"
PUBLIC_URL = "https://lighter.160.236.128.88.sslip.io/dashboard.html"

SSH_KEY = Path("/home/box/.ssh/mexc_vps")
SSH_HOST = "root@160.236.128.88"
SSH_HOST_MARKER = "160.236.128.88"
R_MARKER = f"-R {REMOTE_BIND}:{LOCAL_PORT}:{REMOTE_BIND}:{LOCAL_PORT}"


def utc_now() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def log(msg: str) -> None:
    DATA.mkdir(parents=True, exist_ok=True)
    line = f"[{utc_now()}] {msg}"
    print(line, flush=True)
    with LOG.open("a") as f:
        f.write(line + "\n")


def emit(kind: str, detail: str) -> None:
    DATA.mkdir(parents=True, exist_ok=True)
    with EVENTS.open("a") as f:
        f.write(
            json.dumps({"ts": utc_now(), "kind": kind, "detail": detail}, ensure_ascii=False)
            + "\n"
        )
    log(f"{kind}: {detail}")


def cmdline_of(pid: int) -> str:
    try:
        return (
            Path(f"/proc/{pid}/cmdline")
            .read_bytes()
            .replace(b"\0", b" ")
            .decode("utf-8", "replace")
            .strip()
        )
    except Exception:
        return "?"


def list_ssh_reverse_procs() -> list[tuple[int, str]]:
    """Local ssh processes that reverse-forward 8765 to the VPS."""
    out: list[tuple[int, str]] = []
    try:
        for p in Path("/proc").iterdir():
            if not p.name.isdigit():
                continue
            cmd = cmdline_of(int(p.name))
            if not is_legit_reverse(cmd):
                continue
            out.append((int(p.name), cmd))
    except Exception as e:
        emit("proc_scan_fail", str(e))
    return out


def is_legit_reverse(cmd: str) -> bool:
    if not cmd or cmd == "?" or "guard_dashboard_reverse" in cmd:
        return False
    if re.search(r"(^|/|\s)ssh(\s|$)", cmd) is None:
        return False
    if SSH_HOST_MARKER not in cmd:
        return False
    if "-R" not in cmd:
        return False
    # accept spaced or compact forms mentioning 8765 both sides
    if f":{LOCAL_PORT}:{REMOTE_BIND}:{LOCAL_PORT}" in cmd:
        return True
    if f":{LOCAL_PORT}:127.0.0.1:{LOCAL_PORT}" in cmd:
        return True
    # -R 8765:127.0.0.1:8765
    if re.search(rf"-R\s+({REMOTE_BIND}:)?{LOCAL_PORT}:{REMOTE_BIND}:{LOCAL_PORT}", cmd):
        return True
    return False


def decide_action(local_ok: bool, remote_ok: bool) -> str:
    """keep | recreate"""
    if local_ok and remote_ok:
        return "keep"
    return "recreate"


def remote_dashboard_ok() -> bool:
    """Probe VPS localhost:8765/dashboard.html via a short ssh."""
    if not SSH_KEY.exists():
        return False
    cmd = [
        "ssh",
        "-i",
        str(SSH_KEY),
        "-o",
        "BatchMode=yes",
        "-o",
        "ConnectTimeout=8",
        SSH_HOST,
        "curl -sS -o /dev/null -w '%{http_code}' --max-time 4 http://127.0.0.1:8765/dashboard.html",
    ]
    try:
        r = subprocess.run(cmd, capture_output=True, text=True, check=False, timeout=20)
    except Exception as e:
        emit("remote_probe_fail", str(e))
        return False
    code = (r.stdout or "").strip()
    return code == "200"


def local_dashboard_ok() -> bool:
    try:
        r = subprocess.run(
            [
                "curl",
                "-sS",
                "-o",
                "/dev/null",
                "-w",
                "%{http_code}",
                "--max-time",
                "4",
                f"http://127.0.0.1:{LOCAL_PORT}/dashboard.html",
            ],
            capture_output=True,
            text=True,
            check=False,
            timeout=8,
        )
    except Exception as e:
        emit("local_probe_fail", str(e))
        return False
    return (r.stdout or "").strip() == "200"


def ssh_reverse_cmd() -> list[str]:
    return [
        "ssh",
        "-i",
        str(SSH_KEY),
        "-o",
        "BatchMode=yes",
        "-o",
        "ConnectTimeout=10",
        "-o",
        "ServerAliveInterval=15",
        "-o",
        "ServerAliveCountMax=3",
        "-o",
        "ExitOnForwardFailure=yes",
        "-f",
        "-N",
        "-R",
        f"{REMOTE_BIND}:{LOCAL_PORT}:{REMOTE_BIND}:{LOCAL_PORT}",
        SSH_HOST,
    ]


def kill_pids(pids: list[int], reason: str) -> None:
    for pid in pids:
        emit("kill", f"reason={reason} pid={pid}")
        try:
            os.kill(pid, signal.SIGTERM)
        except ProcessLookupError:
            pass
        except Exception as e:
            emit("kill_fail", f"pid={pid} {e}")
    if pids:
        time.sleep(0.5)
        for pid in pids:
            try:
                os.kill(pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
            except Exception:
                pass
        time.sleep(0.3)


def start_reverse() -> bool:
    if not SSH_KEY.exists():
        emit("missing_key", str(SSH_KEY))
        return False
    if not local_dashboard_ok():
        emit("local_down", "serve_dashboard not answering; skip reverse start")
        return False
    cmd = ssh_reverse_cmd()
    emit("start_reverse", " ".join(cmd))
    try:
        r = subprocess.run(cmd, capture_output=True, text=True, check=False, timeout=25)
    except Exception as e:
        emit("start_fail", str(e))
        return False
    if r.returncode != 0:
        emit("start_fail", f"rc={r.returncode} err={(r.stderr or '')[:200]}")
        return False
    time.sleep(1.2)
    ok = remote_dashboard_ok()
    emit("reverse_up" if ok else "reverse_still_down", f"remote_ok={ok}")
    if ok:
        URL_FILE.write_text(PUBLIC_URL + "\n" + utc_now() + "\n")
    return ok


def tick() -> None:
    procs = list_ssh_reverse_procs()
    local_ok = local_dashboard_ok()
    remote_ok = remote_dashboard_ok() if local_ok else False
    # if remote ok, local ssh list may still be empty after box reboot of scan — treat remote as source of truth
    local_ssh_ok = bool(procs)
    action = decide_action(local_ssh_ok or remote_ok, remote_ok)
    if action == "keep":
        return
    emit(
        "recreate_needed",
        f"local_ssh={local_ssh_ok} remote_ok={remote_ok} local_dash={local_ok} n_ssh={len(procs)}",
    )
    if procs:
        kill_pids([p for p, _ in procs], "stale_or_broken_reverse")
    # also clear remote listener if orphaned (sshd holding -R from dead client)
    if not remote_ok and local_ok:
        # starting a new -R usually fails if port still held; try free remote port
        subprocess.run(
            [
                "ssh",
                "-i",
                str(SSH_KEY),
                "-o",
                "BatchMode=yes",
                "-o",
                "ConnectTimeout=8",
                SSH_HOST,
                "fuser -k 8765/tcp 2>/dev/null || true",
            ],
            capture_output=True,
            text=True,
            check=False,
            timeout=15,
        )
        time.sleep(0.4)
    start_reverse()


def main() -> int:
    DATA.mkdir(parents=True, exist_ok=True)
    lock_f = LOCK.open("w")
    try:
        fcntl.flock(lock_f.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
    except BlockingIOError:
        print("another dashboard reverse guard holds the lock; exit", flush=True)
        return 0
    PIDFILE.write_text(str(os.getpid()))
    emit("guard_start", f"interval={INTERVAL_S}s public={PUBLIC_URL}")
    URL_FILE.write_text(PUBLIC_URL + "\n" + utc_now() + "\n")

    def _stop(*_a: object) -> None:
        emit("guard_stop", "signal")
        sys.exit(0)

    signal.signal(signal.SIGTERM, _stop)
    signal.signal(signal.SIGINT, _stop)

    # ensure once at boot
    tick()
    while True:
        time.sleep(INTERVAL_S)
        try:
            tick()
        except Exception as e:
            emit("tick_error", str(e))


if __name__ == "__main__":
    raise SystemExit(main())