#!/usr/bin/env python3
"""Helper for ssh-menu.sh: parse and edit ~/.ssh/config.

Host metadata lives in comments directly above a `Host` line:
    # group: Proxmox
    # env: prod
    Host pm
"""
import os
import re
import shutil
import socket
import subprocess
import sys
import time
import zlib
from collections import Counter
from concurrent.futures import ThreadPoolExecutor

HOME = os.path.expanduser("~")
CFG = os.path.join(HOME, ".ssh", "config")
STATE = os.path.join(os.environ.get("XDG_STATE_HOME", HOME + "/.local/state"), "ssh-menu")
HIST = os.path.join(STATE, "history")
CACHE = os.path.join(STATE, "status.tsv")
BACKUPS = os.path.join(STATE, "backups")
os.makedirs(BACKUPS, exist_ok=True)

G, R, Y, D, B, X = "\x1b[32m", "\x1b[31m", "\x1b[33m", "\x1b[2m", "\x1b[1;31m", "\x1b[0m"


def read_lines():
    try:
        return open(CFG).read().split("\n")
    except FileNotFoundError:
        return []


def backup():
    if os.path.exists(CFG):
        shutil.copy2(CFG, os.path.join(BACKUPS, "config-" + time.strftime("%Y%m%d-%H%M%S")))
        old = sorted(os.listdir(BACKUPS))[:-20]
        for f in old:
            os.remove(os.path.join(BACKUPS, f))


def write_lines(lines):
    backup()
    text = "\n".join(lines)
    if not text.endswith("\n"):
        text += "\n"
    with open(CFG, "w") as f:
        f.write(text)
    os.chmod(CFG, 0o600)


def parse():
    hosts, meta = [], {}
    for i, l in enumerate(read_lines()):
        m = re.match(r"^\s*#\s*(group|env|url)\s*:\s*(.+?)\s*$", l, re.I)
        if m:
            meta[m.group(1).lower()] = m.group(2)
            continue
        m = re.match(r"^\s*host\s+(.+?)\s*$", l, re.I)
        if m:
            for a in m.group(1).split():
                if not re.search(r"[*?!]", a):
                    hosts.append(dict(alias=a, group=meta.get("group", ""), env=meta.get("env", ""), url=meta.get("url", ""), line=i + 1))
            meta = {}
        elif l.strip() == "" or not l.lstrip().startswith("#"):
            meta = {}
    return hosts


def find(alias):
    return next((h for h in parse() if h["alias"] == alias), None)


def resolve(alias):
    """ssh -G -> dict(user, hostname, port, proxyjump, identity)"""
    info = dict(user="", hostname=alias, port="22", proxyjump="", identity="")
    try:
        out = subprocess.run(["ssh", "-G", alias], capture_output=True, text=True, timeout=3).stdout
    except Exception:
        return info
    ids = []
    for l in out.splitlines():
        k, _, v = l.partition(" ")
        if k == "user": info["user"] = v
        elif k == "hostname": info["hostname"] = v
        elif k == "port": info["port"] = v
        elif k == "proxyjump": info["proxyjump"] = v
        elif k == "identityfile": ids.append(os.path.expanduser(v))
    info["identity"] = next((p for p in ids if os.path.exists(p)), "")
    return info


def probe(info):
    if info["proxyjump"]:
        return "jump"
    try:
        socket.create_connection((info["hostname"], int(info["port"])), timeout=0.6).close()
        return "up"
    except Exception:
        return "down"


def addr(info):
    return f'{info["user"]}@{info["hostname"]}:{info["port"]}'


def load_cache():
    c = {}
    try:
        for l in open(CACHE):
            a, ad, st = l.rstrip("\n").split("\t")
            c[a] = (ad, st)
    except Exception:
        pass
    return c


def uses():
    try:
        return Counter(open(HIST).read().split())
    except FileNotFoundError:
        return Counter()


PALETTE = [(248, 187, 217), (125, 207, 255), (158, 206, 106), (224, 175, 104), (187, 154, 247), (115, 218, 202)]


def rgb(r, g, b, bold=False):
    return f"\x1b[{'1;' if bold else ''}38;2;{r};{g};{b}m"


FIXED = {"proxmox": (248, 187, 217), "services": (125, 207, 255), "git": (158, 206, 106), "legacy": (130, 120, 125)}


def group_color(name):
    if name.lower() in FIXED:
        return FIXED[name.lower()]
    return PALETTE[zlib.crc32(name.lower().encode()) % len(PALETTE)]


def group_cell(name, width):
    if not name:
        return " " * width
    return rgb(*group_color(name)) + f"{name:<{width}}" + X


def addr_cell(ad):
    """user@host:port with the host highlighted and the rest dimmed"""
    m = re.match(r"^(.*)@(.*):(\d+)$", ad or "")
    if not m:
        return D + (ad or "") + X
    u, h, p = m.groups()
    return f"{rgb(120, 108, 112)}{u}@{X}{rgb(205, 195, 200)}{h}{X}{rgb(120, 108, 112)}" + ("" if p == "22" else f":{p}") + X


def cmd_list(live):
    hosts = parse()
    cnt = uses()
    hosts.sort(key=lambda h: (-cnt.get(h["alias"], 0), h["alias"].lower()))
    if live:
        def work(h):
            i = resolve(h["alias"])
            return h["alias"], addr(i), probe(i)
        with ThreadPoolExecutor(32) as ex:
            rows = list(ex.map(work, hosts))
        with open(CACHE, "w") as f:
            for a, ad, st in rows:
                f.write(f"{a}\t{ad}\t{st}\n")
        cache = {a: (ad, st) for a, ad, st in rows}
    else:
        cache = load_cache()
    dots = {"up": G + "●" + X, "down": "\x1b[1;31m●" + X, "jump": Y + "●" + X}
    for h in hosts:
        ad, st = cache.get(h["alias"], ("", "?"))
        dot = dots.get(st, D + "·" + X)
        prod = rgb(255, 92, 108, bold=True) + "PROD" + X if h["env"].lower() == "prod" else "    "
        print(f'{dot}  {h["alias"]:<24} {group_cell(h["group"], 11)}  {prod}    {addr_cell(ad)}')


def cmd_preview(alias):
    lab = lambda t: rgb(120, 108, 112) + f"{t:<7}" + X
    val = lambda t: rgb(235, 225, 230) + str(t) + X
    h = find(alias)
    if not h:
        print(f"{alias}\n\n{D}not in ~/.ssh/config:\nwill connect as typed{X}")
        return
    i = resolve(alias)
    st = probe(i)
    col = {"up": (120, 235, 90), "down": (255, 92, 108), "jump": (224, 175, 104)}[st]
    prod = "  " + rgb(255, 92, 108, bold=True) + "PROD" + X if h["env"].lower() == "prod" else ""
    print(rgb(*col, bold=True) + "● " + alias + X + prod)
    print(f'{D}{addr(i)}{X}\n')
    if i["proxyjump"]:
        print(lab("via") + val(i["proxyjump"]))
    print(lab("key") + val(os.path.basename(i["identity"]) or "default / agent"))
    if h["group"]:
        print(lab("group") + rgb(*group_color(h["group"])) + h["group"] + X)
    if h["url"]:
        print(lab("web") + val(h["url"]))
    print(lab("used") + val(f'{uses().get(alias, 0)}x'))
    if st != "jump":
        try:
            out = subprocess.run(["ping", "-c1", "-W1", i["hostname"]], capture_output=True, text=True, timeout=2).stdout
            m = re.search(r"time=([\d.]+)", out)
            print(lab("ping") + (val(m.group(1) + " ms") if m else rgb(255, 92, 108) + "no reply" + X))
        except Exception:
            print(lab("ping") + "-")
    print(f"\n{D}~/.ssh/config:{h['line']}{X}")
    for l in read_lines()[h["line"] - 1:h["line"] + 12]:
        if l.strip() == "" or (re.match(r"^\s*host\s", l, re.I) and not l.lower().lstrip().startswith("host " + alias.lower())):
            break
        print(D + l + X)


def host_index(lines, alias):
    for i, l in enumerate(lines):
        m = re.match(r"^\s*host\s+(.+?)\s*$", l, re.I)
        if m and alias in m.group(1).split():
            return i
    return -1


def set_meta(alias, key, value):
    lines = read_lines()
    idx = host_index(lines, alias)
    if idx < 0:
        sys.exit(f"{alias}: not found")
    j = idx
    while j > 0 and lines[j - 1].lstrip().startswith("#"):
        j -= 1
    block = [l for l in lines[j:idx] if not re.match(rf"^\s*#\s*{key}\s*:", l, re.I)]
    if value:
        block.append(f"# {key}: {value}")
    lines[j:idx] = block
    write_lines(lines)


def cmd_add(alias, hostname, user, port, jump, group, env, identity):
    if find(alias) or re.search(r"[\s*?!]", alias):
        sys.exit("bad or duplicate alias")
    lines = read_lines()
    while lines and lines[-1].strip() == "":
        lines.pop()
    lines.append("")
    if group: lines.append(f"# group: {group}")
    if env: lines.append(f"# env: {env}")
    lines.append(f"Host {alias}")
    lines.append(f"    HostName {hostname or alias}")
    if user: lines.append(f"    User {user}")
    if port and port != "22": lines.append(f"    Port {port}")
    if jump: lines.append(f"    ProxyJump {jump}")
    if identity:
        lines.append(f"    IdentityFile {identity}")
        lines.append("    IdentitiesOnly yes")
    write_lines(lines)


def cmd_run(cmd, hosts):
    """Run cmd on all hosts in parallel; hosts with identical output are merged into one block."""
    def one(h):
        try:
            p = subprocess.run(["ssh", "-T", "-o", "BatchMode=yes", "-o", "ConnectTimeout=6", h, cmd],
                               stdin=subprocess.DEVNULL, capture_output=True, text=True, timeout=60)
            return h, p.returncode, (p.stdout + p.stderr).rstrip("\n")
        except subprocess.TimeoutExpired:
            return h, -1, "timeout"
        except Exception as e:
            return h, -1, str(e)
    with ThreadPoolExecutor(16) as ex:
        res = list(ex.map(one, hosts))
    groups, order = {}, []
    for h, rc, out in res:
        k = (rc, out)
        if k not in groups:
            groups[k] = []
            order.append(k)
        groups[k].append(h)
    print(f"{D}$ {cmd}{X}\n")
    for rc, out in order:
        hs = groups[(rc, out)]
        note = "" if rc == 0 else f"  {D}{'ssh error' if rc in (-1, 255) else f'exit {rc}'}{X}"
        print(f'{(G if rc == 0 else R)}●{X} {rgb(235, 225, 230, bold=True)}{", ".join(hs)}{X}{note}')
        for l in out.splitlines():
            print("  " + l)
        print()


def url_default(alias):
    i, h = resolve(alias), find(alias)
    port = ""
    if h and h["group"].lower() == "proxmox":
        port = ":8007" if alias == "pbs" else ":8006"
    return f'https://{i["hostname"]}{port}'


def main():
    a = sys.argv[1:]
    c = a[0] if a else "list"
    if c == "list": cmd_list(False)
    elif c == "live": cmd_list(True)
    elif c == "preview": cmd_preview(a[1])
    elif c == "addr":
        i = resolve(a[1]); print(addr(i), "jump" if i["proxyjump"] else "")
    elif c == "env":
        h = find(a[1]); print(h["env"] if h else "")
    elif c == "line":
        h = find(a[1]); print(h["line"] if h else 1)
    elif c == "groups":
        print("\n".join(sorted({h["group"] for h in parse() if h["group"]})))
    elif c == "url":
        h = find(a[1]); print(h["url"] if h else "")
    elif c == "url-default": print(url_default(a[1]))
    elif c == "run": cmd_run(a[1], a[2:])
    elif c == "has": sys.exit(0 if find(a[1]) else 1)
    elif c == "set": set_meta(a[1], a[2], a[3] if len(a) > 3 else "")
    elif c == "toggle-prod":
        h = find(a[1])
        if h: set_meta(a[1], "env", "" if h["env"].lower() == "prod" else "prod")
    elif c == "add": cmd_add(*a[1:9])


main()