#!/usr/bin/env python3
import logging
from pathlib import Path

from nornir import InitNornir
from nornir.plugins.tasks import networking, text
from nornir.plugins.tasks.data import load_yaml
from nornir.plugins.functions.text import print_title, print_result
from nornir.core.filter import F

# import ipdb; ipdb.set_trace()


def configure_bgp(task):
    if Path(f'inventory/bgp/{task.host.name}.yaml').is_file():
        # Load additional data (interfaces)
        data = task.run(
            task=load_yaml,
            file=f'inventory/bgp/{task.host.name}.yaml',
            severity_level=logging.DEBUG,
        )

        # Check for route-maps
        for item in data.result['bgp']['address_families']:
            for neighbor in item['neighbors']:
                if 'route_maps' in neighbor:
                    for rm in neighbor['route_maps']:
                        task.run(
                            task=configure_route_map,
                            data=rm,
                        )

        # Transform inventory data to configuration via a template file
        template_config = task.run(
            task=text.template_file,
            name="BGP Configuration",
            template="bgp.j2",
            path=f"templates/{task.host.platform}",
            bgp=data.result['bgp'],
            severity_level=logging.DEBUG,
        )

        # Save the compiled configuration into a host variable
        task.host["config"] = template_config.result

        # Deploy that configuration to the device using NAPALM
        task.run(
            task=networking.napalm_configure,
            name="Loading Configuration on the device",
            configuration=task.host["config"],
            replace=False,
            severity_level=logging.INFO,
        )
    else:
        print(f'No BGP data found for {task.host.name}!')


def configure_route_map(task, data):
    # Transform inventory data to configuration via a template file
    template_config = task.run(
        task=text.template_file,
        name=f"ROUTE-MAP '{data['name']}' Configuration",
        template=f"{data['name']}.j2".lower(),
        path=f"templates/{task.host.platform}/route-maps",
        route_map=data,
        severity_level=logging.DEBUG,
    )

    # Save the compiled configuration into a host variable
    task.host["config"] = template_config.result

    # Deploy that configuration to the device using NAPALM
    task.run(
        task=networking.napalm_configure,
        name="Loading Configuration on the device",
        configuration=task.host["config"],
        replace=False,
        severity_level=logging.INFO,
    )


def main():
    nr = InitNornir(
        config_file="config.yaml",
        # dry_run=True,
        # core={"num_workers": 1},
    )

    devices = nr.filter(
        F(groups__contains="mpls-lab") & ~F(role="switch"),
    )

    print_title(f'Hello, my name is title')

    r = devices.run(
        task=configure_bgp,
    )
    print_result(r)


if __name__ == "__main__":
    main()