#!/usr/bin/env python3

import argparse
from os import path, makedirs
from shutil import rmtree


def folder(default_search, user_input):
    if path.sep in user_input:
        return user_input
    else:
        the_path = path.dirname(path.dirname(path.abspath(__file__)))
        the_path = path.join(the_path, default_search, user_input)
        return the_path


parser = argparse.ArgumentParser(description='''
    Extract an "original" source code tree from an HFS disk image, or
    (on macOS) a folder. If a destination is given that does not contain
    a path separator, it will created as a subdirectory of ../base/
''')

parser.add_argument('src', metavar='SOURCE', action='store', help='Disk image, archive, folder, whatever')
parser.add_argument('--out', '-o', metavar='DEST', action='store', default='SuperMarioProj.1994-02-09', type=lambda x: folder('base', x), help='Processed source tree (default=SuperMarioProj.1994-02-09)')

args = parser.parse_args()


if path.splitext(args.src)[1].lower() in ('.dmg', '.img', '.dsk'):
    from machfs import Volume

    with open(args.src, 'rb') as f:
        v = Volume()
        v.read(f.read())

    # slight hack: remove nonprinting chars from text files
    for parent, child_dirnames, child_filenames in v.walk():
        for c in child_filenames:
            o = v[parent + (c,)]
            if o.type == b'TEXT':
                nudata = bytes(b for b in o.data if b >= 32 or b == 9 or b == 13)
                if nudata != o.data:
                    print(parent + (c,), len(o.data), len(nudata))
                o.data = nudata

    try:
        rmtree(args.out)
    except FileNotFoundError:
        pass
    makedirs(args.out, exist_ok=True)
    v.write_folder(args.out)

else:
    raise NotImplementedError()