#!/usr/bin/env python3
import sys
import os
import json
import shutil
import tempfile
import subprocess
import time
import signal
import socket
import threading
from pathlib import Path

# --- Configuration & Defaults ---
CONFIG_FILE = Path.home() / ".config" / "distributed-llama.json"
DEFAULT_MODELS_DIR = Path.home() / "models"

# Binary paths: override via env vars or persisted config
DEFAULT_LLAMA_SERVER_BIN = os.getenv("LLAMA_SERVER_BIN", "llama-server")
DEFAULT_RPC_SERVER_BIN = os.getenv("LLAMA_RPC_SERVER_BIN", "rpc-server")

REMOTE_SSH_PORT = os.getenv("REMOTE_PORT", "22")
RPC_PORT = int(os.getenv("RPC_PORT", "50052"))
LOCAL_HOST_PORT = os.getenv("LOCAL_HOST_PORT", "8080")

MODES = ["llama-server", "llama-cli", "llama-bench"]
DEFAULT_MODE = "llama-server"

DEFAULT_HOSTS: list = []

KV_CACHE_QUANT_VALUES = ("q8_0", "q5_1", "q5_0", "q4_1", "q4_0", "iq4_nl")
KV_CACHE_OPTIONS = {
    "off": "Disabled (full precision)",
    "q8_0": "Q8_0 (recommended)",
    "q5_1": "Q5_1",
    "q5_0": "Q5_0",
    "q4_1": "Q4_1",
    "q4_0": "Q4_0 (aggressive)",
    "iq4_nl": "IQ4_NL",
}

HEALTH_CHECK_INTERVAL = 10   # seconds between liveness polls
RESTART_COOLDOWN = 5         # seconds to wait before attempting restart
PORT_WAIT_TIMEOUT = 30       # seconds to wait for a (re)started node's port


# --- Helpers ---

def check_dependencies():
    missing = [t for t in ("dialog", "ssh") if not shutil.which(t)]
    if missing:
        print(f"Error: missing required tools: {', '.join(missing)}")
        sys.exit(1)


def run_dialog(args):
    with tempfile.NamedTemporaryFile(mode="w+") as tf:
        try:
            res = subprocess.run(["dialog"] + args, stderr=tf, check=False)
            tf.seek(0)
            return tf.read().strip(), res.returncode
        except Exception:
            return None, 1


def show_msg(title, msg):
    run_dialog(["--title", title, "--msgbox", msg, "10", "60"])


def _port_open(ip, port, timeout=2):
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(timeout)
        s.connect((ip, port))
        s.close()
        return True
    except Exception:
        return False


def _wait_for_port(ip, port, stop_event, timeout=PORT_WAIT_TIMEOUT):
    for _ in range(timeout):
        if stop_event.is_set():
            return False
        if _port_open(ip, port):
            return True
        time.sleep(1)
    return False


# --- File Picker (GGUF only) ---

def _dir_contents(path):
    try:
        entries = os.listdir(path)
    except PermissionError:
        return [], []
    dirs = sorted(e for e in entries if os.path.isdir(os.path.join(path, e)))
    files = sorted(e for e in entries if e.endswith(".gguf"))
    return dirs, files


def custom_file_picker(start_path):
    current = os.path.abspath(start_path if os.path.isdir(start_path) else os.path.dirname(start_path))
    while True:
        dirs, files = _dir_contents(current)
        items = []
        if current != "/":
            items += ["..", "Parent Directory"]
        for d in dirs:
            items += [d + "/", "<DIR>"]
        for f in files:
            items += [f, "<GGUF>"]
        if not items:
            items += [".", "Empty Directory"]

        pretty = current if len(current) <= 50 else "..." + current[-47:]
        sel, code = run_dialog([
            "--title", "Select GGUF File",
            "--backtitle", f"Current: {pretty}",
            "--menu", "Navigate and select a .gguf file:", "20", "70", "12",
            *items,
        ])
        if code != 0:
            return None
        sel = sel.strip()
        if sel == "..":
            current = os.path.dirname(current)
        elif sel.endswith("/"):
            current = os.path.join(current, sel[:-1])
        elif sel != ".":
            return os.path.join(current, sel)


# --- Node Monitor (health check + auto-restart) ---

class NodeMonitor(threading.Thread):
    """Watches one RPC node and restarts it via SSH if the port goes dark."""

    def __init__(self, ip, rpc_port, ssh_port, remote_bin, stop_event):
        super().__init__(daemon=True, name=f"monitor-{ip}")
        self.ip = ip
        self.rpc_port = rpc_port
        self.ssh_port = ssh_port
        self.remote_bin = remote_bin   # path to rpc-server on the remote machine
        self.stop_event = stop_event
        self.restart_count = 0
        self.status = "ok"             # "ok" | "down" | "restarting" | "unreachable"

    # --- internals ---

    def _ssh_start(self):
        """Kill any stale rpc-server, start a fresh one, return PID string or None."""
        script = (
            f"pkill -9 -f '{self.remote_bin}' 2>/dev/null || true\n"
            f"sleep 1\n"
            f"nohup {self.remote_bin} -H 0.0.0.0 -p {self.rpc_port} "
            f"> /tmp/rpc-server-{self.ip}.log 2>&1 < /dev/null &\n"
            f"echo $!\n"
        )
        res = subprocess.run(
            [
                "ssh", "-p", self.ssh_port,
                "-o", "StrictHostKeyChecking=no",
                "-o", "ConnectTimeout=10",
                "-o", "BatchMode=yes",
                self.ip, "bash -s",
            ],
            input=script, text=True, capture_output=True,
        )
        if res.returncode != 0:
            return None
        for line in reversed(res.stdout.splitlines()):
            if line.strip().isdigit():
                return line.strip()
        return None

    def _attempt_restart(self):
        self.status = "restarting"
        print(f"\n[MONITOR] {self.ip}: restarting RPC server...", flush=True)
        pid = self._ssh_start()
        if not pid:
            self.status = "unreachable"
            print(f"[MONITOR] {self.ip}: SSH failed — node unreachable.", flush=True)
            return
        print(f"[MONITOR] {self.ip}: started (PID {pid}), waiting for port...", flush=True)
        if _wait_for_port(self.ip, self.rpc_port, self.stop_event):
            self.restart_count += 1
            self.status = "ok"
            print(f"[MONITOR] {self.ip}: back online (restart #{self.restart_count})", flush=True)
        else:
            self.status = "unreachable"
            print(f"[MONITOR] {self.ip}: port did not open after restart.", flush=True)

    # --- thread entry point ---

    def run(self):
        while not self.stop_event.wait(HEALTH_CHECK_INTERVAL):
            if not _port_open(self.ip, self.rpc_port):
                self.status = "down"
                print(f"\n[MONITOR] {self.ip}:{self.rpc_port} is DOWN.", flush=True)
                time.sleep(RESTART_COOLDOWN)
                if not self.stop_event.is_set():
                    self._attempt_restart()
            else:
                self.status = "ok"


# --- AppState ---

class AppState:
    def __init__(self):
        self.model_path = ""
        self.mode = DEFAULT_MODE
        self.hosts: list = list(DEFAULT_HOSTS)   # [[ip, enabled], ...]
        self.context_size = None
        self.bench_prefill = "512,8192,16384,32768,65536"
        self.bench_gen = "128"
        self.kv_cache_quant = None
        self.extra_args = ""
        self.bench_extra_args = ""
        self.llama_server_bin = DEFAULT_LLAMA_SERVER_BIN
        self.rpc_server_bin = DEFAULT_RPC_SERVER_BIN
        self.load_config()

    @property
    def active_hosts(self):
        return [h[0] for h in self.hosts if h[1]]

    def save_config(self):
        data = {
            "model_path": self.model_path,
            "mode": self.mode,
            "hosts": self.hosts,
            "context_size": self.context_size,
            "bench_prefill": self.bench_prefill,
            "bench_gen": self.bench_gen,
            "kv_cache_quant": self.kv_cache_quant,
            "extra_args": self.extra_args,
            "bench_extra_args": self.bench_extra_args,
            "llama_server_bin": self.llama_server_bin,
            "rpc_server_bin": self.rpc_server_bin,
        }
        try:
            CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True)
            CONFIG_FILE.write_text(json.dumps(data, indent=2))
        except OSError:
            pass

    def load_config(self):
        if not CONFIG_FILE.is_file():
            return
        try:
            data = json.loads(CONFIG_FILE.read_text())
        except (json.JSONDecodeError, OSError):
            return
        if not isinstance(data, dict):
            return

        mp = data.get("model_path", "")
        if isinstance(mp, str) and mp and os.path.isfile(mp):
            self.model_path = mp

        md = data.get("mode")
        if isinstance(md, str) and md in MODES:
            self.mode = md

        hs = data.get("hosts")
        if isinstance(hs, list) and hs:
            valid = [
                e for e in hs
                if isinstance(e, list) and len(e) == 2
                and isinstance(e[0], str) and isinstance(e[1], bool)
            ]
            if valid:
                self.hosts = valid

        cs = data.get("context_size")
        if cs is None or (isinstance(cs, int) and cs > 0):
            self.context_size = cs

        kv = data.get("kv_cache_quant")
        if kv is None or (isinstance(kv, str) and kv in KV_CACHE_QUANT_VALUES):
            self.kv_cache_quant = kv

        for attr in ("bench_prefill", "bench_gen", "extra_args", "bench_extra_args"):
            v = data.get(attr)
            if isinstance(v, str):
                setattr(self, attr, v)

        for attr in ("llama_server_bin", "rpc_server_bin"):
            v = data.get(attr)
            if isinstance(v, str) and v:
                setattr(self, attr, v)


# --- TUI actions ---

def select_model(state):
    start = state.model_path if state.model_path else str(DEFAULT_MODELS_DIR)
    if os.path.isfile(start):
        start = os.path.dirname(start)
    sel = custom_file_picker(start)
    if sel:
        state.model_path = sel


def select_mode(state):
    items = []
    for m in MODES:
        items += [m, ""]
    sel, code = run_dialog([
        "--title", "Select Execution Mode",
        "--menu", "Choose how to run the model:", "12", "50", "5",
        *items,
    ])
    if code == 0 and sel:
        state.mode = sel


def select_context(state):
    if state.mode == "llama-bench":
        sel_p, code = run_dialog([
            "--title", "Bench Prefill Sizes (pp)",
            "--inputbox",
            "Comma-separated prompt sizes (e.g. 512,8192,16384).\nLeave empty to skip:",
            "10", "68", str(state.bench_prefill or ""),
        ])
        if code == 0:
            state.bench_prefill = sel_p.strip() or None
            sel_n, code2 = run_dialog([
                "--title", "Bench Token Generation (tg)",
                "--inputbox",
                "Comma-separated generation lengths (e.g. 128).\nLeave empty to skip:",
                "10", "68", str(state.bench_gen or ""),
            ])
            if code2 == 0:
                state.bench_gen = sel_n.strip() or None
    else:
        sel, code = run_dialog([
            "--title", "Context Size",
            "--inputbox",
            "Enter context size (e.g. 4096, 8192).\nLeave empty for model default:",
            "10", "60", str(state.context_size or ""),
        ])
        if code == 0:
            state.context_size = int(sel.strip()) if sel.strip().isdigit() else None


def select_kv_cache(state):
    items = []
    for k, v in KV_CACHE_OPTIONS.items():
        items += [k, v]
    sel, code = run_dialog([
        "--title", "KV Cache Quantization",
        "--menu", "Quantize the KV cache to reduce memory usage:", "14", "55", "5",
        *items,
    ])
    if code == 0 and sel:
        state.kv_cache_quant = None if sel == "off" else sel


def edit_extra_args(state):
    is_bench = state.mode == "llama-bench"
    current = state.bench_extra_args if is_bench else state.extra_args
    label = "(Bench)" if is_bench else "(Server/CLI)"
    sel, code = run_dialog([
        "--title", f"Extra Arguments {label}",
        "--inputbox",
        f"Extra CLI args for {state.mode} (appended as-is).\nLeave empty for none:",
        "12", "65", current,
    ])
    if code == 0:
        if is_bench:
            state.bench_extra_args = sel.strip()
        else:
            state.extra_args = sel.strip()


def edit_binaries(state):
    sel_s, code = run_dialog([
        "--title", "Local llama-server Binary",
        "--inputbox", "Path or name of the local llama-server binary:", "10", "65",
        state.llama_server_bin,
    ])
    if code == 0 and sel_s.strip():
        state.llama_server_bin = sel_s.strip()

    sel_r, code = run_dialog([
        "--title", "Remote rpc-server Binary",
        "--inputbox",
        "Path or name of rpc-server on the REMOTE machines.\n"
        "(e.g. /opt/llama/rpc-server or just rpc-server if in PATH):",
        "11", "65", state.rpc_server_bin,
    ])
    if code == 0 and sel_r.strip():
        state.rpc_server_bin = sel_r.strip()


def add_server(state):
    sel, code = run_dialog(["--title", "Add Server", "--inputbox", "Enter new server IP address:", "10", "50"])
    if code == 0 and sel.strip():
        state.hosts.append([sel.strip(), True])


def remove_server(state):
    if not state.hosts:
        show_msg("Info", "No servers to remove.")
        return
    items = []
    for i, (ip, _) in enumerate(state.hosts):
        items += [str(i), ip]
    sel, code = run_dialog(["--title", "Remove Server", "--menu", "Select server to remove:", "15", "50", "5", *items])
    if code == 0 and sel:
        idx = int(sel)
        if 0 <= idx < len(state.hosts):
            del state.hosts[idx]


def edit_server(state):
    if not state.hosts:
        show_msg("Info", "No servers to edit.")
        return
    items = []
    for i, (ip, _) in enumerate(state.hosts):
        items += [str(i), ip]
    sel, code = run_dialog(["--title", "Edit Server", "--menu", "Select server to edit:", "15", "50", "5", *items])
    if code == 0 and sel:
        idx = int(sel)
        if 0 <= idx < len(state.hosts):
            new_ip, code2 = run_dialog([
                "--title", "Edit Server IP",
                "--inputbox", "Enter new IP address:", "10", "50", state.hosts[idx][0],
            ])
            if code2 == 0 and new_ip.strip():
                state.hosts[idx][0] = new_ip.strip()


def toggle_servers(state):
    if not state.hosts:
        show_msg("Info", "No servers to configure. Add some first.")
        return
    items = []
    for i, (ip, enabled) in enumerate(state.hosts):
        items += [str(i), ip, "on" if enabled else "off"]
    sel_str, code = run_dialog([
        "--title", "Toggle Active Servers",
        "--checklist", "Space to toggle active nodes:", "15", "50", "5",
        *items,
    ])
    if code == 0:
        for h in state.hosts:
            h[1] = False
        if sel_str:
            for x in sel_str.split():
                idx = int(x.strip('"'))
                if 0 <= idx < len(state.hosts):
                    state.hosts[idx][1] = True


def configure_servers(state):
    while True:
        menu = ["1", "Toggle Active Servers", "2", "Add Server", "3", "Remove Server", "4", "Edit Server", "5", "Back"]
        sel, code = run_dialog(["--title", "Manage Remote Servers", "--menu", "Choose an action:", "15", "50", "5", *menu])
        if code != 0 or sel == "5":
            break
        {"1": toggle_servers, "2": add_server, "3": remove_server, "4": edit_server}.get(sel, lambda s: None)(state)


# --- Core: start one RPC node and return its PID ---

def _ssh_start_rpc(ip, rpc_port, ssh_port, remote_bin):
    script = (
        f"pkill -9 -f '{remote_bin}' 2>/dev/null || true\n"
        f"sleep 1\n"
        f"nohup {remote_bin} -H 0.0.0.0 -p {rpc_port} "
        f"> /tmp/rpc-server-{ip}.log 2>&1 < /dev/null &\n"
        f"echo $!\n"
    )
    res = subprocess.run(
        [
            "ssh", "-p", ssh_port,
            "-o", "StrictHostKeyChecking=no",
            "-o", "ConnectTimeout=10",
            "-o", "BatchMode=yes",
            ip, "bash -s",
        ],
        input=script, text=True, capture_output=True,
    )
    if res.returncode != 0:
        return None, res.stderr.strip()
    for line in reversed(res.stdout.splitlines()):
        if line.strip().isdigit():
            return line.strip(), None
    return None, f"unexpected output: {res.stdout!r}"


def _kill_rpc(ip, ssh_port, remote_bin, pid=None):
    parts = []
    if pid:
        parts.append(f"kill -9 {pid} 2>/dev/null || true")
    parts.append(f"pkill -9 -f '{remote_bin}' 2>/dev/null || true")
    subprocess.run(
        ["ssh", "-p", ssh_port, "-o", "ConnectTimeout=5", "-o", "BatchMode=yes", ip, "; ".join(parts)],
        stderr=subprocess.DEVNULL,
    )


# --- Run ---

def run_distributed(state):
    if not state.model_path or not os.path.exists(state.model_path):
        show_msg("Error", f"Model file not found:\n{state.model_path}")
        return
    if not state.active_hosts:
        show_msg("Error", "No remote servers selected.")
        return

    subprocess.run(["clear"])
    print("=== Distributed llama.cpp ===")
    print(f"Model:      {state.model_path}")
    print(f"Mode:       {state.mode}")
    print(f"Local bin:  {state.llama_server_bin}")
    print(f"Remote bin: {state.rpc_server_bin}")
    if state.mode == "llama-bench":
        print(f"Prefill:    {state.bench_prefill or '-'}")
        print(f"Gen:        {state.bench_gen or '-'}")
    else:
        print(f"Context:    {state.context_size or 'default'}")
    print(f"KV Cache:   {state.kv_cache_quant or 'off'}")
    print(f"Hosts:      {state.active_hosts}")
    print("---")

    active_ips = state.active_hosts
    remote_pids: dict[str, str | None] = {}
    stop_event = threading.Event()
    monitors: list[NodeMonitor] = []

    def cleanup():
        print("\nStopping monitors and cleaning up remote RPC servers...")
        stop_event.set()
        for ip in active_ips:
            pid = remote_pids.get(ip)
            print(f"  Killing RPC on {ip} (PID: {pid or '?'})...")
            _kill_rpc(ip, REMOTE_SSH_PORT, state.rpc_server_bin, pid)

    def on_signal(sig, frame):
        cleanup()
        sys.exit(0)

    signal.signal(signal.SIGINT, on_signal)
    signal.signal(signal.SIGTERM, on_signal)

    try:
        rpc_endpoints = []

        # 1. Start RPC servers on all remote nodes
        for ip in active_ips:
            print(f"-> Starting RPC server on {ip}...", end=" ", flush=True)
            pid, err = _ssh_start_rpc(ip, RPC_PORT, REMOTE_SSH_PORT, state.rpc_server_bin)
            if not pid:
                print(f"FAILED ({err})")
                cleanup()
                return
            remote_pids[ip] = pid
            print(f"PID {pid}", end=" ", flush=True)

            print(f"| waiting for :{RPC_PORT}...", end=" ", flush=True)
            if not _wait_for_port(ip, RPC_PORT, stop_event):
                print("TIMEOUT")
                cleanup()
                return
            print("OK")
            rpc_endpoints.append(f"{ip}:{RPC_PORT}")

        rpc_arg = ",".join(rpc_endpoints)
        print(f"All nodes ready. RPC: {rpc_arg}")

        # 2. Start health monitor threads
        for ip in active_ips:
            m = NodeMonitor(ip, RPC_PORT, REMOTE_SSH_PORT, state.rpc_server_bin, stop_event)
            m.start()
            monitors.append(m)
        print(f"[MONITOR] Watching {len(monitors)} node(s) every {HEALTH_CHECK_INTERVAL}s.")
        print("---")

        # 3. Build and run local command
        import shlex
        cmd = [state.llama_server_bin, "-m", state.model_path, "--rpc", rpc_arg]

        if state.mode == "llama-server":
            cmd += ["--host", "0.0.0.0", "--port", LOCAL_HOST_PORT]
            if state.context_size:
                cmd += ["-c", str(state.context_size)]

        elif state.mode == "llama-cli":
            cmd += ["-cnv", "-p", "You are a helpful assistant."]
            if state.context_size:
                cmd += ["-c", str(state.context_size)]

        elif state.mode == "llama-bench":
            if state.bench_prefill:
                cmd += ["-p", state.bench_prefill.strip()]
            if state.bench_gen:
                cmd += ["-n", state.bench_gen.strip()]

        if state.kv_cache_quant:
            cmd += ["--cache-type-k", state.kv_cache_quant, "--cache-type-v", state.kv_cache_quant]

        extra = state.bench_extra_args if state.mode == "llama-bench" else state.extra_args
        if extra:
            cmd += shlex.split(extra)

        print(f"CMD: {' '.join(cmd)}")
        print("---")
        proc = subprocess.Popen(cmd)
        proc.wait()

    except Exception as e:
        print(f"\n[ERROR] {e}")
    finally:
        cleanup()

    input("\nDone. Press Enter to return to menu...")


# --- Main Menu ---

def main_menu():
    state = AppState()

    while True:
        model_display = Path(state.model_path).name if state.model_path else "(None)"
        servers_display = f"{len(state.active_hosts)} active"

        if state.mode == "llama-bench":
            ctx_display = f"pp=[{state.bench_prefill or '-'}] tg={state.bench_gen or '-'}"
            ctx_label = "Bench:    "
            run_label = "RUN BENCHMARK"
        else:
            ctx_display = str(state.context_size) if state.context_size else "default"
            ctx_label = "Context:  "
            run_label = "RUN DISTRIBUTED SERVER"

        kv_display = state.kv_cache_quant or "off"
        extra = (state.bench_extra_args if state.mode == "llama-bench" else state.extra_args) or "(none)"

        menu_args = [
            "--clear", "--backtitle", "Distributed llama.cpp",
            "--title", "Main Menu",
            "--menu", "Configure and launch distributed inference:", "24", "70", "11",
            "1", f"Model:    {model_display}",
            "2", f"Servers:  {servers_display}",
            "3", f"Mode:     {state.mode}",
            "4", f"{ctx_label}{ctx_display}",
            "5", f"KV Cache: {kv_display}",
            "6", f"Extra:    {extra}",
            "7", f"Binaries: local={Path(state.llama_server_bin).name}  remote={Path(state.rpc_server_bin).name}",
            "8", run_label,
            "9", "Exit",
        ]

        choice, code = run_dialog(menu_args)
        if code != 0:
            break

        {
            "1": select_model,
            "2": configure_servers,
            "3": select_mode,
            "4": select_context,
            "5": select_kv_cache,
            "6": edit_extra_args,
            "7": edit_binaries,
            "8": run_distributed,
        }.get(choice, lambda s: None)(state)

        if choice == "9":
            break

        state.save_config()

    state.save_config()
    subprocess.run(["clear"])


if __name__ == "__main__":
    check_dependencies()
    try:
        main_menu()
    except KeyboardInterrupt:
        subprocess.run(["clear"])
        sys.exit(0)