#!/usr/bin/env python3

from json import load
from os import environ
from os.path import join
from sys import argv, exit
from tempfile import gettempdir

from mako.template import Template

from bundlewrap.repo import Repository
from bundlewrap.utils.text import bold
from bundlewrap.utils.ui import io

NTP_SERVERS = {
# pool.ntp.org
    '148.251.54.81',
    '162.159.200.123',
    '213.209.109.44',
    '54.36.110.36',
}

try:
    node_name = argv[1]
except Exception:
    print(f'Usage: {argv[0]} <node name>')
    exit(1)

path = environ.get('BW_REPO_PATH', '.')
repo = Repository(path)
node = repo.get_node(node_name)

try:
    io.activate()

    interfaces = {}
    users = {}
    vlans = {
        'default': {
            'id': None,
            'ip_address': '169.254.254.254/24',
        },
    }

    tmpfile = join(gettempdir(), f'{node.name}.conf')

    gw_split = node.hostname.split('.')
    gw_split[3] = '1'
    gateway = '.'.join(gw_split)

    with io.job('reading netbox_dump.json'):
        with open(join(repo.path, 'netbox_dump.json'), 'r') as f:
            json = load(f)[node.metadata.get('location')]

        for vlan, vid in json['vlans'].items():
            vlans[vlan] = {
                'id': vid,
                'ip_address': None,
            }

        for iface, iconfig in json['devices'][node.name].items():
            if iface in vlans:
                # If the interface name is the same as a vlan name, this
                # means the ip assigned to this interface should get
                # assigned to that vlan.
                vlans[iface]['ip_address'] = iconfig['ip_addresses'][0]
            else:
                interfaces[iface] = {
                    'enabled': bool(
                        iconfig['enabled']
                        and iconfig['mode']
                        and (
                            iconfig['vlans']['tagged']
                            or iconfig['vlans']['untagged']
                        )
                    ),
                    'description': iconfig['description'],
                    'untagged_vlan': iconfig['vlans']['untagged'],
                }

                if iconfig['mode'] and iconfig['mode'].startswith('tagged'):
                    interfaces[iface]['mode'] = 'trunk'
                else:
                    interfaces[iface]['mode'] = 'access'

                tagged_vlans = set()
                for vlan in iconfig['vlans']['tagged']:
                    tagged_vlans.add(str(vlans[vlan]['id']))
                interfaces[iface]['tagged_vlans'] = tagged_vlans

    with io.job('reading users.json'):
        with open(join(repo.path, 'users.json'), 'r') as f:
            json = load(f)

        users = {}
        for uname, config in json.items():
            if config.get('is_admin', False):
                users[uname] = {
                    'password': repo.vault.password_for(f'{node.name} {uname} login'),
                    'ssh_pubkey': set(config['ssh_pubkey']),
                }


    with io.job(f'{bold(node.name)}  rendering config template to {tmpfile}'):
        with open(join(repo.path, 'configs', 'junos-template.conf')) as f:
            template = Template(
                f.read().encode('utf-8'),
                input_encoding='utf-8',
                output_encoding='utf-8',
            )
        content = template.render(
            gateway=gateway,
            interfaces=interfaces,
            node=node,
            ntp_servers=NTP_SERVERS,
            repo=repo,
            users=users,
            vlans=vlans,
        )
        with open(tmpfile, 'w+') as f:
            f.write(content.decode('utf-8'))

    with io.job(f'{bold(node.name)}  updating configuration on device'):
        node.upload(tmpfile, '/tmp/bundlewrap.conf')

        result = node.run(
            'configure exclusive ; load override /tmp/bundlewrap.conf ; commit',
            log_output=True,
        )

        if 'commit complete' in result.stdout.decode():
            node.run(
                'request system configuration rescue save',
                log_output=True,
            )
finally:
    io.deactivate()