#!/usr/bin/env python3
#
# Copyright VyOS maintainers and contributors <maintainers@vyos.io>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 or later as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

# T7760: Remove per VRF setting for system-as option and replace it with
#        local-as if required.

from vyos.configtree import ConfigTree

def migrate(config: ConfigTree) -> None:
    vrf_base = ['vrf', 'name']
    bgp_base = ['protocols', 'bgp']

    if not config.exists(vrf_base):
        return

    global_asn = None
    if config.exists(bgp_base + ['system-as']):
        global_asn = config.return_value(bgp_base + ['system-as'])

    for vrf in config.list_nodes(vrf_base):
        # bail out early if there is no per VRF BGP instance defined
        vrf_bgp_base = vrf_base + [vrf] + bgp_base
        if not config.exists(vrf_bgp_base):
            continue

        # This is a mandatory node in the old design but keep it optional
        # if one want's to load a weird config
        if config.exists(vrf_bgp_base + ['system-as']):
            system_as = config.return_value(vrf_bgp_base + ['system-as'])
            config.delete(vrf_bgp_base + ['system-as'])

            # If there is no existing global BGP instance - start one
            if not config.exists(bgp_base):
                config.set(bgp_base + ['system-as'], value=system_as)
                global_asn = system_as

            vrf_neighbor_base = vrf_bgp_base + ['neighbor']
            vrf_peer_group_base = vrf_bgp_base + ['peer-group']
            if config.exists(vrf_neighbor_base):
                for neighbor in config.list_nodes(vrf_neighbor_base):
                    # Neighbor already has local-as option set, do not touch it
                    if config.exists(vrf_neighbor_base + [neighbor, 'local-as']):
                        #print(f'VRF {vrf} BGP neighbor {neighbor} has local-as set - do not migrate!')
                        continue

                    # Check if the neighbor uses a peer-group which has the local-as
                    # option set, do not touch it either
                    peer_group = None
                    if config.exists(vrf_neighbor_base + [neighbor, 'peer-group']):
                        peer_group = config.return_value(vrf_neighbor_base + [neighbor, 'peer-group'])
                        # Check if the peer-group has a local-as option set
                        if config.exists(vrf_peer_group_base + [peer_group, 'local-as']):
                            #print(f'VRF {vrf} BGP neighbor {neighbor} uses peer-group {peer_group} which has local-as set - do not migrate!')
                            continue

                    # BGP local-as option is only allowed for eBGP speakers
                    if global_asn == system_as:
                        continue

                    config.set(vrf_neighbor_base + [neighbor, 'local-as', system_as, 'no-prepend', 'replace-as'])
                    config.set_tag(vrf_neighbor_base + [neighbor, 'local-as'])

                    # We do also need to take care about BGP internas. When using local-as routes with our own AS
                    # previously in the path will get rejected:
                    # bgpd: x.x.x.x(Unknown) rcvd UPDATE about 192.0.2.0/24 IPv4 unicast -- DENIED due to: as-path contains our own AS;
                    # Set allowas-in option
                    allowas_in_numer = ['allowas-in', 'number']
                    for afi in ['ipv4-labeled-unicast', 'ipv4-multicast', 'ipv4-unicast', 'ipv4-vpn',
                                'ipv6-labeled-unicast', 'ipv6-multicast', 'ipv6-unicast', 'ipv6-vpn']:
                        afi_neighbor_base = vrf_neighbor_base + [neighbor, 'address-family', afi]
                        afi_peer_group_base = vrf_peer_group_base + [peer_group, 'address-family', afi]

                        # No need to change anything on an AFI not in use
                        if not config.exists(afi_neighbor_base) and not config.exists(afi_peer_group_base):
                            continue

                        allowas_in_value = 0
                        print(neighbor, afi)
                        # Check if there is any allowas-in definition for a peer-group
                        if peer_group and config.exists(afi_peer_group_base + allowas_in_numer):
                            allowas_in_value = int(config.return_value(afi_peer_group_base + allowas_in_numer))
                            #print(f'peer-group {peer_group} allowas-in for {afi} is {allowas_in_value}')

                        # Per neighbor "allowas-in" definition takes higher precendence
                        if config.exists(afi_neighbor_base + allowas_in_numer):
                            allowas_in_value = int(config.return_value(afi_neighbor_base + allowas_in_numer))
                            #print(f'neighbor {neighbor} allowas-in for {afi} is {allowas_in_value}')

                        # Increment allowas-in by 1 as we now have one more entry
                        allowas_in_value += 1
                        # Clip allowas-in to 10 - max supported by FRR platform
                        if allowas_in_value > 10: allowas_in_value = 10

                        # Set per neighbor allowas-in which always takes precedence over the peer-group definition
                        config.set(afi_neighbor_base + allowas_in_numer, value=allowas_in_value, replace=True)