#!/usr/bin/env python3
#
# Copyright (C) 2024 VyOS maintainers and contributors
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 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/>.

import sys
import typing
import vyos.opmode

from vyos.configquery import ConfigTreeQuery
from vyos.utils.process import call
from vyos.utils.commit import commit_in_progress

config = ConfigTreeQuery()

service_map = {
    'dhcp' : {
        'systemd_service': 'kea-dhcp4-server',
        'path': ['service', 'dhcp-server'],
    },
    'dhcpv6' : {
        'systemd_service': 'kea-dhcp6-server',
        'path': ['service', 'dhcpv6-server'],
    },
    'dns_dynamic': {
        'systemd_service': 'ddclient',
        'path': ['service', 'dns', 'dynamic'],
    },
    'dns_forwarding': {
        'systemd_service': 'pdns-recursor',
        'path': ['service', 'dns', 'forwarding'],
    },
    'igmp_proxy': {
        'systemd_service': 'igmpproxy',
        'path': ['protocols', 'igmp-proxy'],
    },
    'ipsec': {
        'systemd_service': 'strongswan',
        'path': ['vpn', 'ipsec'],
    },
    'mdns_repeater': {
        'systemd_service': 'avahi-daemon',
        'path': ['service', 'mdns', 'repeater'],
    },
    'reverse_proxy': {
        'systemd_service': 'haproxy',
        'path': ['load-balancing', 'reverse-proxy'],
    },
    'router_advert': {
        'systemd_service': 'radvd',
        'path': ['service', 'router-advert'],
    },
    'snmp' : {
        'systemd_service': 'snmpd',
    },
    'ssh' : {
        'systemd_service': 'ssh',
    },
    'suricata' : {
        'systemd_service': 'suricata',
    },
    'vrrp' : {
        'systemd_service': 'keepalived',
        'path': ['high-availability', 'vrrp'],
    },
    'webproxy' : {
        'systemd_service': 'squid',
    },
}
services = typing.Literal['dhcp', 'dhcpv6', 'dns_dynamic', 'dns_forwarding', 'igmp_proxy', 'ipsec', 'mdns_repeater', 'reverse_proxy', 'router_advert', 'snmp', 'ssh', 'suricata' 'vrrp', 'webproxy']

def _verify(func):
    """Decorator checks if DHCP(v6) config exists"""
    from functools import wraps

    @wraps(func)
    def _wrapper(*args, **kwargs):
        config = ConfigTreeQuery()
        name = kwargs.get('name')
        human_name = name.replace('_', '-')

        if commit_in_progress():
            print(f'Cannot restart {human_name} service while a commit is in progress')
            sys.exit(1)

        # Get optional CLI path from service_mapping dict
        # otherwise use "service name" CLI path
        path = ['service', name]
        if 'path' in service_map[name]:
            path = service_map[name]['path']

        # Check if config does not exist
        if not config.exists(path):
            raise vyos.opmode.UnconfiguredSubsystem(f'Service {human_name} is not configured!')
        if config.exists(path + ['disable']):
            raise vyos.opmode.UnconfiguredSubsystem(f'Service {human_name} is disabled!')
        return func(*args, **kwargs)

    return _wrapper

@_verify
def restart_service(raw: bool, name: services, vrf: typing.Optional[str]):
    systemd_service = service_map[name]['systemd_service']
    if vrf:
        call(f'systemctl restart "{systemd_service}@{vrf}.service"')
    else:
        call(f'systemctl restart "{systemd_service}.service"')

if __name__ == '__main__':
    try:
        res = vyos.opmode.run(sys.modules[__name__])
        if res:
            print(res)
    except (ValueError, vyos.opmode.Error) as e:
        print(e)
        sys.exit(1)