#!/usr/bin/env python3

import subprocess
import os
import re
import argparse

FSTAB_PATH = '/etc/fstab'

def parse_fstab():
    mounts = []
    with open(FSTAB_PATH, 'r') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            parts = re.split(r'\s+', line)
            if len(parts) >= 3 and parts[2] == 'nfs':
                source, target, fstype = parts[:3]
                mounts.append(target)
    return mounts

def get_mounted():
    result = subprocess.run(['mount'], stdout=subprocess.PIPE, text=True)
    mounts = set()
    for line in result.stdout.splitlines():
        parts = line.split()
        if len(parts) >= 3:
            mounts.add(parts[2])
    return mounts

def restart_services(services):
    for service in services:
        print(f"Restarting service: {service}")
        try:
            subprocess.run(['systemctl', 'restart', service], check=True)
        except subprocess.CalledProcessError as e:
            print(f"Failed to restart {service}: {e}")

def mount_missing(services_to_restart):
    fstab_mounts = parse_fstab()
    mounted = get_mounted()
    something_mounted = False

    for mountpoint in fstab_mounts:
        if not os.path.ismount(mountpoint):
            print(f"Mounting missing NFS mount: {mountpoint}")
            try:
                subprocess.run(['mount', mountpoint], check=True)
                something_mounted = True
            except subprocess.CalledProcessError as e:
                print(f"Failed to mount {mountpoint}: {e}")

    if something_mounted and services_to_restart:
        restart_services(services_to_restart)

def main():
    parser = argparse.ArgumentParser(description='Mount missing NFS from fstab and optionally restart systemd services.')
    parser.add_argument('--restart-service', action='append', help='Systemd service to restart if any mount succeeds. Can be used multiple times.', default=[])
    args = parser.parse_args()

    mount_missing(args.restart_service)

if __name__ == '__main__':
    main()