path: root/vgstash

#!/usr/bin/env python3

# TODO: Consider putting all help messages into a dict for easier management
# TODO: Decide on docstring delimiter and stick to it
# TODO: refactor

import sys
import os
import sqlite3
import argparse
import yaml
import subprocess # to handle the piping use case

__version__ = '0.2'

DB_LOCATION = ''
OWNERSHIP = 1
PROGRESS = 1
TABLE_WIDTH = 80

def safe_print(line):
    try:
        print(line, flush=True)
    # We're catching this in case the other end of a pipe exits early
    except BrokenPipeError:
        sys.stderr.close()

def set_env():
    '''Ensures environment variables are respected. Sets defaults if
    they're not present. If the defaults are not usable, it throws an
    AssertionError and alerts the user.'''
    # This makes outside-scope referencing and assignment possible
    global DB_LOCATION, OWNERSHIP, PROGRESS, TABLE_WIDTH

    # Precedence = $VGSTASH_DB_LOCATION, $HOME/.vgstash.db, ./.vgstash.db
    DB_LOCATION = os.getenv('VGSTASH_DB_LOCATION', os.path.join(os.getenv('HOME', '.'), '.vgstash.db'))

    # Can't decide what to do with this yet; the first run of vgstash doesn't let you init the db if we uncomment this
    # try:
    #     assert(os.path.isfile(DB_LOCATION) and os.path.exists(DB_LOCATION))
    # except AssertionError:
    #     print("VGSTASH_DB_LOCATION is not a file. Unset it to fall back to $HOME/.vgstash.db, or correct the environment variable.")
    #     sys.exit()

    OWNERSHIP = int(os.getenv('VGSTASH_DEFAULT_OWNERSHIP', OWNERSHIP))
    try:
        assert(OWNERSHIP == 0 or OWNERSHIP == 1)
    except AssertionError:
        print("VGSTASH_DEFAULT_OWNERSHIP is not zero or one. Ensure your variable is set correctly.")
        sys.exit()

    PROGRESS = int(os.getenv('VGSTASH_DEFAULT_PROGRESS', PROGRESS))
    try:
        assert(PROGRESS >= 0 and PROGRESS <= 3)
    except AssertionError:
        print("VGSTASH_DEFAULT_PROGRESS is not between zero and three. Ensure your variable is set correctly.")
        sys.exit()

    TABLE_WIDTH = int(os.getenv('VGSTASH_TABLE_WIDTH', TABLE_WIDTH))
    try:
        assert(TABLE_WIDTH >= 50 or TABLE_WIDTH == 0)
    except AssertionError:
        print("VGSTASH_TABLE_WIDTH must be a number that's zero, or a number that's 50 or greater, to facilitate readability. A value of zero will use your terminal's width.")
        sys.exit()

def init_db(args):
    '''Creates the database schema and the relevant views, preparing it
    for use with vgstash. Will not overwrite an extant database.'''

    # Confirm we can create the file
    try:
        assert(os.path.isfile(DB_LOCATION) == False)
        conn = sqlite3.connect(DB_LOCATION)
        print("DB successfully connected!")
    except sqlite3.OperationalError:
        print("The database could not be created and/or connected to.")
        sys.exit()
    except AssertionError:
        print("The database already exists! Delete or move '%s' and init the database again." % DB_LOCATION)
        sys.exit()

    # Now let's run a bunch of queries! Fun...
    try:
        conn.execute("CREATE TABLE games (\
        title TEXT NOT NULL,\
        system TEXT NOT NULL,\
        ownership INTEGER NOT NULL DEFAULT 1,\
        progress INTEGER NOT NULL DEFAULT 1\
        );")
        print("Table created.")
        # TODO: Consider executescript()
        conn.execute("CREATE VIEW allgames AS SELECT rowid,* FROM games ORDER BY system,title ASC;")
        conn.execute("CREATE VIEW arcade AS SELECT rowid,* FROM games WHERE ownership >= 1 AND progress = -1 ORDER BY system,title ASC;")
        conn.execute("CREATE VIEW backlog AS SELECT rowid,* FROM games WHERE ownership >= 1 AND progress < 2 ORDER BY system,title ASC;")
        conn.execute("CREATE VIEW borrowing AS SELECT rowid,* FROM games WHERE ownership = 0 AND progress = 1 ORDER BY system,title ASC;")
        conn.execute("CREATE VIEW complete AS SELECT rowid,* FROM games WHERE progress = 3 ORDER BY system,title ASC;")
        conn.execute("CREATE VIEW digital AS SELECT rowid,* FROM games WHERE ownership >= 2 ORDER BY system,title ASC;")
        conn.execute("CREATE VIEW done AS SELECT rowid,* FROM games WHERE progress >= 2 ORDER BY system,title ASC;")
        conn.execute("CREATE VIEW incomplete AS SELECT rowid,* FROM games WHERE ownership >= 1 AND progress = 2 ORDER BY system,title ASC;")
        conn.execute("CREATE VIEW new AS SELECT rowid,* FROM games WHERE progress = 0 ORDER BY system,title ASC;")
        conn.execute("CREATE VIEW owned AS SELECT rowid,* FROM games WHERE ownership >= 1 ORDER BY system,title ASC;")
        conn.execute("CREATE VIEW physical AS SELECT rowid,* FROM games WHERE ownership = 1 OR ownership = 3 ORDER BY system,title ASC;")
        conn.execute("CREATE VIEW playlog AS SELECT rowid,* FROM games WHERE ownership >= 1 AND progress = 1 ORDER BY system,title ASC;")
        conn.execute("CREATE VIEW unowned AS SELECT rowid,* FROM games WHERE ownership = 0 ORDER BY system,title ASC;")
        conn.execute("CREATE VIEW wishlist AS SELECT rowid,* FROM games WHERE ownership = 0 AND progress = 0 ORDER BY system,title ASC;")
        print("Views created.")
        conn.commit()
    except sqlite3.OperationalError as e:
        print("Query error:", e)
    finally:
        conn.close()

def export_db(args):
    '''Exports a YAML representation of the database to a given file.'''
    try:
        if args.path == None:
            fp = sys.stdout
        else:
            fp = open(args.path, 'w')
    except PermissionError:
        print("The path is not writable by the current user. Please double-check permissions.")
        sys.exit()
    else:
        conn = sqlite3.connect(DB_LOCATION)
        conn.row_factory = sqlite3.Row
        c = conn.cursor()
        fp.write("# vgstash DB file version 0.2\n")
        db_tree = []
        # The .format() call supports the optional -i flag
        for row in c.execute("SELECT {}title,system,ownership,progress FROM games".format("rowid," if args.ids else '')):
            game = {}
            for field in row.keys():
                game.update({field: row[field]})
            db_tree.append(game)
        yaml.dump(db_tree, fp, default_flow_style=False, indent=4, allow_unicode=True)
        fp.close()
        if fp is not sys.stdout:
            print("Database exported to {}.".format(args.path))

def import_db(args):
    '''Read yaml lines from a file (or stdin) to add to the database.'''
    try:
        if args.path == None:
            fp = sys.stdin
        else:
            fp = open(args.path, 'r')
    except PermissionError:
        print("The path is not readable by the current user. Please check permissions.")
        sys.exit()
    # Time for business! :D
    records = 0
    data = yaml.safe_load(fp)
    conn = sqlite3.connect(DB_LOCATION)
    conn.row_factory = sqlite3.Row
    c = conn.cursor()
    # List comprehensions are why I like Python.
    # inline conditionals are badass, too
    c.executemany("INSERT {}INTO games ({}title,system,ownership,progress) VALUES ({}:title, :system, :ownership, :progress)".format("OR REPLACE " if args.ids else "", "rowid," if args.ids else "", ":rowid," if args.ids else ""), (game for game in data))
    conn.commit()
    conn.close()
    print("Imported {} games.".format(c.rowcount))

def add_game(args):
    '''Adds a game to the database. Requires at least a title and
    system.

    Ownership can be 0 to 3:(not owned) or 1 (owned). The default is 1.

    Completion can be -1 to 3:
       -1 (unbeatable)
        0 (fresh, never played)
        1 (in-progress) (default)
        2 (beaten)
        3 (completed 100%)
    '''
    if args.ownership == '-':
        args.ownership = OWNERSHIP
    if args.progress == '-':
        args.progress = PROGRESS
    # Translate our args so they can be added and reflected correctly
    args.ownership = translate_ownership(args.ownership)
    args.progress = translate_progress(args.progress)

    conn = sqlite3.connect(DB_LOCATION)
    game = (args.title, args.system, args.ownership, args.progress)
    conn.execute("INSERT INTO games VALUES(:title, :system, :ownership, :progress)", game)
    conn.commit()
    conn.close()
    qual = (
        "don't own it",
        "own it physically",
        "own it digitally",
        "own it physically and digitally"
    )
    comp = {
        -1: "it's unbeatable",
        0: "it's new",
        1: "you're playing it",
        2: "it's been beaten",
        3: "it's been completed"
    }
    print("Added {0} for {1}. You {2} and {3}.".format(args.title, args.system, qual[int(args.ownership)], comp[int(args.progress)]))

def delete_game(args):
    '''Removes a game from the database.'''
    target = game_found(args.id)
    if target != None:
        try:
            conn = sqlite3.connect(DB_LOCATION)
            c = conn.cursor()
            conn.execute("DELETE FROM games WHERE rowid = :id", {'id': args.id})
            conn.commit()
            print("Removed {0} for {1} from your database.".format(target[0], target[1]))
        except sqlite3.OperationalError:
            print("Could not remove game from the database. Check the file and ensure you have write access.")
        finally:
            conn.close()
    else:
        print("That game ID does not exist in the database.")

def translate_progress(arg):
    '''Translate a letter progress value into a numeric value for the database.'''
    try:
        if len(arg) != 1:
            print("Argument must be a single character.")
            sys.exit()
        if not arg.isnumeric():
            vals = {
                'u': -1,
                'n': 0,
                'p': 1,
                'b': 2,
                'c': 3,
            }
            try:
                return vals[arg]
            except KeyError:
                # Note to self, doubling a brace escapes it. This was tucked away in the Python docs...
                print("Value '{}' not valid. Try one of {{{}}}.".format(arg, ', '.join(vals)))
                sys.exit()
        else:
            return arg
    except TypeError:
        return arg

def translate_ownership(arg):
    '''Translate a letter ownership value into a numeric value for the database.'''
    try:
        if len(arg) != 1:
            print("Argument must be a single character.")
            sys.exit()
        if not arg.isnumeric():
            vals = {
                'n': 0,
                'p': 1,
                'd': 2,
                'b': 3
            }
            try:
                return vals[arg]
            except KeyError:
                print("Value '{}' not valid. Try one of {{{}}}.".format(arg, ', '.join(vals)))
                sys.exit()
        else:
            return arg
    except TypeError:
        return arg

def update_game(args):
    '''Changes a specific field for a game. Mostly used for correcting
    mistakes and updating game status as it changes.'''
    target = game_found(args.gid)
    # Let's accept anything remotely matching the field names
    opts = ['system', 'title', 'ownership', 'progress']
    for i in opts:
        if i.startswith(args.field):
            args.field = i
    if not (args.field in opts):
        print("Invalid field name indicated! Please choose from {}".format(', '.join(opts)))
        sys.exit()
    # Translate y, n, f, i, b, and c
    if args.field == 'ownership':
        args.value = translate_ownership(args.value)
    if args.field == 'progress':
        args.value = translate_progress(args.value)
    # We need this workaround because execute() doesn't like variable column names
    update_stmt = "UPDATE games SET {} = :val WHERE rowid = :id".format(args.field)
    conn = sqlite3.connect(DB_LOCATION)
    c = conn.cursor()
    c.execute(update_stmt, {'id': args.gid, 'val': args.value})
    conn.commit()
    if c.rowcount == 1:
        own_msg = ('not owned', 'physically owned', 'digitally owned', 'physically and digitally owned')
        prog_msg = {
            -1: 'unbeatable',
            0: 'new',
            1: 'playing',
            2: 'beaten',
            3: 'complete'
        }
        update_msg = {
            'title': "{ot} on {os} is now named {val}.",
            'system': "{ot} on {os} is now on {val}.",
            'ownership': "{ot} on {os} is now marked {val}.",
            'progress': "{ot} on {os} is now marked {val}."
        }
        print(update_msg[args.field].format(ot = target['title'], os = target['system'], val = args.value if args.field == 'title' or args.field == 'system' else own_msg[int(args.value)] if args.field == 'ownership' else prog_msg[int(args.value)]))
    else:
        print("Could not update game information. Check the DB's permissions.")
    conn.close()

def list_games(args):
    '''Filters games according to preset queries. Internally,
    they are SQLite VIEWs

    Each filter targets different games:

        all:           everything, sorted by system then title
        backlog:     * have not been beaten or completed
        complete:      have been completed
        done:          are beaten or complete
        fresh:         haven't been started yet
        incomplete:  * beaten, but not completed
        owned:         are marked as 'owned'
        unowned:       are marked as 'unowned'
        wishlist:      are both 'unowned' and 'fresh' or 'in-progress'

    * Will only target games that are owned.
    '''
    if args.filter == 'all':
        args.filter = 'allgames'
    # DBs don't allow table names to be parameterized. This little hack
    # works around that limitation.
    conn = sqlite3.connect(DB_LOCATION)
    conn.row_factory = sqlite3.Row
    select_stmts = {
            'allgames':   "SELECT * FROM allgames",
            'arcade':     "SELECT * FROM arcade",
            'backlog':    "SELECT * FROM backlog",
            'borrowing':  "SELECT * FROM borrowing",
            'complete':   "SELECT * FROM complete",
            'digital':    "SELECT * FROM digital",
            'done':       "SELECT * FROM done",
            'incomplete': "SELECT * FROM incomplete",
            'new':        "SELECT * FROM new",
            'owned':      "SELECT * FROM owned",
            'physical':   "SELECT * FROM physical",
            'playlog':    "SELECT * FROM playlog",
            'unowned':    "SELECT * FROM unowned",
            'wishlist':   "SELECT * FROM wishlist",
            }
    # We're emulating a do-while loop
    first_pass = True
    for row in conn.execute(select_stmts[args.filter]):
        if args.raw:
            # Use this for raw output, for use with other tools
            print('|'.join(map(str,row)))
            continue
        else:
            row_format(row, header=first_pass)
            first_pass = False


def game_found(gid):
    '''Confirms the given ID is in the database and, if found, returns it.
    Used in conjunction with delete_game() and update_game().'''
    conn = sqlite3.connect(DB_LOCATION)
    # Give us access to .keys()
    conn.row_factory = sqlite3.Row
    c = conn.cursor()
    c.execute("SELECT title,system,ownership,progress FROM games WHERE rowid = :id", {'id': gid})
    row = c.fetchone()
    conn.close()
    return row

def row_format(args, header):
    # There's another way to do this, involving gathering the entire results and *then* formatting them
    # That is incredibly wasteful of resources imo, so it may need some testing to see if it's better.
    # Ideally, we'd only make the table as wide as needed; that can't happen unless we know the longest
    # title's length...
    # Get our maximum width so we can toy around with things.
    if TABLE_WIDTH >= 50:
        maxwidth = TABLE_WIDTH
    else:
        try:
            maxwidth = os.get_terminal_size().columns
        except OSError:
            with open('/dev/tty') as tty:
                curwidth = int(subprocess.check_output(['stty', 'size']).split()[1])
            maxwidth = curwidth
    twidth = maxwidth - 37
    if header == True:
        print("{:^4s} | {:<{w}s} | {:<8s} | {:^3s} | {:<7s}".format("ID", "Title", "System", "Own", "Status", w=twidth))
        print("-" * maxwidth)

    gidstr = "{: >4d}".format(args['rowid'])
    titlestr = "{: <{w}s}".format(args['title'][:twidth], w=twidth)
    systemstr = "{: ^8s}".format(args['system'][:8])
    ownltr = [' ', 'P', '  D', 'P D']
    ownstr = "{: <3s}".format(ownltr[args['ownership']])
    statltr = {
        -1: 'U',
        0: 'N',
        1: 'P',
        2: 'B',
        3: 'C'
    }
    statstr = "{: <7s}".format((" " * args['progress'] * 2) + statltr[args['progress']])
    """
      ID | Title           | System   | Own | Status
    ---------------------------------------------------
    1234 | This is a title | Wii U VC |  *  | U N P B C
    """
    safe_print(" | ".join((gidstr, titlestr, systemstr, ownstr, statstr)))


def main():
    # Establish our important variables
    set_env()

    # Set up the command parsers
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers()

    # 'add' command
    parser_add = subparsers.add_parser('add', help="Add a game to your database")
    parser_add.add_argument('title', type=str, help="The title of the game you're adding")
    parser_add.add_argument('system', type=str, help="The system your game is for")
    parser_add.add_argument('ownership', type=str, nargs='?', default='-', help="0 or 'n' for no, 1 or 'y' for yes. Defaults to {0}. Accepts '-' for default, in case you want to set progress but not ownership".format(OWNERSHIP))
    parser_add.add_argument('progress', type=str, nargs='?', default='-', help="Indicates completion level. 0 = [f]resh, never played; 1 = [i]n progress; 2 = [b]eaten; 3 = [c]ompleted. Defaults to {0}".format(PROGRESS))
    parser_add.set_defaults(func=add_game)

    # 'init' command
    parser_init = subparsers.add_parser('init', help="Initialize the database so it can be used")
    parser_init.set_defaults(func=init_db)

    # 'delete' command
    parser_del = subparsers.add_parser('delete', help="Remove a game from your database")
    parser_del.add_argument('id', type=int, help="The ID of the game to be deleted")
    parser_del.set_defaults(func=delete_game)

    # 'list' command
    parser_list = subparsers.add_parser('list', help="List your games with preset views")
    list_filters = ['all', 'arcade', 'backlog', 'borrowing', 'complete', 'digital', 'done', 'incomplete', 'new', 'owned', 'physical', 'unowned', 'wishlist', 'playlog']
    parser_list.add_argument('filter', nargs='?', choices=list_filters, default='all', help='Filter games accerding to preset queries. Valid filters: {}'.format(', '.join(list_filters)))
    parser_list.add_argument('-r', '--raw', action="store_true", help="Output the list in a machine-readable format, separated by pipe characters.")
    parser_list.set_defaults(func=list_games)

    # 'update' command
    parser_up = subparsers.add_parser('update', help="Update a specific game")
    parser_up.add_argument('gid', type=int, help="The game's ID (found in 'vgstash list' output)")
    parser_up.add_argument('field', help="The field you wish to update")
    parser_up.add_argument('value', type=str, help="The new value for the field")
    parser_up.set_defaults(func=update_game)

    # 'export' command
    parser_ex = subparsers.add_parser('export', help="Export the database's information to a YAML file.")
    parser_ex.add_argument('path', nargs='?', default=None, help="The full path of the file you want to write to. Defaults to stdout")
    parser_ex.add_argument('-i', '--ids', action="store_true", help="Export game IDs along with game information. Useful if preparing a batch update.")
    parser_ex.set_defaults(func=export_db)

    # 'import' command
    parser_im = subparsers.add_parser('import', help="Import a YAML file's content into the database.")
    parser_im.add_argument('path', nargs='?', default=None, help="The full path of the file you want to read from. Defaults to stdin")
    parser_im.add_argument('-i', '--ids', action="store_true", help="Update via game IDs in the database instead of appending")
    parser_im.set_defaults(func=import_db)

    # Let'er rip!
    args = parser.parse_args()
    try:
        args.func(args)
    except AttributeError:
        # Handle "no arguments" case
        args = ['-h']
        parser.parse_args(args)
        args.func(args)
    # Oddity in Python needed if you're going to play nice with piping
    sys.stderr.close()

# Our usual incantation
if __name__ == "__main__":
    main()