#!/usr/bin/env python3
"""
/root/scripts/agent-health-check.py — Consolidated Agent Health Verification v2

Verifies: LiteLLM keys (agent-specific), GPU port conflicts, agent Zulip streaming,
gateway liveness, gateway log health, CT liveness, config YAML integrity,
wrapper/CLI integrity, vault secret non-emptiness. NEVER restarts anything.

Usage:
  python3 /root/scripts/agent-health-check.py          # Full check
  python3 /root/scripts/agent-health-check.py --json   # Machine-readable
  python3 /root/scripts/agent-health-check.py --quiet  # Only output on failure

Cron: */10 * * * * python3 /root/scripts/agent-health-check.py --quiet

Changelog:
  v2 (2026-07-26): Added CT liveness, config validation, wrapper integrity,
    vault secret emptiness check. Fixed Koby/Koonimo SSH hosts and agent key
    name format ({NAME}_LITELLM_API_KEY not LITELLM_API_KEY_{NAME}).
    Fleet roster: tanko (.122), mumuni (.123), koby (.129), koonimo (.114),
    abiba (.24).
"""

import subprocess, json, sys, os, time
from datetime import datetime

LITELLM = "http://192.168.68.116:80"
INFISICAL_PROJECT = "322fceab-39da-4854-a55a-568e76c0f13f"
INFISICAL_ENV = "prod"

# PVE node IPs for CT liveness checks
PVE_NODES = {
    "hwepve":  "192.168.68.4",
    "amdpve":  "192.168.68.15",
    "minipve": "192.168.68.12",
    "storepve": "192.168.68.6",
    "acerpve": "192.168.68.9",
    "ocupve":  "192.168.68.5",
}

# Agent definitions: ct, host, user, pve_node, vault_key_name
AGENTS = {
    "tanko":   {"ct": 112, "host": "192.168.68.122", "user": "jerome", "pve": "amdpve",  "vault_key": "TANKO_LITELLM_API_KEY"},
    "abiba":   {"ct": 100, "host": "192.168.68.24",  "user": "root",   "pve": "hwepve",  "vault_key": None},  # Pi agent + Mumuni Zulip, no vault key
    "koby":    {"ct": 111, "host": "192.168.68.129", "user": "root",   "pve": "amdpve",  "vault_key": "KOBY_LITELLM_API_KEY"},
    "koonimo": {"ct": 113, "host": "192.168.68.114", "user": "root",   "pve": "amdpve",  "vault_key": "KOONIMO_LITELLM_API_KEY"},
}

GPU_HOSTS = {
    "gpu-rtx3090 (.8)":     {"host": "192.168.68.8",   "port": 8080, "service": "llama-server"},
    "gpu-rtx5070 (.110)":   {"host": "192.168.68.110", "port": 8080, "service": "llama-server"},
    "gpu-strixhalo (.15)":  {"host": "192.168.68.15",  "port": 8080, "service": "ornith-server"},
}

FAIL = []

INFISICAL_TOKEN = os.environ.get("INFISICAL_TOKEN")
INFISICAL_API_URL = os.environ.get("INFISICAL_API_URL", "https://vault.sysloggh.net")

# ── Helpers ──────────────────────────────────────────────────────────

def ssh(host, cmd, user="root"):
    """Execute a command on a remote host, return stdout or None."""
    try:
        result = subprocess.run(
            ["ssh", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=8",
             f"{user}@{host}", cmd],
            capture_output=True, text=True, timeout=15
        )
        return result.stdout.strip() if result.returncode == 0 else None
    except:
        return None

def http_get(url, headers=None, timeout=5):
    """Return HTTP status code as string."""
    try:
        cmd = ["curl", "-sfk", "--connect-timeout", str(timeout), "-o", "/dev/null", "-w", "%{http_code}"]
        if headers:
            for k, v in headers.items():
                cmd.extend(["-H", f"{k}: {v}"])
        cmd.append(url)
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout+3)
        return result.stdout.strip() or "000"
    except:
        return "timeout"

def http_json(url, headers=None, timeout=5):
    """Return parsed JSON from URL, or None."""
    try:
        cmd = ["curl", "-sfk", "--connect-timeout", str(timeout)]
        if headers:
            for k, v in headers.items():
                cmd.extend(["-H", f"{k}: {v}"])
        cmd.append(url)
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout+3)
        return json.loads(result.stdout) if result.returncode == 0 and result.stdout else None
    except:
        return None

def run_infisical(args, quiet=True):
    """Run infisical CLI with env-based auth, return stdout or None."""
    env = os.environ.copy()
    env["INFISICAL_API_URL"] = INFISICAL_API_URL
    if INFISICAL_TOKEN:
        env["INFISICAL_TOKEN"] = INFISICAL_TOKEN
    try:
        result = subprocess.run(
            ["/usr/bin/infisical"] + args,
            capture_output=True, text=True, timeout=15, env=env
        )
        return result.stdout.strip() if result.returncode == 0 else None
    except:
        return None

# ── KEY LOOKUP FIX ───────────────────────────────────────────────────

def _get_agent_key(agent_name, vault_key_name):
    """Retrieve agent-specific key from Infisical vault.
    
    Uses {NAME}_LITELLM_API_KEY format (e.g., TANKO_LITELLM_API_KEY,
    KOONIMO_LITELLM_API_KEY) which matches actual vault key names.
    """
    if not vault_key_name:
        return None

    # Primary: get the agent-specific key by name
    key = run_infisical([
        "secrets", "get", vault_key_name,
        "--projectId=" + INFISICAL_PROJECT,
        "--env=" + INFISICAL_ENV,
        "--plain",
    ])
    if key and key.startswith("sk-"):
        return key

    # Fallback: export all and search for the key name
    try:
        export = run_infisical([
            "export",
            "--projectId=" + INFISICAL_PROJECT,
            "--env=" + INFISICAL_ENV,
            "--format=dotenv",
        ])
        if export:
            for line in export.splitlines():
                if line.startswith(vault_key_name + "="):
                    value = line.split("=", 1)[1].strip().strip('"').strip("'")
                    if value.startswith("sk-"):
                        return value
    except:
        pass

    return None

# Inject keys from vault for each agent
for agent_name in AGENTS:
    info = AGENTS[agent_name]
    key = _get_agent_key(agent_name, info.get("vault_key"))
    AGENTS[agent_name]["key"] = key


# ═══════════════════════════════════════════════════════════════════
# CHECK 1: LiteLLM Key Validation (agent-specific keys)
# ═══════════════════════════════════════════════════════════════════

def check_keys():
    for name, agent in AGENTS.items():
        key = agent.get("key")
        if not key:
            print(f"  ❌ {name}: NO KEY FOUND (vault empty or unreachable)")
            FAIL.append(f"key:{name}:no-key")
            continue
        data = http_json(f"{LITELLM}/v1/models",
                         headers={"Authorization": f"Bearer {key}"})
        if data and data.get("data"):
            model = data["data"][0].get("id", "?")
            print(f"  ✅ {name}: key valid → {model}")
        else:
            print(f"  ❌ {name}: KEY FAILURE — auth rejected or unreachable")
            FAIL.append(f"key:{name}")


# ═══════════════════════════════════════════════════════════════════
# CHECK 2: GPU Port Conflict Detection (unchanged)
# ═══════════════════════════════════════════════════════════════════

def check_gpu_ports():
    for label, gpu in GPU_HOSTS.items():
        host = gpu["host"]
        port = gpu["port"]
        svc = gpu["service"]

        svc_status = ssh(host, f"systemctl is-active {svc}")
        port_owner = ssh(host, f"ss -tlnp 2>/dev/null | grep -Po ':{port}\\s+.*pid=\\K[0-9]+' | head -1")

        if not svc_status:
            print(f"  ❌ {label}: UNREACHABLE")
            FAIL.append(f"gpu-unreachable:{host}")
            continue

        if not port_owner:
            print(f"  ❌ {label}: PORT {port} NOT LISTENING (svc={svc_status})")
            FAIL.append(f"gpu-no-port:{label}")
        elif svc_status != "active":
            svc_pid = ssh(host, f"systemctl show {svc} -p MainPID 2>/dev/null | cut -d= -f2")
            if svc_pid and port_owner != svc_pid:
                print(f"  ❌ {label}: GHOST PROCESS — port owned by pid {port_owner}, svc pid {svc_pid} (svc={svc_status})")
                FAIL.append(f"gpu-ghost:{label}:{port_owner}")
            else:
                print(f"  ⚠️  {label}: svc={svc_status}, port owned by {port_owner}")
        else:
            health = ssh(host, f"curl -s --max-time 5 http://localhost:{port}/health")
            if health and '"status":"ok"' in health:
                print(f"  ✅ {label}: healthy (pid={port_owner})")
            elif health and '"status":"no slot available"' in health:
                print(f"  ✅ {label}: healthy (loading, pid={port_owner})")
            elif health and '"error"' in health.lower():
                print(f"  ⚠️  {label}: error response (pid={port_owner}): {health[:80]}")
            else:
                print(f"  ⚠️  {label}: unknown health (pid={port_owner}): {str(health)[:80]}")


# ═══════════════════════════════════════════════════════════════════
# CHECK 3: Agent Gateway Liveness + Streaming (now covers all agents)
# ═══════════════════════════════════════════════════════════════════

def check_agents():
    for name, agent in AGENTS.items():
        host = agent.get("host")
        user = agent.get("user")
        ct = agent["ct"]

        if not host or not user:
            print(f"  ⬜ {name} (CT {ct}): cannot SSH — skip liveness check")
            continue

        # Gateway process
        pid = ssh(host, "pgrep -f 'hermes_cli.main gateway run' | grep -v infisical | head -1", user=user)
        if not pid:
            # Try alternate binary name
            pid = ssh(host, "pgrep -f 'hermes.*gateway' | grep -v infisical | grep -v bash | head -1", user=user)
        if not pid:
            print(f"  ❌ {name}: GATEWAY NOT RUNNING")
            FAIL.append(f"gateway-down:{name}")
            continue

        # Gateway state file
        state = ssh(host, "cat ~/.hermes/gateway_state.json 2>/dev/null", user=user)
        if state:
            try:
                st = json.loads(state)
                gw_state = st.get("gateway_state", "?")
                zulip = st.get("platforms", {}).get("zulip", {}).get("state", "?")
            except:
                gw_state, zulip = "corrupt", "?"
        else:
            gw_state, zulip = "no-state-file", "?"

        # Zulip streaming check
        adapter_paths = [
            "~/.hermes/plugins/zulip-platform/adapter.py",
            "~/.hermes/plugins/platforms/zulip/adapter.py",
        ]
        streaming = "no"
        for p in adapter_paths:
            has_edit = ssh(host, f"grep -c 'async def edit_message' {p} 2>/dev/null", user=user)
            if has_edit and has_edit != "0":
                streaming = "yes"
                break

        # Recent errors
        recent_errors = ssh(host,
            r"journalctl --user -u hermes-gateway --since '10 min ago' -o cat --no-pager 2>/dev/null "
            r"| grep -ci 'error\|traceback\|exception\|401\|403\|500' || echo 0",
            user=user)
        recent_errors = (recent_errors or "0").strip().split("\n")[-1]

        print(f"  {'✅' if gw_state == 'running' and zulip == 'connected' else '⚠️'} "
              f"{name}: gw={gw_state} zulip={zulip} streaming={streaming} "
              f"errors_10m={recent_errors.strip() or '0'} pid={pid}")


# ═══════════════════════════════════════════════════════════════════
# CHECK 4: CT Liveness (NEW)
# ═══════════════════════════════════════════════════════════════════

def check_ct_liveness():
    """Check that all agent CTs are running on their PVE nodes."""
    for name, agent in AGENTS.items():
        ct = agent["ct"]
        pve_node = agent.get("pve")
        if not pve_node:
            print(f"  ⬜ {name} (CT {ct}): no PVE node mapped — skip")
            continue

        pve_ip = PVE_NODES.get(pve_node)
        if not pve_ip:
            print(f"  ⬜ {name}: unknown PVE node '{pve_node}' — skip")
            continue

        status = ssh(pve_ip, f"pct status {ct} 2>/dev/null", user="root")
        if not status:
            print(f"  ❌ {name} (CT {ct} on {pve_node}): PVE UNREACHABLE")
            FAIL.append(f"ct-unreachable:{name}:{pve_ip}")
        elif "running" in status:
            print(f"  ✅ {name} (CT {ct} on {pve_node}): running")
        elif "stopped" in status:
            print(f"  ❌ {name} (CT {ct} on {pve_node}): STOPPED")
            FAIL.append(f"ct-stopped:{name}")
        else:
            print(f"  ⚠️  {name} (CT {ct} on {pve_node}): {status.strip()}")


# ═══════════════════════════════════════════════════════════════════
# CHECK 5: Config YAML Integrity (NEW)
# ═══════════════════════════════════════════════════════════════════

def check_config_integrity():
    """Verify agent config.yaml parses as valid YAML."""
    for name, agent in AGENTS.items():
        host = agent.get("host")
        user = agent.get("user")
        if not host or not user:
            print(f"  ⬜ {name}: cannot SSH — skip config check")
            continue

        # Check YAML parses
        yaml_ok = ssh(host,
            "python3 -c "
            '"import yaml; yaml.safe_load(open(\'/root/.hermes/config.yaml\')); print(\'OK\')" '
            "2>&1 || echo 'FAIL'",
            user=user)
        if not yaml_ok:
            print(f"  ❌ {name}: SSH UNREACHABLE (config check skipped)")
            FAIL.append(f"config-unreachable:{name}")
        elif "OK" in yaml_ok:
            print(f"  ✅ {name}: config.yaml valid YAML")
        else:
            print(f"  ❌ {name}: config.yaml YAML ERROR — {yaml_ok[:120]}")
            FAIL.append(f"config-yaml-error:{name}")


# ═══════════════════════════════════════════════════════════════════
# CHECK 6: Wrapper/CLI Integrity (NEW)
# ═══════════════════════════════════════════════════════════════════

def check_wrapper_integrity():
    """Verify the hermes CLI wrapper exists and can reach hermes-real."""
    for name, agent in AGENTS.items():
        host = agent.get("host")
        user = agent.get("user")
        if not host or not user:
            print(f"  ⬜ {name}: cannot SSH — skip wrapper check")
            continue

        # Check wrapper exists
        wrapper = ssh(host, "ls -la /root/.local/bin/hermes 2>/dev/null", user=user)
        if not wrapper:
            # Check alternate wrapper locations
            wrapper = ssh(host, "which hermes 2>/dev/null; command -v hermes 2>/dev/null", user=user)
            if not wrapper:
                print(f"  ❌ {name}: NO HERMES CLI WRAPPER FOUND")
                FAIL.append(f"wrapper-missing:{name}")
                continue
            else:
                print(f"  ⚠️  {name}: hermes at {wrapper.strip()} (not ~/.local/bin/hermes)")

        # Check wrapper has correct infisical path
        infisical_path_valid = ssh(host,
            "head -20 /root/.local/bin/hermes 2>/dev/null | grep -q '/usr/bin/infisical' && echo OK || echo MISS",
            user=user)
        if infisical_path_valid == "MISS":
            # Check if infisical exists on path
            inf_actual = ssh(host, "command -v infisical 2>/dev/null", user=user)
            if not inf_actual:
                print(f"  ❌ {name}: INFISICAL NOT INSTALLED (wrapper broken)")
                FAIL.append(f"wrapper-no-infisical:{name}")
            else:
                print(f"  ⚠️  {name}: wrapper infisical path may be wrong (infisical at {inf_actual})")
                FAIL.append(f"wrapper-infisical-path:{name}")

        # Check hermes-real exists
        hermes_real = ssh(host,
            "ls -la /root/.local/bin/hermes-real 2>/dev/null || echo MISS",
            user=user)
        if not hermes_real or hermes_real.strip() == "MISS":
            # Check venv path
            hermes_real = ssh(host,
                "ls -la /usr/local/lib/hermes-agent/venv/bin/hermes 2>/dev/null || echo MISS",
                user=user)
            if not hermes_real or hermes_real.strip() == "MISS":
                print(f"  ❌ {name}: hermes-real NOT FOUND (wrapper broken)")
                FAIL.append(f"wrapper-no-hermes-real:{name}")
            else:
                print(f"  ✅ {name}: hermes-real at alt path")

        # Check the .env file has the key
        env_has_key = ssh(host,
            "grep -c 'LITELLM_API_KEY' /root/.hermes/.env 2>/dev/null || echo 0",
            user=user)
        if env_has_key and env_has_key.strip() not in ("", "0"):
            print(f"  ✅ {name}: wrapper + .env key present")
        else:
            print(f"  ⚠️  {name}: .env may be missing LITELLM_API_KEY entry")


# ═══════════════════════════════════════════════════════════════════
# CHECK 7: Vault Secret Non-Emptiness (NEW)
# ═══════════════════════════════════════════════════════════════════

def check_vault_secrets():
    """Verify agent-specific vault secrets are non-empty and start with sk-."""
    for name, agent in AGENTS.items():
        vault_key_name = agent.get("vault_key")
        if not vault_key_name:
            continue

        key = agent.get("key")
        if not key:
            print(f"  ❌ {name}: vault secret {vault_key_name} MISSING or EMPTY")
            FAIL.append(f"vault-empty:{name}:{vault_key_name}")
        elif not key.startswith("sk-"):
            print(f"  ❌ {name}: vault secret {vault_key_name} WRONG FORMAT (starts '{key[:8]}...')")
            FAIL.append(f"vault-bad-format:{name}:{vault_key_name}")
        else:
            print(f"  ✅ {name}: vault {vault_key_name}=sk-...{key[-4:]}")


# ═══════════════════════════════════════════════════════════════════
# DEPLOY: copy updated script to /root/scripts/ on local host
# ═══════════════════════════════════════════════════════════════════

def deploy_self():
    """Copy this script to /root/scripts/agent-health-check.py if out of date."""
    dest = "/root/scripts/agent-health-check.py"
    try:
        with open(__file__, "r") as f:
            current = f.read()
        if os.path.isfile(dest):
            with open(dest, "r") as f:
                existing = f.read()
            if current == existing:
                return  # Already deployed
        # Write new version
        with open(dest, "w") as f:
            f.write(current)
        os.chmod(dest, 0o755)
        print(f"  📦 Deployed updated script to {dest}")
    except:
        pass  # Not fatal if deploy fails


# ═══════════════════════════════════════════════════════════════════
# MAIN
# ═══════════════════════════════════════════════════════════════════

def main():
    quiet = "--quiet" in sys.argv
    as_json = "--json" in sys.argv

    # Self-deploy to canonical location
    if not quiet and "--no-deploy" not in sys.argv:
        deploy_self()

    if not quiet:
        print(f"🏥 Agent Health Check v2 — {datetime.now().strftime('%Y-%m-%d %H:%M UTC')}")
        print()

    print("🔑 LiteLLM Keys:")
    check_keys()
    print()

    print("🎮 GPU Port Health:")
    check_gpu_ports()
    print()

    print("🤖 Agent Gateways:")
    check_agents()
    print()

    print("🖥️  CT Liveness:")
    check_ct_liveness()
    print()

    print("📝 Config Integrity:")
    check_config_integrity()
    print()

    print("🔌 Wrapper/CLI Integrity:")
    check_wrapper_integrity()
    print()

    print("🔐 Vault Secrets:")
    check_vault_secrets()

    if FAIL:
        print(f"\n❌ {len(FAIL)} FAILURE(S): {' | '.join(FAIL)}")
        if quiet:
            print(f"ALERT agent-health:{','.join(FAIL)}")
    elif not quiet:
        print("\n✅ All checks passed")

    if as_json:
        print(json.dumps({"timestamp": datetime.now().isoformat(),
                          "failures": FAIL, "healthy": len(FAIL) == 0}))

    sys.exit(1 if FAIL else 0)

if __name__ == "__main__":
    main()