#!/usr/bin/env python3

import os
import re
import shutil
import subprocess
import sys


usage = f"""Usage: {sys.argv[0]} <command>

Commands:
  list         List all NTFS volumes available for mounting
  all          Mount all NTFS volumes via ntfs-3g
  <disk id>    Mount a specific NTFS volume via ntfs-3g
"""


def cli(command):
    volumes = get_ntfs_volumes()

    if command == 'list':
        if len(volumes) == 0:
            print("No NTFS volumes found.")

        for id, volume in volumes.items():
            name = f'{id}: {volume["name"]} [{volume["size"]}]'
            details = (
                'mounted: ' + ('yes' if volume['mounted'] else 'no')
                + (' (read-only)' if volume['read-only'] else '')
            )

            print(f"{name} -- {details}")

        sys.exit(0)

    if command == 'all':
        run_checks()
        print(f'Found {len(volumes)} NTFS volume(s).')

        for id, volume in volumes.items():
            print()
            mount(volume)

        sys.exit(0)

    if command in volumes:
        run_checks()
        ok = mount(volumes[command])
        sys.exit(0 if ok else 1)

    print(f'{sys.argv[0]}: Invalid command or disk id.')
    print()
    print(usage)

    sys.exit(1)


def run_checks():
    if shutil.which('ntfs-3g') is None:
        sys.exit('ERROR: Could not find ntfs-3g.')

    if not os.geteuid() == 0:
        sys.exit('ERROR: Need root privilages to mount via ntfs-3g.')


def get_ntfs_volumes():
    list_out = subprocess.run(['diskutil', 'list'], capture_output=True, check=True).stdout.decode()

    disk_ids = [
        re.search(r'\S+$', line)[0]
        for line in list_out.split('\n')
        if re.match(r'^\s*\d+:\s*Windows_NTFS ', line)
    ]

    volumes = {}

    for id in disk_ids:
        info_out = subprocess.run(['diskutil', 'info', id], capture_output=True, check=True).stdout.decode()

        info = {
            line.split(':', 1)[0].strip(): line.split(':', 1)[1].strip()
            for line in info_out.split('\n') if line != ''
        }

        if (
            info['Type (Bundle)'] == 'ntfs'
            and info['File System Personality'] == 'NTFS'
            and info['Read-Only Media'] == 'No'
        ):
            volumes[id] = {
                'id': id,
                'node': info['Device Node'],
                'name': info['Volume Name'] if info['Volume Name'] != '' else 'Untitled',
                'mounted': info['Mounted'] == 'Yes',
                'size': re.match(r'\d+(\.\d+)? \S+', info['Disk Size'])[0],
                'read-only': info['Read-Only Volume'] == 'Yes',
            }

    return volumes


def mount(volume):
    print(f'Volume: {volume["name"]} [{volume["size"]}]')

    if volume['mounted']:
        if not volume['read-only']:
            print(f'{volume["name"]} is already writable.')
            return False

        print('Unmounting...')
        try:
            subprocess.run(['diskutil', 'unmount', volume['id']], check=True)
        except:
            return False

    path = f'/Volumes/{volume["name"]}'

    if os.path.exists(path):
        counter = 1
        while os.path.exists(path + ' ' + str(counter)):
            counter += 1

        path = path + ' ' + str(counter)

    user_id = os.getenv('SUDO_UID', os.getuid())
    group_id = os.getenv('SUDO_GID', os.getgid())

    mount_cmd = [
        'ntfs-3g',
        '-o', f'volname={volume["name"]}',
        '-o', 'local',
        '-o', 'allow_other',
        '-o', 'auto_xattr',
        '-o', f'uid={user_id}',
        '-o', f'gid={group_id}',
        '-o', 'windows_names',
        volume['node'], path,
    ]

    print(f'Mounting on `{path}` via ntfs-3g...')
    try:
        subprocess.run(mount_cmd, check=True)
        print(f'Successfully mounted {volume["name"]}.')
        return True
    except:
        print(f'Failed to mount {volume["name"]}.')

        if volume['mounted']:
            print('Remounting via macOS...')
            subprocess.run(['diskutil', 'mount', volume['id']], check=True)

        return False


def main():
    if len(sys.argv) < 2:
        print(usage)
        sys.exit(1)

    cli(sys.argv[1])


if __name__ == '__main__':
    main()
