import os
import sys
from unifi_controller_api import UnifiController, UnifiDevice
from unifi_controller_api.exceptions import UnifiAuthenticationError, UnifiAPIError
from pprint import pprint

import pynetbox
from pynetbox.core.response import Record

class Db():
    def __init__(self):
        self.devices = []
        self.interfaces = []
        self.ips = []
        self.macs = []

class Ip():
    def __init__(self, vrf):
        self.address = address
        self.vrf = vrf
        self.ips = ips

class Query():
    def __init__(self, query):
        self.query = query

def find_device(nb, name: str):
    return Query(lambda: nb.dcim.devices.get(name=name).id)

def find_interface(nb, device: str, name: str):
    return Query(lambda: nb.dcim.interfaces.get(device=device, name=name).id)

class NetboxCache():
    def __init__(self, nb):
        self.nb = nb
        self.device_roles = {}
        self.device_types = {}
        self.ip_addresses = {}

    def get_device_role(self, slug):
        dt = self.device_roles.get(slug)

        if dt is not None:
            return dt

        dt = self.nb.dcim.device_roles.get(slug=slug)
        if dt is None:
            raise Exception(f"No such device type: {slug}")

        self.device_roles[slug] = dt
        return dt

    def get_device_type(self, slug):
        dt = self.device_types.get(slug)

        if dt is not None:
            return dt

        dt = self.nb.dcim.device_types.get(slug=slug)
        if dt is None:
            raise Exception(f"No such device type: {slug}")

        self.device_types[slug] = dt
        return dt

    def get_or_create_ip_address(self, addr: str, vrf: Record | None, data):
        vrf_id = vrf.id if vrf is not None else None
        key = (addr, vrf_id)
        ip = self.ip_addresses.get(key)
        if ip is not None:
            return ip

        ip = self.nb.ipam.ip_addresses.get(address=addr, vrf_id=vrf_id)
        if ip is not None:
            print(f"Found IP address {ip.id} address={ip.address}, vrf={ip.vrf}")
            ip.update(data)
            ip = self.nb.ipam.ip_addresses.get(address=addr, vrf_id=vrf_id)
            self.ip_addresses[key] = ip
            return ip

        ip = self.nb.ipam.ip_addresses.create(address=addr, vrf=vrf_id, status="active")
        self.ip_addresses[key] = ip
        return ip

def create_or_update_device(nb, d):
    device = nb.dcim.devices.get(name = d["name"])
    if device is None:
        device = nb.dcim.devices.create(d)
        print(f"Created device id={device.id}, name={device.name}")

        return device

    print(f"Updating device id={device.id}, name={device.name}")
    device.update(d)

    return nb.dcim.devices.get(id=device.id)

def create_or_update_interface(nb, i):
    iface = nb.dcim.interfaces.get(device_id=i["device"], name=i["name"])
    if iface is None:
        iface = nb.dcim.interfaces.create(i)
        print(f"Created interface id={iface.id}, name={iface.name}")

        return iface

    print(f"Updating interface id={iface.id}, name={iface.name}")
    iface.update(i)
    return nb.dcim.interfaces.get(id=iface.id)

def create_or_update_mac_address(nb, data):
    ma = nb.dcim.mac_addresses.get(mac_address=data["mac_address"])
    if ma is None:
        ma = nb.dcim.mac_addresses.create(data)
        print(f"Created MAC address id={ma.id}, address={ma.mac_address}")

        return ma

    print(f"Updating MAC address id={ma.id}, address={ma.mac_address}")
    ma.update(data)
    return nb.dcim.mac_addresses.get(id=ma.id)

def create_or_update_ip_address(nb, data):
    ip = nb.ipam.ip_addresses.get(address=data["address"])
    if ip is None:
        ip = nb.ipam.ip_addresses.create(data)
        print(f"Created IP address id={ip.id}, ip={ip.address}")

        return ip

    print(f"Updating IP address id={ip.id}, ip={ip.address}")
    ip.update(data)
    return nb.ipam.ip_addresses.get(id=ip.id)

def process_switch(d: UnifiDevice, db: Db, nb: NetboxCache, site, vrf):
    db.devices.append({
        "name": d.name,
        "device_type": nb.get_device_type("ubiquiti-us-8-150w").id,
        "role": nb.get_device_role("switch").id,
        "serial": d.serial,
        "site_name": site,
    })

    db.interfaces.append({
        "device": find_device(nb.nb, name=d.name),
        "name": "switch0",
        "type": "virtual",
    })

    db.ips.append({
        "address": f"{d.ip}/32",
        "is_primary": "true",
        "vrf": vrf.id,
#        "assigned_object": iface,
        "assigned_object_id": find_interface(nb.nb, device=d.name, name="switch0"),
        "assigned_object_type": "dcim.interface",
    })

    db.macs.append({
        "assigned_object_id": find_interface(nb.nb, device=d.name, name="switch0"),
        "assigned_object_type": "dcim.interface",
        "mac_address": d.mac,
#        "is_primary": "true", does not work
    })

    # TODO: add cables

def sync_db(db: Db, nb):
    def resolve_query(data):
        for k, v in data.items():
            if not isinstance(v, Query):
                continue

            data[k] = v.query()

    for device in db.devices:
        resolve_query(device)
        create_or_update_device(nb, device)

    for iface in db.interfaces:
        resolve_query(iface)
        create_or_update_interface(nb, iface)

    for mac in db.macs:
        resolve_query(mac)
        create_or_update_mac_address(nb, mac)

    for ip in db.ips:
        resolve_query(ip)
        create_or_update_ip_address(nb, ip)

def main():
    unifi_url=os.getenv("UNIFI_URL")
    unifi_username=os.getenv("UNIFI_USERNAME")
    unifi_password=os.getenv("UNIFI_PASSWORD")
    unifi_site=os.getenv("UNIFI_SITE")

    netbox_url=os.getenv("NETBOX_URL")
    netbox_token=os.getenv("NETBOX_TOKEN")
    netbox_vrf_name=os.getenv("NETBOX_VRF")
    netbox_site_name=os.getenv("NETBOX_SITE")

    controller = controller_login(unifi_url, unifi_username, unifi_password)

    (nb, netbox_site, netbox_vrf) = netbox_login(netbox_url, netbox_token, netbox_site_name, netbox_vrf_name)
    status = nb.status()
    print(f"NetBox status: {status}")

    devices = collect_devices(controller, unifi_site)

    nb_cache = NetboxCache(nb)
    db = Db()
    for d in devices:
#        pprint(d)
        if d.model == "US8P150":
            process_switch(d, db, nb_cache, netbox_site, netbox_vrf)

    sync_db(db, nb)

def controller_login(url, username, password) -> UnifiController:
#    try:
        controller = UnifiController(
            controller_url=url,
            username=username,
            password=password,
            is_udm_pro=False,
            verify_ssl=True,
        )

        # Just to check that there is a valid authentication
        controller.get_unifi_site(include_health=False, raw=False)

        return controller
#    except UnifiAuthenticationError:
#        print("Authentication failed - please check your UniFi Controller credentials and URL.")
#    except UnifiAPIError as e:
#        print(f"UniFi API error: {e}")
#    except Exception as e:
#        print(f"An unexpected error occurred: {e}")

def collect_devices(controller: UnifiController, site_name: str) -> list[UnifiDevice]:
    try:
        return controller.get_unifi_site_device(site_name=site_name, detailed=True, raw=False)
    except UnifiAPIError as e:
        print(f"Error fetching device information: {e}")
    except Exception as e:
        print(f"An unexpected error occurred: {e}")

def netbox_login(url: str, token: str, site_name: str, vrf_name: str) -> pynetbox.core.api.Api:
    nb = pynetbox.api(url, token=token)

    site = nb.dcim.sites.get(name=site_name)
    if site is None:
        site = nb.dcim.sites.get(slug=site_name)
        if site is None:
            print(f"Could not look up site by name or slug: {site_name}")
            exit(1)
    print(f"NetBox site {site.name}")

    vrf = None
    vrf_id = None
    if vrf_name is not None:
        vrf = nb.ipam.vrfs.get(site=site, name=vrf_name)
        if vrf is None:
            print(f"Could not look up VRF by slug: {vrf_name}")
            exit(1)
        vrf_id = vrf.id

    return nb, site, vrf

if __name__ == "__main__":
    main()