import os,sys

from pathlib import Path
from collections import deque, defaultdict

import argparse

from yaml import load, YAMLError
try:
    from yaml import CLoader as Loader
except ImportError:
    from yaml import Loader


SOPS_TYPES = frozenset({'kms', 'gcp_kms', 'azure_kv', 'hc_vault', 'age', 'pgp'})


class BooleanAction(argparse.Action):
    def __init__(self, option_strings, dest, nargs=None, **kwargs):
        super(BooleanAction, self).__init__(option_strings, dest, nargs=0, **kwargs)

    def __call__(self, parser, namespace, values, option_string=None):
        setattr(namespace, self.dest, False if option_string.startswith('--no') else True)


def main():
    default_base = os.getenv('SOPS_INVENTORY_BASE', default=[])
    if default_base:
        default_base = Path(default_base)

    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--list-files', '--no-list-files', action=BooleanAction, default=False, help='Only list sops files')
    parser.add_argument('path', metavar='PATH', nargs='?' if default_base else None, type=Path, default=default_base, help='Base directory to take inventory of')
    args = parser.parse_args()

    inventory = defaultdict(set)

    queue = deque([args.path])
    while queue:
        baseDir = queue.popleft()
        for child in baseDir.iterdir():
            if child.is_dir():
                queue.append(child)
            else:
                try:
                    with child.open(mode='r') as fh:
                        yaml = load(fh, Loader=Loader)
                        if not yaml:
                            raise ValueError('Could not parse YAML')
                        if not isinstance(yaml, dict) or not 'sops' in yaml:
                            raise ValueError('Did not find "sops" key')
                        sops = yaml['sops']

                        key_info = set()
                        for k in SOPS_TYPES:
                            if k in sops:
                                v = sops[k]
                                if not v:
                                    continue

                                match k:
                                    case 'pgp':
                                      for r in v:
                                          key_info.add(r['fp'])
                                    case 'age':
                                      for r in v:
                                          key_info.add(r['recipient'])
                                    case _:
                                      raise NotImplementedError
                        inventory[frozenset(key_info)].add(child.relative_to(args.path))
                except (YAMLError, ValueError) as e:
                    pass

    if not args.list_files:
        for keys, files in inventory.items():
            print(','.join(keys) + ':')
            for file in files:
                print('  - ' + str(file))
    else:
        for _, files in inventory.items():
            for file in files:
                print(file)

if __name__ == '__main__':
    os.exit(main())