#!/usr/bin/env python3
"""
teleport_prune_by_hostname_v2.py

- Groups nodes by spec.hostname
- In each duplicate group, removes the node with the *oldest version*
- If versions are tied, removes the one with the earliest metadata.expires (oldest heartbeat)
- Dry-run by default; pass --execute to delete with `tctl rm node/<metadata.name> --confirm`
"""

import argparse, json, shutil, subprocess, sys
from collections import defaultdict
from datetime import datetime
from typing import Any, Dict, List, Tuple, Optional

# ---------- tiny shell helpers ----------
def ensure_binary(name: str):
    if shutil.which(name) is None:
        sys.exit(f"ERROR: '{name}' not found on PATH.")

def run(cmd: List[str]) -> subprocess.CompletedProcess:
    try:
        return subprocess.run(cmd, check=True, capture_output=True, text=True)
    except subprocess.CalledProcessError as e:
        sys.stderr.write(f"\nCommand failed: {' '.join(cmd)}\nSTDERR:\n{e.stderr}\n")
        raise

# ---------- JSON field helpers ----------
def get_hostname(n: Dict[str, Any]) -> Optional[str]:
    return (n.get("spec") or {}).get("hostname") or None

def get_uuid(n: Dict[str, Any]) -> Optional[str]:
    return (n.get("metadata") or {}).get("name") or None

def get_version(n: Dict[str, Any]) -> str:
    return (n.get("spec") or {}).get("version") or "0.0.0"

def get_expires(n: Dict[str, Any]) -> Optional[datetime]:
    exp = (n.get("metadata") or {}).get("expires")
    if not exp:
        return None
    # tolerate fractional seconds + Z
    try:
        # Example: 2025-10-26T17:45:26.132967655Z
        exp = exp.rstrip("Z")
        # trim to microseconds if too long
        if "." in exp:
            head, tail = exp.split(".", 1)
            tail = ''.join(ch for ch in tail if ch.isdigit())
            tail = (tail + "000000")[:6]
            exp = f"{head}.{tail}"
        return datetime.fromisoformat(exp)
    except Exception:
        return None

def parse_version_tuple(v: str) -> Tuple[int, ...]:
    parts = []
    for token in (v or "").split("."):
        num = ""
        for ch in token:
            if ch.isdigit(): num += ch
            else: break
        parts.append(int(num) if num else 0)
    return tuple(parts) if parts else (0,)

# ---------- core ----------
def list_nodes(tsh_bin: str, tsh_flags: List[str]) -> List[Dict[str, Any]]:
    cp = run([tsh_bin, "ls", "--format=json", *tsh_flags])
    try:
        data = json.loads(cp.stdout)
    except json.JSONDecodeError:
        sys.exit("ERROR: Could not parse JSON from `tsh ls --format=json`.")
    if not isinstance(data, list):
        sys.exit("ERROR: Unexpected JSON (expected list).")
    return data

def decide_deletions(nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    groups: Dict[str, List[Dict[str, Any]]] = defaultdict(list)
    for n in nodes:
        host = get_hostname(n)
        print(host)
        if not host:  # ignore ungroupable entries
            continue
        entry = {
            "uuid": get_uuid(n),
            "hostname": host,
            "version": get_version(n),
            "expires": get_expires(n),
        }
        entry["version_key"] = parse_version_tuple(entry["version"])
        groups[host].append(entry)

    to_delete: List[Dict[str, Any]] = []
    for host, items in groups.items():
        if len(items) <= 1:
            continue
        # Sort by: newest version first, then latest expiry first
        items_sorted = sorted(
            items,
            key=lambda x: (x["version_key"], x["expires"] or datetime.min),
            reverse=True
        )
        # choose the *oldest* by version; if tie, earliest expiry
        victim = items_sorted[-1]
        to_delete.append(victim)
    return to_delete

def main():
    ap = argparse.ArgumentParser(description="Remove the oldest Teleport node per duplicated hostname (Teleport v2 JSON).")
    ap.add_argument("--tsh", default="tsh")
    ap.add_argument("--tctl", default="tctl")
    ap.add_argument("--tsh-flag", action="append", default=[], help="Extra flags for tsh (repeatable)")
    ap.add_argument("--tctl-flag", action="append", default=[], help="Extra flags for tctl (repeatable)")
    ap.add_argument("--execute", action="store_true", help="Actually delete (default is dry-run)")
    args = ap.parse_args()

    ensure_binary(args.tsh); ensure_binary(args.tctl)

    print("Listing nodes via tsh...")
    nodes = list_nodes(args.tsh, args.tsh_flag)

    deletions = decide_deletions(nodes)
    if not deletions:
        print("No duplicate hostnames found. Nothing to do.")
        return

    print("\nPlanned deletions (oldest per duplicate hostname):")
    for d in deletions:
        exp = d["expires"].isoformat() if d["expires"] else "unknown"
        print(f"  host={d['hostname']:<30} uuid={d['uuid']}  version={d['version']}  expires={exp}")

    if not args.execute:
        print("\nDRY-RUN: No changes made. Re-run with --execute to delete the above nodes.")
        return

    print("\nDeleting via tctl…")
    for d in deletions:
        if not d["uuid"]:
            print(f"SKIP (no UUID): host={d['hostname']} version={d['version']}")
            continue
        cmd = [args.tctl, *args.tctl_flag, "rm", f"node/{d['uuid']}"]
        print("  ", " ".join(cmd))
        try:
            run(cmd); print("    -> OK")
        except Exception: print("    -> FAILED")

    print("\nDone.")

if __name__ == "__main__":
    main()