from dcim.models import Device, Interface
from extras.models import CustomField
from extras.scripts import *

class CreateMLAGScript(Script):
    class Meta:
        name = 'Create MLAG'
        description = 'Create and configure a new MLAG on paired Cumulus leaf switches'
        scheduling_enabled = False

    switch = ObjectVar(model=Device, query_params={'role': 'switch'},
            description='One of the peer switches')
    iface = ObjectVar(label='Interface', model=Interface, query_params={'device_id': '$switch'},
            description='Switch interface (same for both switches)')
    bond_name = StringVar(label='MLAG name', regex='[0-9a-z-]', max_length=15)
    mode = ChoiceVar(choices=CustomField.objects.get(name='bond_mode').choices, required=False)

    def run(self, data, commit):
        switch = data['switch']
        iface_name = data['iface'].name
        bond_name = data['bond_name']
        mode = data['mode']

        peer_name = switch.local_context_data.get('peer')
        if not peer_name:
            raise AbortScript(f'Switch {switch} has no peer defined in local context')

        for device in (switch, Device.objects.get(name=peer_name)):
            bridge = device.interfaces.get(name='bridge')
            iface = device.interfaces.get(name=iface_name)

            bond_iface, new = Interface.objects.get_or_create(device=device, name=bond_name, type='lag', bridge=bridge)
            if new:
                self.log_info(f'created new bond {bond_iface} on {device}')
            if mode:
                bond_iface.custom_field_data['bond_mode'] = mode
                self.log_info(f'setting mode for bond {bond_iface} to {mode}')
            bond_iface.full_clean()
            bond_iface.save()

            iface.enabled = True
            iface.lag = bond_iface
            self.log_info(f'setting {bond_iface} as LAG for interface {iface}')
            iface.full_clean()
            iface.save()