#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Streamlined Wizard for dumping games with the INL Retro cartridge dumper"""

__author__ = "Stephan Sokolow (deitarion/SSokolow)"
__appname__ = "INL Retro Dumping Helper"
__version__ = "0.1"
__license__ = "MIT"

# pylint: disable=bad-builtin
import csv, hashlib, itertools, math, os, re, shlex, subprocess, sys
from xml.etree import ElementTree as ET

if sys.version_info.major < 3:
    print("This script requires Python 3.x")
    sys.exit(1)

# Get width of the terminal window on Linux. Fall back to assuming 80 columns.
try:
    TERM_COLUMNS = int(os.environ.get('COLUMNS'))
except (TypeError, ValueError):
    TERM_COLUMNS = 80

# Used for hashing
CHUNK_SIZE = 2**16

basedir = os.path.dirname(os.path.abspath(__file__))
# TODO: Make the info_db_path values at least somewhat configurable

# -- Helper Functions --

def col_menu(prompt, options, label_getter=str):
    """Code to generate an arbitrary columnar menu using provided choices

    TODO: On Linux, this can be inelegantly wide and short. Add support for
          a `min_height` option which restricts the maximum number of columns
          to satisfy a constraint on the minimum number of rows.
    """
    # Don't bother displaying a menu if there's only one choice
    if len(options) == 1:
        return options[0]

    # Calculate the maximum length of an option name
    labels = [label_getter(x) for x in options]
    opt_width = max(len(x) for x in labels)

    # Calculate how many columns will fit per row, and how wide they should be
    col_width = len('00) ') + opt_width + len('  ')
    col_count = TERM_COLUMNS // col_width

    while True:
        # Calculate how many rows there will be, given that number of columns
        # FIXME: Less confusing variable names for this math here
        row_len = math.ceil(len(labels) / col_count)
        rows = math.ceil(len(labels) / row_len)

        # Adjust for aesthetics (avoid very wide one- or two-row menus)
        if col_count < 2 or row_len > 2:
            break
        else:
            col_count -= 1

    # Convert the list of names into a list of (number, name) tuples
    enumerated = [pair for pair in enumerate(labels)]

    # Break the list of pairs up into column-sized chunks
    chunked = [enumerated[i:i + row_len]
               for i in range(0, len(enumerated), row_len)]

    # Use a matrix transpose (rotate the 2D array about its diagonal axis)
    # to switch from row-major order ("1,2,3" runs along the first row) to
    # column-major order ("1,2,3" runs down the first column).
    #
    # (missing cells will be filled with `None`)
    rows = itertools.zip_longest(*chunked)

    # Render the menu
    print("")
    for row in rows:
        row_formatted = []
        for pair in row:
            if not pair:
                continue

            row_formatted.append(
                "{:>2}) {:<{width}}".format(pair[0] + 1, pair[1],
                                            width=opt_width))

        print('  '.join(row_formatted))

    # Prompt for input
    return options[prompt_for(prompt,
                              convert=int,
                              test=lambda x: 1 <= x <= len(options)) - 1]

def hash_file(path, hasher=hashlib.sha1, seek=0):
    """Generate a hash for a potentially long file.
    Accepts paths and file-like objects.

    Digesting will obey CHUNK_SIZE to conserve memory.
    """
    with open(path, 'rb') as fobj:
        fhash = hasher()
        fobj.seek(seek)

        # Chunked digest generation (conserve memory)
        for block in iter(lambda: fobj.read(CHUNK_SIZE), b''):
            fhash.update(block)

        return fhash.hexdigest()

def make_filename(title):
    """Generate a filename that's valid on Windows from a title"""
    win32_special_names = ['CON', 'PRN', 'AUX', 'NUL', 'COM1', 'COM2', 'COM3',
        'COM4', 'COM5', 'COM6', 'COM7', 'COM8', 'COM9', 'LPT1', 'LPT2', 'LPT3',
        'LPT4', 'LPT5', 'LPT6', 'LPT7', 'LPT8', 'LPT9']

    # Replace disallowed characters with underscores
    fname = re.sub(r'[\x00-\x1f<>:|?*"/\\]+', '_', title)

    # Append an underscore to the name portion if it's a reserved name
    base, ext = os.path.splitext(fname)
    if base in win32_special_names:
        fname = base + '_' + ext

    return fname

def make_search_slug(substr):
    """Lazy attempt at fuzzy matching preprocessing"""
    return re.sub(r"\W+", "", substr).lower()

def mappers_for(system, inlretro_dir):
    """Convenience wrapper for querying supported mappers

    (Hides the actual mechanism from code elsewhere for easy improvement)

    NOTE: This heuristic isn't *guaranteed* to keep working, but actually
    extracting the list of valid named would require parsing Lua and,
    at the moment, there's a 1-to-1 correlation between valid mapper names
    and the names of the scripts which implement them.
    """

    # system.lower() because Linux filesystems are case-sensitive and the
    # scripts/... folders are lowercase but proper capitalizations like "NES"
    # should be valid `system` values.
    # TODO: Redesign so the script can inventory its dependencies on start
    script_dir = os.path.join(inlretro_dir, 'scripts', system.lower())
    scripts = [os.path.splitext(fname)[0] for fname in os.listdir(script_dir)]
    scripts.sort()

    # Omit the template for writing new mapper scripts
    # It's neither a valid choice nor functional
    if "blank" in scripts:
        scripts.remove("blank")

    return scripts

def prompt_for(prompt, convert=lambda x: x, test=lambda x: x, err_msg=None):
    """Convenience helper for validity-checking input prompts

    After leading/trailing whitespace has been stripped, `convert(string)`
    will be called to produce the value that will be returned.
    It may raise TypeError or ValueError to trigger redisplay of the prompt.

    `test(converted)` must return a truthy value for the prompt loop to exit
    but, otherwise, has no effect on what gets returned.
    """
    prompt += " "
    choice, converted = None, None
    while not choice:
        try:
            converted = convert(input('\n' + prompt).strip())
        except (TypeError, ValueError):
            pass

        if converted is not None and test(converted):
            return converted
        elif err_msg:
            print(err_msg)

def prompt_yn(prompt):
    """DRY wrapper for asking yes/no questions

    Adds ` (y/n)` for you.
    """
    return prompt_for(prompt + ' (y/n)',
                      lambda x: x.lower(), lambda x: x[:1] in 'yn')[0] == 'y'

# -- System Definitions --

class System(object):
    """Unified definition of behaviour for a dumpable system"""
    name = None       # Human-friendly name for the system
    system_id = None  # The name used for the scripts/ folder and -c argument

    # The extensions for the ROM and save file, without leading period
    rom_ext = None
    save_ext = None  # Optional. Leave on `None` to omit `-a` flag to inlretro

    # Set these to false to skip prompting
    auto_mapper = False
    auto_size = False

    # Set this to enable semi-automatic lookup of required metadata
    info_db_path = None
    info_source = None
    unhashed_header_len = 0  # Length of header to skip before hashing ROM

    def __init__(self, **kwargs):
        """Shorthand for setting member variables"""

        for key, value in kwargs.items():
            setattr(self, key, value)

        if not self.rom_ext:
            self.rom_ext = self.system_id

        self.meta_db = None

    def __str__(self):
        return self.name

    def parse_database(self):  # pylint: disable=no-self-use
        """Override this to automate retrieval of mapper/size info"""
        return None

    def get_size(self, meta_entry=None):  # pylint: disable=R0201,W0613
        """Override to customize prompting for size.

        Return a set of arguments to be appended to inlretro's command line.
        """
        # TODO: Research more stringent validity checks for non-NES sizes
        return ['-k', prompt_for("Size of ROM to be dumped (in KiB):",
                                 int, lambda x: x > 0,
                                 "Value must be a positive integer.")]

    def get_meta(self):
        """Prompt for the name and attempt to look up metadata from it"""
        if self.info_db_path and not self.meta_db:
            if os.path.exists(self.info_db_path):
                self.meta_db = self.parse_database()
            else:
                print("\nCannot find %r to automatically query metadata" %
                      os.path.normpath(self.info_db_path))
                if self.info_source:
                    print("You may download it from %s\n" % self.info_source)

        if not self.meta_db:
            return {'name': prompt_for("Name of cartridge to be dumped: ",
                                       err_msg="Please type a name")}

        while True:
            slug = prompt_for("Please type part of the cartridge's name:",
                   make_search_slug, lambda x: x, "Please type a name")

            matches = [x for x in self.meta_db
                       if slug in make_search_slug(x['name']) or
                      slug == make_search_slug(x.get('catalog', ''))]

            if matches:
                def formatter(x):
                    """Callback for formatting col_menu entries"""
                    if 'catalog' in x:
                        return "{} ({})".format(x['name'], x['catalog'])
                    else:
                        return x['name']

                matches.append({'name': '← Revise Search', 'go_back': True})

                result = col_menu("Please select your cartridge:", matches,
                                 formatter)

                if not result.get('go_back'):
                    return result
            else:
                # TODO: Support falling back to all-manual details
                print("No matches found")

    def choose_mapper(self, meta, inlretro_dir):
        """Split out the logic for resolving any ambiguity in mapper choice"""
        # If we get a single mapper from meta, use it.
        # Otherwise, if not doing auto-detection, fallback to the list provided
        # by the database and then to the list of supported mappers
        mappers = meta.get('mappers', [])
        if mappers and len(mappers) == 1:
            return ['-m', meta['mappers'][0]]
        elif not self.auto_mapper:
            if mappers and len(mappers) > 1:
                mapper_id = meta.get('mapper_id')
                if mapper_id:
                    msg = "Select script to use for mapper %r:" % mapper_id
                else:
                    msg = "Select script to use for this mapper:"

                return ['-m', col_menu(msg, mappers)]
            else:
                return ['-m', col_menu("Select mapper used by this cart:",
                    mappers_for(self.system_id, inlretro_dir))]
        return []

    def dump_game(self, path_base, inlretro_path):
        """Dump a cartridge for the specified system.

        `path_base` should be the path the ROM and, if applicable, the save
        file should be written to, minus extension.
        """

        args = []
        meta = self.get_meta()
        inlretro_dir, inlretro_cmd = os.path.split(inlretro_path)

        if meta.get('size'):
            args += ['-k', str(math.ceil(meta['size'] / 1024))]
        elif not self.auto_size:
            args += self.get_size(meta)

        args += self.choose_mapper(meta, inlretro_dir)

        rom_path = os.path.join(path_base,
            make_filename('{}.{}'.format(meta['name'], self.rom_ext)))

        cmd = [os.path.join(os.curdir, inlretro_cmd),
               "-c", self.system_id,
               "-s", os.path.join('scripts', 'inlretro2.lua'),
               "-d", rom_path,
               ] + list(args)

        if self.save_ext:
            cmd.extend(["-a", make_filename('{}{}{}.{}'.format(path_base,
                os.sep, meta['name'], self.save_ext))])

        while True:
            try:
                print("\nRunning: " + ' '.join(shlex.quote(x) for x in cmd))
                print("With working directory: {}\n".format(inlretro_dir))
                subprocess.check_call(cmd, cwd=inlretro_dir)
            except subprocess.CalledProcessError:
                if prompt_yn("inlretro returned an error. Try again?"):
                    continue
                else:
                    return
            else:
                if not os.path.exists(rom_path):
                    if prompt_yn("No ROM file was produced. Try again?"):
                        continue
                    else:
                        return

            # TODO: Have an option to dump the save file multiple times and
            #       compare to catch corruption introduced in the dumping
            #       process.

            retry = None
            for hash_type in ('sha1', 'md5'):
                if hash_type in meta and meta[hash_type]:
                    hash_got = self.hash_rom(rom_path, hash_type).lower()
                    hash_expected = meta[hash_type]
                    if isinstance(hash_expected, str):
                        hash_expected = [hash_expected]

                    if any(hash_got == x.lower() for x in hash_expected):
                        print("Success! Dump matches expected hash.")
                        return
                    else:
                        print("ROM doesn't match %s hash on file. (%s != %s)"
                              "" % (hash_type, hash_got, meta[hash_type]))
                        retry = prompt_yn("Try again?")
                        break

            if retry is True:
                continue
            elif retry is False:
                return

            print("No good hash on file. Cannot check dump success.")
            return

    def hash_rom(self, path, hash_type):
        return hash_file(path,
                         hasher=getattr(hashlib, hash_type),
                         seek=self.unhashed_header_len)

class GameboySystem(System):
    """Definition for dumping Gameboy cartridges"""
    name = 'Gameboy'
    system_id = 'gb'

    info_db_path = os.path.join(basedir, os.pardir, 'docs',
        'CartridgeList.csv')
    info_source = (
        "https://github.com/gbdev/awesome-gbdev/blob/master/CartridgeList.csv")
    # TODO: Look into processing the data from
    #       https://gbhwdb.gekkio.fi/cartridges/ into a datfile with mapper
    #       information AND nice titles
    #       (See https://github.com/Gekkio/gb-hardware-db)

    MAPPERS = {
        'MBC1': ['mbc1'],
        'ROM': ['romonly', 'romonly_paul'],
    }

    def parse_database(self):
        meta_db = []

        with open(self.info_db_path) as fobj:
            for line in csv.reader(fobj.readlines()):
                try:
                    mapper_id = line[0]
                    meta_db.append({
                        'mapper_id': mapper_id,
                        'mappers': self.MAPPERS.get(mapper_id, None),
                        'size': int(line[1]) * 1024,
                        'name': line[5],
                    })
                except (TypeError, ValueError):
                    pass  # TODO: Report failure somehow

        return meta_db


class GbaSystem(System):
    """Definition for dumping Gameboy Advance cartridges"""
    name = 'Gameboy Advance'
    system_id = 'gba'
    auto_mapper = True

    # TODO: Try to find a source with Nintendo catalogue numbers
    info_source = "https://datomatic.no-intro.org/?page=download"
    info_db_path = os.path.join(basedir, os.pardir, 'docs',
        'Nintendo - Game Boy Advance (20190815-231257).dat')
    # FIXME: Don't depend on a specific version's filename

    def parse_database(self):
        meta_db = []
        old_meta = None
        tree = ET.parse(self.info_db_path)

        for game in tree.iterfind('game'):
            for rom in game.iterfind('.//rom'):
                meta = {
                    'name': game.get('name'),
                    'md5': rom.get('md5'),
                    'sha1': rom.get('sha1'),
                    'size': int(rom.get('size')),
                }

                # Ensure no duplicate entries show up
                if old_meta != meta:
                    meta_db.append(meta)

                old_meta = meta

        return meta_db


class N64System(System):
    """Definition for dumping Nintendo 64 cartridges"""
    name = 'Nintendo 64'
    system_id = 'n64'
    auto_mapper = True

    info_db_path = os.path.join(basedir, os.pardir, 'docs',
        'Nintendo - Nintendo 64 - Dump Status (BigEndian) (2019-07-24).csv')

    def parse_database(self):
        meta_db = []

        with open(self.info_db_path) as fobj:
            for line in csv.reader(fobj.readlines()[2:], delimiter=';'):
                if len(line) >= 3:
                    try:
                        meta_db.append({
                            'name': line[1],
                            'size': int(line[2]),
                            'md5': line[4],
                            'catalog': line[5],
                            'mappers': None,
                        })
                    except (TypeError, ValueError):
                        pass  # TODO: Report failure somehow

        return meta_db


class NesSystem(System):
    """Definition for dumping NES and Famicom cartridges"""
    name = 'NES/Famicom'
    system_id = 'nes'

    # FIXME: Don't depend on a specific version's filename
    info_db_path = os.path.join(basedir, os.pardir, 'docs',
        'NesCarts (2017-08-21).xml')
    info_source = 'http://bootgod.dyndns.org:7777/xml.php'

    # Thanks to https://forums.nesdev.com/viewtopic.php?f=2&t=9425
    # for pointing out why the hashes weren't matching
    unhashed_header_len = 16

    MAPPERS = {
        0: ['nrom'],  # Based on a "Popeye no Eigo Asobi" comment in nrom.lua
        1: ['mmc1'],
        # unrom_tsop is excluded from the results for mapper 2 because I don't
        # want to inconvenience and potentially confuse 99.999%+ use cases
        # with the off-chance that somebody is trying to dump an unlicensed
        # clone of a legit cart and would benefit from datfile-based detection
        2: ['unrom'],
        3: ['cnrom'],
        4: ['mmc3'],
        5: ['mmc5'],
        9: ['mmc9'],
        10: ['mmc4'],
        11: ['cdream'],
        28: ['action53', 'action53_tsop'],
        30: ['mapper30', 'mapper30v2'],  # TODO: Check that this is what I want
        # I *think* easyNSF is #31: wiki.nesdev.com/w/index.php/INES_Mapper_031
        34: ['bnrom'],
        69: ['fme7'],
        111: ['gtrom'],
        # TODO: Figure out what the cninja and dualport mappers are for.
    }

    def parse_database(self):
        meta_db = []

        tree = ET.parse(self.info_db_path)
        old_meta = None
        for game in tree.iterfind('game'):
            seen_hashes = {
                'crc': [],
                'sha1': [],
            }
            for cartridge in game.iterfind('cartridge'):
                for hash_type in ('crc', 'sha1'):
                    hash_str = cartridge.get(hash_type, None)
                    if hash_str and hash_str not in seen_hashes[hash_type]:
                        seen_hashes[hash_type].append(hash_str)

                for board in game.iterfind('.//board'):
                    meta = {
                        'name': game.get('name'),
                        'catalog': game.get('catalog'),
                        'mappers': None,
                        'prg': 0,
                        'chr': 0,
                        'wram': None,
                    }

                    meta['mapper_id'] = board.get('mapper')
                    try:
                        m_type = int(meta['mapper_id'])
                    except (TypeError, ValueError):
                        m_type = None

                    if m_type in self.MAPPERS:
                        meta['mappers'] = self.MAPPERS[m_type]

                    for ctype in ('prg', 'chr', 'wram'):
                        for rom in board.iterfind('.//' + ctype):
                            rom_size = rom.get('size', '')
                            assert re.match(r'\d+k', rom_size)
                            meta[ctype] = int(rom_size[:-1])

                    # Skip duplicate entries show up
                    if old_meta != meta:
                        tmp_meta = {}
                        tmp_meta.update(meta)
                        tmp_meta.update(seen_hashes)
                        meta_db.append(tmp_meta)
                    old_meta = meta

        return meta_db

    def get_size(self, meta_entry=None):
        args, meta_entry = [], meta_entry or {}

        for key, arg in (('prg', '-x'), ('chr', '-y')):
            value = meta_entry.get(key)
            if not value:
                value = prompt_for(
                    "Enter %s ROM size in KiB (0 for none):" % key.upper(),
                    int, lambda x: x == 0 or (x >= 8 and x % 4 == 0),
                    "ROM size must be at least 8 and a multiple of 4.")
            if value:
                args.extend([arg, str(value)])

        if 'wram' in meta_entry:
            has_save = meta_entry.get('wram')
        else:
            has_save = prompt_yn(
                "Does this game have battery-backed saves (WRAM)?")

        if has_save:
            args.extend(['-w', '8'])
            self.save_ext = 'sav'
        else:
            self.save_ext = None

        return args


SYSTEMS = [
    NesSystem(),
    System(name='SNES/Super Famicom',
        system_id='snes', rom_ext='sfc', save_ext='srm',
        auto_mapper=True, auto_size=True),
    N64System(),
    GameboySystem(),
    GbaSystem(),
    System(name='Sega Genesis/Mega Drive', system_id='genesis', rom_ext='bin',
           auto_mapper=True, auto_size=True),
]

# -- Code Here --

def main():
    """The main entry point, compatible with setuptools entry points."""
    default_inlretro_path = os.path.join(basedir, 'inlretro')
    if os.name == 'nt':
        default_inlretro_path += '.exe'

    from argparse import ArgumentParser, RawDescriptionHelpFormatter
    parser = ArgumentParser(formatter_class=RawDescriptionHelpFormatter,
            description=__doc__.replace('\r\n', '\n').split('\n--snip--\n')[0])
    parser.add_argument('--version', action='version',
            version="%%(prog)s v%s" % __version__)
    parser.add_argument('--inlretro-path', action='store', metavar="PATH",
            default=os.path.relpath(default_inlretro_path),
            help="Path to the `inlretro` binary (default: %(default)s)")

    parser.add_argument('output_dir', nargs='?',
            default=os.path.relpath(os.path.join(basedir, 'ignore')),
            help="Directory to write dumps to (default: %(default)s)")

    args = parser.parse_args()

    if not os.path.exists(args.inlretro_path):
        print("Could not find %r. Exiting." % args.inlretro_path)
        sys.exit(1)

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    while True:
        print(__appname__)
        system = col_menu("Select cartridge type:", SYSTEMS)
        system.dump_game(args.output_dir, os.path.abspath(args.inlretro_path))

        if not prompt_yn("Dump another?"):
            break


if __name__ == '__main__':
    main()

# vim: set sw=4 sts=4 expandtab :