from dcim.models import Device
from ipam.models import IPAddress
from virtualization.models import VirtualMachine
from startup_script_utils import load_yaml
import sys

def link_primary_ip(assets, asset_model):
  for params in assets:
    primary_ip_fields = set(params) & {'primary_ip4', 'primary_ip6'}
    if not primary_ip_fields:
      continue

    for assoc, details in optional_assocs.items():
      if assoc in params:
        model, field = details
        query = { field: params.pop(assoc) }

        try:
          params[assoc] = model.objects.get(**query)
        except model.DoesNotExist:
          primary_ip_fields -= {assoc}
          print(f"⚠️ IP Address '{query[field]}' not found")

    asset = asset_model.objects.get(name=params['name'])
    for field in primary_ip_fields:
      if getattr(asset, field) != params[field]:
        setattr(asset, field, params[field])
        print(f"🔗 Define primary IP '{params[field].address}' on '{asset.name}'")
    asset.save()

devices = load_yaml('/opt/netbox/initializers/devices.yml')
virtual_machines = load_yaml('/opt/netbox/initializers/virtual_machines.yml')

optional_assocs = {
    'primary_ip4': (IPAddress, 'address'),
    'primary_ip6': (IPAddress, 'address')
}

if devices is None and virtual_machines is None:
    sys.exit()
elif devices is not None:
  link_primary_ip(devices, Device)
elif virtual_machines is not None:
  link_primary_ip(virtual_machines, VirtualMachine)
else:
  link_primary_ip(devices, Device)
  link_primary_ip(virtual_machines, VirtualMachine)