from extras.scripts import *
from django.core.exceptions import ValidationError
from django.contrib.contenttypes.models import ContentType

from dcim.models import Cable, CableTermination, Device, DeviceType, Location, DeviceRole, Site, Interface
from dcim.choices import InterfaceModeChoices, InterfaceTypeChoices

from ipam.models import VLANGroup, VLAN, Role, Prefix, IPAddress
from ipam.choices import PrefixStatusChoices

DEFAULT_SITE = Site.objects.get(name='Vikingskipet')
DEFAULT_DEVICE_TYPE = DeviceType.objects.get(model='EX2200-48T-4G')
DEFAULT_DEVICE_ROLE = DeviceRole.objects.get(slug='access-switch')
DEFAULT_TG_DNS_SUFFIX = "tg25.tg.no"
DEFAULT_UPLINK_SWITCH = Device.objects.get(name='d1.ring')

DEVICE_ROLE_LEAF = "leaf"
DEVICE_ROLE_DISTRO = "distro"
DEVICE_ROLE_UTSKUTT_DISTRO = "utskutt-distro"

# VLAN Group to allocate VLANs from
FABRIC_VLAN_GROUP = VLANGroup.objects.get(slug='client-vlans')

# Vlan role for fabric clients
FABRIC_CLIENTS_ROLE = Role.objects.get(slug='clients')

# Client networks allocated from here
FABRIC_V4_CLIENTS_PREFIX = Prefix.objects.get(prefix='10.25.0.0/16')
FABRIC_V6_CLIENTS_PREFIX = Prefix.objects.get(prefix='2a06:5844:e::/48')

# Switch mgmt allocates from here
FABRIC_V4_JUNIPER_MGMT_PREFIX = Prefix.objects.get(prefix='185.110.149.0/25')
FABRIC_V6_JUNIPER_MGMT_PREFIX = Prefix.objects.get(prefix='2a06:5841:f::/64')

UPLINK_PORTS = {
    'EX2200-48T-4G': ["ge-0/0/44", "ge-0/0/45", "ge-0/0/46", "ge-0/0/47"],
}

UPLINK_TYPES = (
    (InterfaceTypeChoices.TYPE_10GE_SFP_PLUS, '10G SFP+'),
    (InterfaceTypeChoices.TYPE_1GE_FIXED, '1G RJ45'),
    (InterfaceTypeChoices.TYPE_10GE_FIXED, '10G RJ45')
)


def generatePrefix(prefix, length):
    firstPrefix = prefix.get_first_available_prefix()
    out = list(firstPrefix.subnet(length, count=1))[0]
    return out


class CreateSwitch(Script):
    class Meta:
        name = "Create Switch"
        description = "Provision a new switch"
        commit_default = True
        field_order = ['site_name', 'switch_count', 'switch_model']
        fieldsets = ""
        scheduling_enabled = False

    switch_name = StringVar(
        description="Switch name. Remember, e = access switch, d = distro switch",
        required=True,
        default="e1.test"  # default during development
    )
    device_type = ObjectVar(
        description="Device model",
        model=DeviceType,
        default=DEFAULT_DEVICE_TYPE.id,
    )
    device_role = ObjectVar(
        description="Device role",
        model=DeviceRole,
        default=DEFAULT_DEVICE_ROLE.id,
    )
    location = ObjectVar(
        model=Location,
        required=True,
        default=Location.objects.get(name="Ringen")  # Default during development
    )
    uplink_type = ChoiceVar(
        label='Uplink Type',
        required=True,
        description="What type of interface should this switch be delivered on",
        choices=UPLINK_TYPES,
        default=InterfaceTypeChoices.TYPE_1GE_FIXED
    )
    destination_device_a = ObjectVar(
        description="Uplink device (A)",
        required=True,
        model=Device,
        query_params={
            'role': [DEVICE_ROLE_LEAF, DEVICE_ROLE_DISTRO, DEVICE_ROLE_UTSKUTT_DISTRO],
        },
        default=DEFAULT_UPLINK_SWITCH
    )
    destination_device_b = ObjectVar(
        description="If connected to leaf pair - Uplink device (B)",
        required=False,
        model=Device,
        query_params={
            'role': [DEVICE_ROLE_LEAF],
        },
    )
    # If leaf pair we assume same port. This input only cares about cases with single device.
    destination_interfaces = MultiObjectVar(
        description="Destination interface(s). \n\n IF You're looking at d1.ring: ge-{PLACEMENT}/x/x. Placements: 0 = South, 1 = Log, 2 = Swing, 3 = North, 4 = noc, 5 = tele",
        model=Interface,
        query_params={
            'device_id': '$destination_device_a',
            'occupied': False,
            'type': '$uplink_type'
        }
    )

    def run(self, data, commit):
        if not data['switch_name'].startswith("e") and not data['switch_name'].startswith("d"):
            raise ValidationError("Switch name must start whit e or d")

        switch = self.create_switch(data)

        vlan = self.create_vlan(switch)

        v4_prefix, v6_prefix = self.allocate_prefixes(vlan)

        self.set_traffic_vlan(switch, vlan)

        self.connect_switch(data, switch, vlan)

        self.log_success(f"✅ Script completed successfully.")
        self.log_success(f"🔗 Switch:     <a href=\"{switch.get_absolute_url()}\">{switch}</a>")
        self.log_success(f"🔗 v6 Prefix:  <a href=\"{v6_prefix.get_absolute_url()}\">{v6_prefix}</a>")
        self.log_success(f"🔗 v4 Prefix:  <a href=\"{v4_prefix.get_absolute_url()}\">{v4_prefix}</a>")
        self.log_success(f"🔗 VLAN:       <a href=\"{vlan.get_absolute_url()}\">{vlan}</a>")
        self.log_success(f"⚠️ <strong>️Fabric config must be deployed before switch can be fapped.</strong>")

    def allocate_prefixes(self, vlan):
        v6_prefix = Prefix.objects.create(
            prefix=generatePrefix(FABRIC_V6_CLIENTS_PREFIX, 64),
            status=PrefixStatusChoices.STATUS_ACTIVE,
            role=FABRIC_CLIENTS_ROLE,
            vlan=vlan
        )
        v4_prefix = Prefix.objects.create(
            prefix=generatePrefix(FABRIC_V4_CLIENTS_PREFIX, 26),
            status=PrefixStatusChoices.STATUS_ACTIVE,
            role=FABRIC_CLIENTS_ROLE,
            vlan=vlan
        )
        self.log_info("Created network. Created new VLAN and assigned prefixes")
        return v4_prefix, v6_prefix

    def connect_switch(self, data, switch, vlan):
        uplink_description = f"B: {data['destination_device_a'].name}"
        if data['destination_device_b']:
            uplink_description = f"B: {data['destination_device_a'].name} / {data['destination_device_b'].name} - ae{vlan.id}"
        uplink_ae = Interface.objects.create(
            device=switch,
            name="ae0",
            description=uplink_description,
            type=InterfaceTypeChoices.TYPE_LAG,
            mode=InterfaceModeChoices.MODE_TAGGED,
        )
        if data['destination_device_a'].role == DEVICE_ROLE_UTSKUTT_DISTRO:
            self.log_debug(f"{data['destination_device_a']} is utskutt-distro")
            # TODO make sure we add traffic vlan on AE between distro and utskutt-distro as well.
        uplink_ae.tagged_vlans.add(FABRIC_V4_JUNIPER_MGMT_PREFIX.vlan.id)
        uplink_ae.tagged_vlans.add(vlan.id)
        ## We only need this if not connected to leaf (since they are provisioned using AVD)
        if data['destination_device_a'].role != DEVICE_ROLE_LEAF:
            destination_ae = Interface.objects.create(
                device=data['destination_device_a'],
                name=f"ae{vlan.id}",
                description=f'B: {switch.name} ae0',
                type=InterfaceTypeChoices.TYPE_LAG,
                mode=InterfaceModeChoices.MODE_TAGGED,
            )
            destination_ae.save()
            destination_ae.tagged_vlans.add(FABRIC_V4_JUNIPER_MGMT_PREFIX.vlan.id)
            destination_ae.tagged_vlans.add(vlan.id)
            self.log_debug(
                f"Created destination AE <a href=\"{destination_ae.get_absolute_url()}\">{destination_ae}</a>")

            ## TODO support leaf pair
            num_uplinks = len(data['destination_interfaces'])
            uplink_interfaces = list(Interface.objects.filter(device=switch, type=data['uplink_type']))
            if len(uplink_interfaces) < 1:
                raise AbortScript(
                    f"You chose a device type without any {data['uplink_type']} interfaces! Pick another model :)")

            interface_type = ContentType.objects.get_for_model(Interface)
            interfaces_filtered = [interface for interface in uplink_interfaces if
                                   interface.name in UPLINK_PORTS.get(switch.device_type.model, [])]
            for uplink_num in range(0, num_uplinks):
                a_interface = data['destination_interfaces'][uplink_num]
                b_interface = interfaces_filtered[uplink_num]

                self.log_debug(
                    f"Connecting: {data['destination_device_a']} - {a_interface} to {switch} - {b_interface}")

                a_interface.description = f'G: {switch.name} {b_interface.name} (ae0)'
                b_interface.description = f"G: {data['destination_device_a'].name} {a_interface.name} (ae{vlan.id})"

                b_interface.lag = uplink_ae
                b_interface.save()

                a_interface.lag = destination_ae
                a_interface.save()

                cable = Cable.objects.create()
                a = CableTermination.objects.create(
                    cable=cable,
                    cable_end='A',
                    termination_id=a_interface.id,
                    termination_type=interface_type,
                )
                b = CableTermination.objects.create(
                    cable_end='B',
                    cable=cable,
                    termination_id=b_interface.id,
                    termination_type=interface_type,
                )
                cable = Cable.objects.get(id=cable.id)
                # https://github.com/netbox-community/netbox/discussions/10199
                cable._terminations_modified = True
                cable.save()

    def set_traffic_vlan(self, switch, vlan):
        interfaces = list(Interface.objects.filter(device=switch, type=InterfaceTypeChoices.TYPE_1GE_FIXED))
        if len(interfaces) == 0:
            self.log_error(f"no interfaces found")

        for interface in interfaces:
            if interface.name in UPLINK_PORTS.get(switch.device_type.model, []):
                continue
            interface.mode = 'access'
            interface.untagged_vlan = vlan
            interface.description = "C: Clients"
            interface.save()
        self.log_info("Configured traffic vlan on all client ports")

    def create_vlan(self, switch):
        vid = FABRIC_VLAN_GROUP.get_next_available_vid()
        vlan = VLAN.objects.create(
            name=switch.name,
            group=FABRIC_VLAN_GROUP,
            role=FABRIC_CLIENTS_ROLE,
            vid=vid
        )
        vlan.save()
        self.log_info("Created VLAN")
        return vlan

    def create_switch(self, data):
        switch = Device(
            name=data['switch_name'],
            device_type=data['device_type'],
            location=data['location'],
            role=data['device_role'],
            site=DEFAULT_SITE
        )
        switch.save()

        mgmt_interface_name = "irb"
        if switch.device_type.model == "EX2200-48T-4G":
            mgmt_interface_name = "vlan"

        mgmt_vlan_interface = Interface.objects.create(
            device=switch,
            name=f"{mgmt_interface_name}.{FABRIC_V4_JUNIPER_MGMT_PREFIX.vlan.id}",
            description=f'X: Mgmt',
            type=InterfaceTypeChoices.TYPE_VIRTUAL,
            mode=InterfaceModeChoices.MODE_TAGGED,
        )
        v4_mgmt_addr = IPAddress.objects.create(
            address=FABRIC_V4_JUNIPER_MGMT_PREFIX.get_first_available_ip(),
            dns_name=f"{switch.name}.{DEFAULT_TG_DNS_SUFFIX}"

        )
        v6_mgmt_addr = IPAddress.objects.create(
            address=FABRIC_V6_JUNIPER_MGMT_PREFIX.get_first_available_ip(),
            dns_name=f"{switch.name}.{DEFAULT_TG_DNS_SUFFIX}"
        )
        mgmt_vlan_interface.ip_addresses.add(v4_mgmt_addr)
        mgmt_vlan_interface.ip_addresses.add(v6_mgmt_addr)
        switch.primary_ip4 = v4_mgmt_addr
        switch.primary_ip6 = v6_mgmt_addr
        switch.save()

        self.log_info("Created switch")
        self.log_info("Allocated and assigned mgmt addresses on switch")
        return switch


script = CreateSwitch