#!/usr/bin/env python3
import sys
import yaml
from pathlib import Path


class VaultSafeLoader(yaml.SafeLoader):
    pass


def vault_constructor(loader, node):
    return loader.construct_scalar(node)


VaultSafeLoader.add_constructor('!vault', vault_constructor)


def load_sot(path):
    with open(path, 'r') as f:
        return yaml.load(f, Loader=VaultSafeLoader)


def ip_for(info, prefer_desired=False):
    # Default behavior for flat-network operation:
    # prefer current_ip and only use desired_ip as fallback.
    # For migration windows, set prefer_desired=True.
    current_ip = info.get('current_ip')
    desired_ip = info.get('desired_ip')

    if prefer_desired:
        if isinstance(desired_ip, str) and desired_ip:
            return desired_ip
        if isinstance(current_ip, str) and current_ip:
            return current_ip
        return None

    if isinstance(current_ip, str) and current_ip:
        return current_ip
    if isinstance(desired_ip, str) and desired_ip:
        return desired_ip
    return None


def is_retired(info):
    status = str(info.get('lifecycle_status', '')).strip().lower()
    return status in {'retired', 'retired-shutdown', 'shutdown'}


def is_onboarding_tbd(info):
    status = str(info.get('onboarding_status', '')).strip().lower()
    return status.startswith('tbd')


def is_active(info):
    if is_retired(info):
        return False
    if is_onboarding_tbd(info):
        return False
    if info.get('ansible_managed') is False:
        return False
    return True


def main():
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument('--sot', required=True)
    p.add_argument('--out', required=True)
    p.add_argument(
        '--prefer-desired',
        action='store_true',
        help='Prefer desired_ip over current_ip when generating inventory (use during migrations).',
    )
    args = p.parse_args()

    sot = load_sot(args.sot)
    # Prefer the non-reserved key; keep fallback for older SoT files.
    hosts = sot.get('lab_hosts', sot.get('hosts', {})) or {}
    edge_host = (sot.get('edge_routing') or {}).get('edge_host') or {}
    edge_host_ip = edge_host.get('ip')

    out_lines = []
    out_lines.append('# Generated inventory from ../group_vars/all.yml')
    out_lines.append('')

    # watchtower
    out_lines.append('# --- Watchtower (local controller) ---')
    out_lines.append('[watchtower]')
    out_lines.append('localhost ansible_connection=local')
    out_lines.append('')

    # proxmox
    out_lines.append('# --- Proxmox Cluster (management) ---')
    out_lines.append('[proxmox_cluster]')
    for name, info in hosts.items():
        if info.get('role') == 'proxmox' and is_active(info):
            ip = ip_for(info, prefer_desired=args.prefer_desired)
            if ip:
                out_lines.append(
                    f"{name}     ansible_host={ip}     ansible_user=root     "
                    "ansible_ssh_private_key_file=/home/chester/.ssh/id_ed25519     ansible_port=22"
                )
    out_lines.append('')
    out_lines.append('[proxmox_cluster:vars]')
    out_lines.append('ansible_user=root')
    out_lines.append('ansible_become=true')
    out_lines.append('ansible_python_interpreter=/usr/bin/python3')
    out_lines.append('')

    # swarm managers
    out_lines.append('# --- Swarm Managers ---')
    out_lines.append('[swarm_managers]')
    for name, info in hosts.items():
        if (info.get('role') == 'swarm_manager' or name.startswith('swarm-manager')) and is_active(info):
            ip = ip_for(info, prefer_desired=args.prefer_desired)
            if ip:
                out_lines.append(f"{name} ansible_host={ip}")
    out_lines.append('')

    # swarm workers
    out_lines.append('# --- Swarm Workers ---')
    out_lines.append('[swarm_workers]')
    for name, info in hosts.items():
        if (info.get('role') == 'swarm_worker' or name.startswith('swarm-worker')) and is_active(info):
            ip = ip_for(info, prefer_desired=args.prefer_desired)
            if ip:
                out_lines.append(f"{name} ansible_host={ip}")
    out_lines.append('')

    out_lines.append('[swarm_hosts:children]')
    out_lines.append('swarm_managers')
    out_lines.append('swarm_workers')
    out_lines.append('')
    out_lines.append('[swarm_hosts:vars]')
    out_lines.append('ansible_user=chester')
    out_lines.append('ansible_ssh_private_key_file=/home/chester/.ssh/id_ed25519')
    out_lines.append('')

    # standalone ubuntu VMs
    out_lines.append('# --- Standalone Ubuntu VMs ---')
    out_lines.append('[standalone_ubuntu]')
    for name, info in hosts.items():
        if info.get('role') in {'standalone_vm', 'standalone_ubuntu'} and is_active(info):
            ip = ip_for(info, prefer_desired=args.prefer_desired)
            if ip:
                out_lines.append(f"{name} ansible_host={ip}")
    out_lines.append('')

    out_lines.append('[standalone_ubuntu:vars]')
    out_lines.append('ansible_user=chester')
    out_lines.append('ansible_ssh_private_key_file=/home/chester/.ssh/id_ed25519')
    out_lines.append('')

    # heimdall edge host
    out_lines.append('# --- Heimdall (Edge Router / Traefik host) ---')
    out_lines.append('[heimdall_hosts]')
    heimdall_info = hosts.get('heimdall', {})
    heimdall_ip = ip_for(heimdall_info, prefer_desired=args.prefer_desired) or edge_host_ip
    if heimdall_info and is_active(heimdall_info) and heimdall_ip:
        out_lines.append(f"heimdall ansible_host={heimdall_ip}")
    out_lines.append('')

    out_lines.append('[heimdall_hosts:vars]')
    out_lines.append('ansible_user=chester')
    out_lines.append('ansible_ssh_private_key_file=/home/chester/.ssh/id_ed25519')
    out_lines.append('')

    # ai_grid
    out_lines.append('# --- AI Grid ---')
    out_lines.append('[ai_grid]')
    for name, info in hosts.items():
        if (info.get('role') == 'ai_node' or name.startswith('ai-')) and is_active(info):
            ip = ip_for(info, prefer_desired=args.prefer_desired)
            if ip:
                out_lines.append(f"{name} ansible_host={ip}")
    out_lines.append('')

    # docker hosts
    out_lines.append('# --- Docker Hosts ---')
    out_lines.append('[docker_hosts]')
    for name, info in hosts.items():
        if info.get('role') in {'docker_host', 'standalone_vm', 'standalone_ubuntu'} and is_active(info):
            ip = ip_for(info, prefer_desired=args.prefer_desired)
            if ip:
                out_lines.append(f"{name} ansible_host={ip}")
    out_lines.append('')

    # storage
    out_lines.append('# --- Storage ---')
    out_lines.append('[storage]')
    for name, info in hosts.items():
        if (info.get('role') == 'nas' or name in ('synology', 'terramaster')) and is_active(info):
            ip = ip_for(info, prefer_desired=args.prefer_desired)
            if ip:
                out_lines.append(f"{name} ansible_host={ip} ansible_scp_if_ssh=True")
    out_lines.append('')

    out_lines.append('# --- Lifecycle: Onboarding TBD ---')
    out_lines.append('[onboarding_tbd]')
    for name, info in hosts.items():
        if is_onboarding_tbd(info):
            ip = ip_for(info, prefer_desired=args.prefer_desired)
            if ip:
                out_lines.append(f"{name} ansible_host={ip}")
    out_lines.append('')

    out_lines.append('# --- Lifecycle: Retired / Shutdown ---')
    out_lines.append('[retired_hosts]')
    for name, info in hosts.items():
        if is_retired(info):
            ip = ip_for(info, prefer_desired=args.prefer_desired)
            if ip:
                out_lines.append(f"{name} ansible_host={ip}")
    out_lines.append('')

    out_lines.append('# --- Aggregate grouping ---')
    out_lines.append('[ubuntu_lab:children]')
    out_lines.append('swarm_managers')
    out_lines.append('swarm_workers')
    out_lines.append('standalone_ubuntu')
    out_lines.append('ai_grid')
    out_lines.append('docker_hosts')
    out_lines.append('storage')
    out_lines.append('')
    out_lines.append('[ubuntu_lab:vars]')
    out_lines.append('ansible_user=chester')
    out_lines.append('ansible_ssh_private_key_file=/home/chester/.ssh/id_ed25519')

    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text('\n'.join(out_lines) + '\n')


if __name__ == '__main__':
    main()