#!/usr/bin/env python3
"""Generate ~/.ssh/config.d/homelab from inventory.yaml.

Usage:
  python3 ssh/gen-config.py                        # print to stdout
  python3 ssh/gen-config.py --install              # write to ~/.ssh/config.d/homelab

The generated config provides short hostname aliases for every host in
the homelab inventory. LAN IPs are preferred (they work directly on-LAN
and are routed via Netbird 192.168.8.0/24 off-LAN); mesh FQDNs are
available as <name>-mesh fallbacks for roaming workstations.
"""

from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

try:
    import yaml
except ImportError:
    print("PyYAML required (pip install pyyaml)", file=sys.stderr)
    sys.exit(2)

INVENTORY_PATH = Path(os.environ.get(
    "HOMELAB_CONTEXT_DIR", "/opt/homelab-context"
)) / "inventory.yaml"

SSH_CONFIG_DIR = Path.home() / ".ssh" / "config.d"
SSH_CONFIG_FILE = SSH_CONFIG_DIR / "homelab"
SSH_MAIN_CONFIG = Path.home() / ".ssh" / "config"
INCLUDE_LINE = f"Include ~/.ssh/config.d/homelab"

# Hosts to skip in the generated config
SKIP_HOSTS = {"ludo-mini", "authentik"}  # offline / not enrolled


def inventory() -> dict:
    if not INVENTORY_PATH.exists():
        print(f"ERROR: no inventory at {INVENTORY_PATH}", file=sys.stderr)
        sys.exit(1)
    return yaml.safe_load(INVENTORY_PATH.read_text())


def gen_config() -> str:
    inv = inventory()
    hosts = inv.get("hosts", {})
    mesh_globals = inv.get("mesh", {})

    lines: list[str] = []
    lines.append("# Homelab SSH config — generated from inventory.yaml")
    lines.append(f"# Source: {INVENTORY_PATH}")
    lines.append("# Do not edit by hand. Run: homelab ssh-config")
    lines.append("")

    # --- Common defaults (before any Host block) ---
    lines.append("# --- Defaults ---")
    lines.append("Host *")
    lines.append("  IdentityFile ~/.ssh/id_ed25519")
    lines.append("  IdentitiesOnly yes")
    lines.append("  ServerAliveInterval 30")
    lines.append("  StrictHostKeyChecking accept-new")
    lines.append("")

    def add_entry(
        tag: str,
        hostname: str,
        user: str = "root",
        port: int = 22,
        extra_lines: list[str] | None = None,
    ) -> None:
        lines.append(f"Host {tag}")
        lines.append(f"  HostName {hostname}")
        lines.append(f"  User {user}")
        if port != 22:
            lines.append(f"  Port {port}")
        if extra_lines:
            lines.extend(f"  {el}" for el in extra_lines)
        lines.append("")

    # Sort hosts by kind for a logical output order
    def sort_key(item):
        name, h = item
        kind = h.get("kind", "")
        order = {
            "workstation": 0,
            "proxmox-host": 1,
            "lxc": 2,
            "vm": 3,
            "external": 4,
        }
        return (order.get(kind, 9), name)

    sorted_hosts = sorted(
        [(n, h) for n, h in hosts.items() if n not in SKIP_HOSTS],
        key=sort_key,
    )

    for name, h in sorted_hosts:
        kind = h.get("kind", "")
        ssh_config = h.get("ssh", {}) or {}
        user = ssh_config.get("user", "root")
        port = ssh_config.get("port", 22)
        lan_ip = h.get("lan_ip")
        mesh = h.get("mesh", {}) or {}
        nb = mesh.get("netbird", {}) or {}
        ts = mesh.get("tailscale", {}) or {}
        nb_fqdn = nb.get("fqdn") or nb.get("ip") or ""
        nb_port = ssh_config.get("netbird_port")

        # Determine primary address
        if nb_port and not lan_ip:
            # Hosts with netbird SSH but no LAN IP: mesh-only
            primary = nb_fqdn
        elif lan_ip:
            primary = lan_ip
        elif nb_fqdn:
            primary = nb_fqdn
        elif ts.get("fqdn"):
            primary = ts["fqdn"]
        else:
            continue  # no address found

        extra = []
        if kind == "proxmox-host":
            extra = [
                "ControlMaster auto",
                "ControlPath ~/.ssh/cm/%C",
                "ControlPersist 2h",
            ]

        add_entry(name, primary, user, port, extra_lines=extra)

        # Mesh fallback for workstations and hubris
        if nb_fqdn and (kind == "workstation" or nb_port):
            mesh_tag = f"{name}-mesh"
            mesh_port = nb_port or port
            mesh_extra = []
            if kind == "proxmox-host":
                mesh_extra = [
                    "ControlMaster auto",
                    "ControlPath ~/.ssh/cm/%C",
                    "ControlPersist 2h",
                ]
            add_entry(mesh_tag, nb_fqdn, user, mesh_port, extra_lines=mesh_extra)

    # --- *.hubris.network ControlMaster ---
    lines.append("# --- Mesh ControlMaster (speeds up repeated mesh ops) ---")
    lines.append("Host *.netbird.selfhosted")
    lines.append("  ControlMaster auto")
    lines.append("  ControlPath ~/.ssh/cm/%C")
    lines.append("  ControlPersist 2h")
    lines.append("")

    return "\n".join(lines)


def ensure_include() -> bool:
    """Add 'Include ~/.ssh/config.d/homelab' to the main SSH config if missing."""
    if not SSH_MAIN_CONFIG.exists():
        SSH_MAIN_CONFIG.parent.mkdir(parents=True, exist_ok=True)
        with open(SSH_MAIN_CONFIG, "w") as f:
            f.write(f"{INCLUDE_LINE}\n")
        return True

    content = SSH_MAIN_CONFIG.read_text()
    for line in content.splitlines():
        stripped = line.strip()
        if stripped.startswith("Include") and "homelab" in stripped:
            return False  # already present

    # Prepend to existing config
    updated = f"{INCLUDE_LINE}\n\n{content}"
    SSH_MAIN_CONFIG.write_text(updated)
    return True


def install() -> None:
    """Write config to ~/.ssh/config.d/homelab and ensure Include is set."""
    SSH_CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    config = gen_config()
    SSH_CONFIG_FILE.write_text(config)
    SSH_CONFIG_FILE.chmod(0o644)
    changed = ensure_include()
    print(f"Wrote {SSH_CONFIG_FILE} ({len(config.splitlines())} lines)")
    if changed:
        print(f"Added '{INCLUDE_LINE}' to {SSH_MAIN_CONFIG}")
    print("Done. Run: ssh <hostname>  (e.g. ssh gitea)")


def main() -> int:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--install", "-i", action="store_true",
                   help=f"write to {SSH_CONFIG_FILE} and wire Include into main config")
    args = p.parse_args()

    if args.install:
        install()
    else:
        print(gen_config())

    return 0


if __name__ == "__main__":
    sys.exit(main())