#!/usr/bin/env python

from __future__ import print_function
from __future__ import absolute_import
from __future__ import division

import os
import re
import argparse
import sys
from sys import exc_info

try:
    # Python 3 only
    from inspect import getfullargspec as getargspec
except ImportError:
    # Python 2 and Python 3
    from inspect import getargspec

try:
    # Python 2
    from urlparse import urlparse
except ImportError:
    # Python 3
    from urllib.parse import urlparse

from ROOT import PyConfig
PyConfig.IgnoreCommandLineOptions = True
from rat.ratdb import RATDBConnector, print_db_error

try:
    import pgdb as psql
except ImportError:
    print("::WARNING : Failed to import built-in PyGreSQL module.")
    print("            Atempting to load psycopg2, if it is available in the system")
    try:
        import psycopg2 as psql
    except ImportError:
        print("::ERROR : Couldn't find any suitable postgres python interface module. Aborting.")
        sys.exit(0)

if sys.version_info[0] == 2 and sys.version_info[1] < 7:
    print("This code cannot be run with Python older than 2.7")
    sys.exit(0)

# This ugliness is required because python 3 removed basestring.
# However, if string in python 2 is unicode, it is not of type str.
try:
    # Python 2
    is_string = basestring
except NameError:
    # Python 3
    is_string = str

def command_help(arg):
    '''
    Get help for a given command.
    '''
    if arg in commands.keys():
        print(commands[arg].__doc__)
    else:
        raise ValueError('Unknown command [{0}]'.format(arg))

#######################################
###
### Standard code
###
#######################################

def get_run_type(db, args, pass_num, **kwargs):
    """
    Print the run type word for a provided list of runs
        * Mandatory arguments:

          run/run range: any number of run numbers or run ranges defined in the form run_start-run_end

        * Optional arguments:

          pass : print results for requested pass
          debug: print results for all passes available
          tag  : RATDB tag to use

        Default behavior: print result for highest pass available for each run

    Example uses:
      1. Print the run type of the latest pass for a series of runs
        ratdb <auth> get_run_type 14226 14230-14240
      2. Print the run type for all passes of a run
        ratdb <auth> -d get_run_type 14226
      3. Print the run type for pass 1 of a run
        ratdb <auth> -p 1 get_run_type 14226
    """

    # Required to accept arguments that may not be used.
    del kwargs

    result = db.get_run_type(runs=args, pass_num=pass_num)
    # Now parse the result to show
    print("=" * 80)
    for run, contents in sorted(result.items()):
        if db.debug:
            if contents is not None:
                print("RUN : [{0}]".format(run))
                print('-' * 80)
                for np, rt in contents.items():
                    print("       pass {0} : {1} [{2}]".format(np, rt, format(rt, '#08x')))
            else:
                print("RUN : [{0}] : Not found".format(run))
        else:
            if contents is not None:
                print("RUN : [{0}] : {1} [{2}]".format(run, contents, format(contents, '#08x')))
            else:
                print("RUN : [{0}] : Not found".format(run))
        print('=' * 80)

def list_obj_run(db, obj_type, args, index, **kwargs):
    """
    List all available objects of a given type and index that are valid for
    the given run(s).

    Usage: ratdb [options] list_obj_run type[index] <run1> [run2 run3-runN]

    """

    # Required to accept arguments that may not be used.
    del kwargs

    assert len(args) > 0, "list_obj_run : At least one run must be specified"

    result = db.list_objects_type_run(obj_type=obj_type, runs=args, index=index)

    print("=" * 80)
    for run, content in sorted(result.items()):
        if content is None:
            print("RUN {0} : Not found".format(run))
        else:
            for header in content:
                print("{tp:30s}[{idx:s}] range [{rb:10d},{re:10d}] version {vs:2d} pass {ps:3d}   stored {ts}".format(tp=header['type'], idx=header['index'],
                                                                                                                      rb=header['run_range'][0], re=header['run_range'][1], ps=header['pass'], ts=header['storage_time'], vs=header['version']))
    print("=" * 80)

def list_all_objs_run(db, args, **kwargs):
    """ List all the objects that exist in the database for a given run (or list of runs).

    Usage: ratdb [options] list_objects_run run1 [run2 run3 run4-runN]

    This command loops over all the listed runs, connects to the database
    and extracts a list of objects that are valid for each run, printing the respective details

    """

    # Required to accept arguments that may not be used.
    del kwargs

    assert len(args) > 0, "list_objs_run : At least one run is needed"

    result = db.list_all_objs_run(runs=args)

    print("=" * 80)

    for run, contents in sorted(result.items()):
        print("-" * 20)
        print("RUN {0}".format(run))
        print("-" * 20)
        if contents is None:
            print("NOT FOUND")
        else:
            for header in contents:
                print("{0:30s} range [{1:10d},{2:10d}] version {vs:3d} pass {3:3d}  stored {ts}".format("{0:s}[{1:s}]".format(header['type'], header['index']),
                                                                                                        header['run_range'][0], header['run_range'][1], header['pass'], ts=header['storage_time'], vs=header['version']))
    print("=" * 80)

def upload(db, args, **kwargs):
    """ Upload a series of tables from RATDB files.

    Usage: ratdb [connection options] upload <file1.ratdb> [f2.ratdb] ... [fN.ratdb]

    One can insert several objects in the same file, and also several different files.
    The objects are inserted in the database from top to bottom in the file and the
    files are walked by the order given. This can be important when uploading
    objects in truncate mode.

    The write mode (truncate, increment_pass) are decided based on the contents of the
    'pass' parameter in each object.

    pass : -1 --> Increment pass mode. If the object overlaps with another, it's pass number increments.
    pass : -2 --> Truncate the exising objects that are valid over the same range.

    Returns a list of dictionaries with details of the objects inserted into RATDB, i.e.

    [{obj1},{obj2},{obj3},{obj4},]

    where objN is

    {"type": str
     "index" str
     "run_range":list
     "pass":int,
     "key":int,
    }

    NOTE: This command is incompatible with the -t/--tag option. One cannot upload data into a tag.

    """

    # Required to accept arguments that may not be used.
    del kwargs

    assert len(args) > 0, "upload : At least one input file is necessary"
    result = db.upload(files=args)
    print("A total of {0} objects have been uploaded.".format(len(result)))
    if len(result) > 0:
        print('=' * 80)
        print("{0:30s}  {1:15s}  {2:4s}  {3:7s}    {4:5s}".format("type[index]", "run_range", "pass", "version", "OID"))
        print('-' * 80)
        for header in result:
            print("{0:30s}  [{1:10d},{2:10d}]  {3:4d}  {4:7d}  {5:5d}".format("{0:s}[{1:s}]".format(header['type'], header['index']),
                                                                              header['run_range'][0], header['run_range'][1], header['pass'], header['version'], header['key']))
    print('=' * 80)

def dump_table(db, obj_type, args, **kwargs):
    """ Dumps a table into a BZip2 output file.

    Usage: ratdb [options] dump_table type[index] <run1> [run2-runN] [run3] ...

    Optional flags:
      -p <pass> : Dump a specific pass. If pass is -1, dump the latest pass (default)
      -d        : Dump all passes available in the database. Ignores -p
      -r        : Dump the table in raw format (as it was uploaded). Use with caution.

    """

    #First thing is check that we have at least 2 arguments: a table name and a run
    assert obj_type is not None, "dump_table : type must be specified"
    assert isinstance(obj_type, is_string), "dump_table : type must be a string"
    assert len(args) > 0, "dump_table : At least one run must be specified"

    result = db.dump_table(obj_type=obj_type, runs=args, **kwargs)

    print("=" * 80)
    for run, content in sorted(result.items()):
        if content is None:
            print("RUN {0} : Not found".format(run))
        else:
            print("RUN {0}".format(run))
            print('-' * 80)
            for header in content:
                print("{tp:s}[{idx:s}] range [{rb:10d},{re:10d}] version {vs:3d} pass {ps:3d}".format(tp=header['type'],
                                                                                                      idx=header['index'], rb=header['run_range'][0], re=header['run_range'][1], ps=header['pass'], vs=header['version']))
    print("=" * 80)

def dump(db, args, raw, **kwargs):
    """ Dumps all the tables available for a given run into a BZip2 output file.

    Usage: ratdb [options] dump  <run1> [run2-runN] [run3] ...

    Optional flags:
      -g     : Debug mode. Dump all available passes
      -r     : Dump the tables in raw mode. Use with care.

    NOTE: By default this tool dumps only the latest pass of each table

    """

    # Required to accept arguments that may not be used.
    del kwargs

    assert len(args) > 0, "dump: Need at least one run"

    result = db.dump(runs=args, raw=raw)

    if len(result) == 0:
        print("dump : No tables were found.")
    else:
        print('=' * 80)
        for run, content in sorted(result.items()):
            if content is None:
                print("Run {0} : NOT FOUND".format(run))
            else:
                print("Run {0}".format(run))
                print('-' * 80)
                for header in content:
                    print("{tp:s}[{idx:s}] range [{rb:10d},{re:10d}] version {vs:3d} pass {ps:3d}".format(tp=header['type'],
                                                                                                          idx=header['index'], rb=header['run_range'][0], re=header['run_range'][1], ps=header['pass'], vs=header['version']))
            print("=" * 80)

def get_obj_structure(db, obj_type, args, index, **kwargs):
    """ Prints the required list of fields for an object of a given type and version.

    Usage: ratdb [options] get_obj_structure type [version]

    If version is not specified, all versions are returned.
    No other modifiers are accepted

    """

    # Required to accept arguments that may not be used.
    del kwargs

    version = None
    if len(args) is not None:
        assert len(args) < 2, "get_obj_structure : Should have either one version or none"
    if len(args) == 1:
        try:
            version = int(args[0])
        except ValueError:
            # Argument cannot be cast into an int
            raise ValueError('get_obj_structure : Can\'t cast argument [{0}] into a version number.'.format(args[0]))

    result = db.get_object_structure(obj_type=obj_type, index=index, version=version)

    if not result: # dictionary is empty
        print("Object {0} not found.".format(obj_type))

    print('=' * 80)
    for version, structure in sorted(result.items()):
        print("V {0} : {1}".format(version, ','.join(structure)))
        print('-' * 80)
    print('=' * 80)

def get_ratdb_tags(db, **kwargs):
    ''' Returns all tags from the database

        It takes no arguments
        Prints a table with tag, id, date inserted and comment.
    '''

    # Required to accept arguments that may not be used.
    del kwargs

    result = db.get_ratdb_tags()

    if len(result) == 0:
        print("get_ratdb_tags : No tags were found in database")
        return

    # else print a fancy table with all the bells and whistles
    print("=" * 100)
    print("{:^25s} | {:^4s} | {:^35} | {:^12} | {:<s}".format('name', 'id', 'date created', 'RAT release', 'comment'))
    print("-" * 100)
    for e in result:
        print("{:<25s}   {:<4d}   {:<35s}   {:<12s}   {:<s}".format(e['name'], e['id'], e['date'], e['rat'], e['comment']))
    print("=" * 100)


def get_all_types(db, **kwargs):
    ''' Returns all types of objects from the database

        It takes no arguments
        Prints a table with type, index and number of objects in the database along with its maximum version number.
    '''

    # Required to accept arguments that may not be used.
    del kwargs

    result = db.get_db_statistics()

    if result is None:
        print("WARNING: There are no objects in the database")
        return

    print("=" * 80)
    print("{0: <35} :  {1: <8}| {2: <8}".format("table[index]", "counts", "Maximum version"))
    print("_" * 80)
    for rtype, indexes in sorted(result.items()):
        for index, stats in sorted(indexes.items()):
            rtype = str(rtype).replace('"', '')
            index = str(index).replace('"', '')
            count = stats['count']
            maxVersion = stats['maxVersion']
            print("{0: <35} :  {1: <8}| {2: <8}".format('{0}[{1}]'.format(rtype, index), count, maxVersion))
    print("=" * 80)

def get_type_run_range(db, obj_type, index, **kwargs):
    ''' Lists for which runs a given object (and optionally index) a given type exists.

        Usage: ratdb [options] get_type_run_range type[index]

        This function accesses RATDB and collects the run validity of all objects of the specified type.
        It then aggregates the runs into patches of continuous ranges

        If the index is not specified, prints the results for each different index.

    '''

    # Required to accept arguments that may not be used.
    del kwargs

    assert isinstance(obj_type, is_string), "get_type_run_range: Type must be a string"
    if index is None:
        index = ''
    result = db.get_type_run_range(obj_type, index)

    if len(result) == 0:
        print("Object {0}[{1}] was not found in database".format(obj_type, index))
        return

    print('=' * 80)
    for ind, r in sorted(result.items()):
        print('-' * 80)
        print('Index : [{0}]'.format(ind))
        print('-' * 80)
        print('{0}'.format(r))
    print('=' * 80)


commands = {
    'get_run_type': get_run_type,             # Returns the run type for a given run
    'list_obj_run': list_obj_run,             # Lists all the available objects for a given type and index and run
    'list_all_objs_run': list_all_objs_run,  # Lists all objects that exist in the database for a given run
    'upload': upload,                        # uploads a table to the database
    'dump_table': dump_table,                # Dumps a single table to disk for a given run/range
    'dump': dump,                            # Dumps all tables to disk for a given run/range
    'get_obj_structure': get_obj_structure,  # Prints the object structure of a requested object
    'get_ratdb_tags': get_ratdb_tags,
    'get_all_types': get_all_types,          # Get all
    'get_type_run_range': get_type_run_range, # Get all runs for which an object exists
    'help': command_help,                    # Prints the help for specific commands
}

def main():
    '''
    Main ratdb function to call.
    '''
    parser = argparse.ArgumentParser(description=("ratdb management script. "
                                                  "This script connects to RATDB and allows "
                                                  "to make basic browsing operations, "
                                                  "data upload, data download into RATDB files. "
                                                  "Specifying the command argument as help "
                                                  "(with no other arguments) "
                                                  "will list all available commands to be run. "
                                                  "You can also use the help option with commands "
                                                  "to get the respective help messages."))

    parser.add_argument('command', type=str)
    parser.add_argument('command_args', type=str, nargs='*')

    parser.add_argument('-s', '--server', dest='server_url', default='',
                        help=('URL to DB server. '
                              'The format of <url> should be '
                              '"postgres://<user>:<password>@<host>:<port>/<database>". '
                              'If none is supplied the code will look for the '
                              'RATDBSERVER environment variable. '
                              'If none is available it will default to '
                              'postgres://snoplus@pgsql.snopl.us:5400/ratdb'))
    parser.add_argument('-p', '--pass', dest="passnum", type=int, default=-1,
                        help='Specify pass number (only supported in select commands)')
    parser.add_argument('-d', '--debug', action='store_true', dest='debug', default=False,
                        help='Enable debug mode (only supported on select commands)')
    parser.add_argument('-v', '--verbose', action='store_true', dest='verbose', default=False,
                        help='Enable verbosity (extra messages printed to screen)')
    parser.add_argument('-r', '--raw', action='store_true', dest='raw', default=False,
                        help='Enable raw object (only used in dump commands)')
    parser.add_argument('-t', '--tag', dest='tag_name', default='',
                        help='RATDB tag to browse. If none is supplied, latest version is used')

    args = parser.parse_args()
    if args.server_url == "":
        try:
            server_url = os.environ['RATDBSERVER']
        except KeyError:
            server_url = "postgres://snoplus@pgsql.snopl.us:5400/ratdb"

    try:
        # Python 2 byte string should be decoded.
        server_url = server_url.decode('utf-8')
    except AttributeError:
        # Python 3 str is unicode, cannot be decoded.
        pass

    # Replace right and left quotation marks in the URL (not sure how this would happen)...
    server_url = server_url.replace('\u201c', '').replace('\u201d', '')

    rat = None

    command = args.command
    command_args = args.command_args

    try:
        if (command == 'help') and (len(command_args) == 0):
            print("Available commands:")
            print("")
            for key in commands:
                print("  {0}".format(key))
            print("")
            return

        if command not in commands.keys():
            raise ValueError('Unknown command:', command)

        if (command == 'help') and (len(command_args) >= 1):
            print("Helping with command: {}".format(command_args[0]))
            command_help(command_args[0])
            sys.exit(0)

        # Reset the argument list to split out the command
        print("")
        result = urlparse(server_url)
        stripped_url = result.scheme + '://'
        stripped_url += result.username
        if result.password is not None:
            stripped_url += ':**********'
        stripped_url += '@' + result.hostname
        if result.port is not None:
            stripped_url += ':' + str(result.port)
        if result.path is not None and len(result.path) > 0:
            stripped_url += '/' + result.path[1:]

        # Strip the password in case someone adds it to the command line
        print(':: Input DB server : [{0}]'.format(stripped_url))
        print(':: Command  : ', command)
        print(':: Pass     : ', args.passnum)
        print(':: Debug    : ', args.debug)
        print(':: Verbose  : ', args.verbose)
        print(':: Tag      : ', args.tag_name)
        print("")

        #Start the work proper
        rat = RATDBConnector(server=server_url,
                             debug=args.debug,
                             verbose=args.verbose,
                             tag=args.tag_name)

        # Now do something special to parse the arguments
        # If type is part of the arguments, then this will *always* be
        # the first argument
        otype = None
        oindex = None
        # Check if this specific command requires a type as argument
        # pylint: disable=deprecated-method
        if 'obj_type' in getargspec(commands[command]).args:
            #split type and index
            reg = re.search(r"(\w+)\[(\w+)\]", command_args[0])
            if reg is None:
                # there is no index
                otype = command_args[0]
                command_args = command_args[1:]
            else:
                # There is type and index
                otype = reg.group(1)
                oindex = reg.group(2)
                command_args = command_args[1:]

        # Now execute the command
        print("ratdb : Executing command...")
        commands[command](rat, pass_num=args.passnum,
                          obj_type=otype, index=oindex,
                          args=command_args, raw=args.raw)

    except psql.Warning as e:
        print("==>Warning caught")
        # Exception raised for important warnings like data truncations while inserting.
        print_db_error(dberror=e)
    except psql.DatabaseError as e:
        print("==> DatabaseError caught with state {0}".format(e.sqlstate))
        # Exception raised for errors that are related to the database
        # In PyGreSQL, this also has a DatabaseError.sqlstate attribute that
        # contains the SQLSTATE error code of this error.
        print_db_error(dberror=e)
    except psql.InterfaceError as e:
        print("==> InterfaceError caught")
        # Exception raised for errors that are related to the database interface
        # rather than the database itself.
        print_db_error(dberror=e)
    except psql.DataError as e:
        print("==> DataError caught")
        # Exception raised for errors that are due to problems with the processed data
        # like division by zero or numeric value out of range.
        print_db_error(dberror=e)
    except psql.OperationalError as e:
        print("==> OperationalError caught")
        # Exception raised for errors that are related to the database's operation and
        # not necessarily under the control of the programmer, e.g. an unexpected disconnect
        # occurs, the data source name is not found, a transaction could not be processed,
        # or a memory allocation error occurred during processing.
        print_db_error(dberror=e)
    except psql.IntegrityError as e:
        print("==> IntegrityError caught")
        #Exception raised when the relational integrity of the database is affected,
        #e.g. a foreign key check fails.
        print_db_error(dberror=e)
    except psql.ProgrammingError as e:
        print("==> ProgrammingError caught")
        # Exception raised for programming errors, e.g. table not found or already exists,
        # syntax error in the SQL statement or wrong number of parameters specified
        print_db_error(dberror=e)
    except psql.InternalError as e:
        print("==> InternalError caught")
        # This is a PyGreSQL specific exception
        print_db_error(dberror=e)
    except psql.NotSupportedError as e:
        print("==> NotSupportedError caught")
        # Exception raised in case a method or database API was used
        # which is not supported by the database
        print_db_error(dberror=e)
    except psql.Error as e:
        print("==> Error caught")
        # Exception that is the base class of all other error exceptions.
        # You can use this to catch all errors with one single except statement.
        # Warnings are not considered errors and thus do not use this class as base.
        # This is put here in case of using a driver that has not implemented the specific
        # exceptions
        print_db_error(dberror=e)
    except RuntimeError as e:
        print("==> Run time exception: {0}".format(e))
    except ValueError as e:
        print("==> Input Exception: {0}".format(e))
    except SystemExit:
        pass
    except Exception:
        e = exc_info()[1]
        print("Unidentified Exception : {0}".format(e))
    finally:
        # Check the status of the server and close it, if necessary
        if rat is not None and rat.backend == "postgres":
            print("::RATDB : Closing database connection")
            rat.close_ratdb_connection()


if __name__ == '__main__':
    main()