#!/usr/bin/env python
# -- coding: utf-8 --
import os
from rat import ratdb
import inspect
import optparse
from sys import exc_info,exit

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)
import re
from types import StringType, IntType
import urlparse


def main_help():
    print """
ratdb management script.

Usage: ratdb [options] <command> [command args]

This script connects to RATDB and allows to make basic browsing operations, data upload, 
data download into RATDB files.

You can also use the help option with commands to get the respective help messages.

Available options:

  -s / --server <url> : Database server to connect to. The format of <url> should be 
                         "postgres://<user>:<password>@<host>:<port>/<database>"
  -p / --pass <pass>  : Select a specific pass to apply the operation (only implemented in some commands)
  -d / --debug        : Activate debug mode in some commands (only implemented in a select few)
  -v / --verbose      : Activate verbose mode (extra messages are printed to screen)

Available commands:
"""
    for key in commands:
      print "  {0}".format(key)
    print ""

def command_help(arg):
  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
        
      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
  """
  result = db.get_run_type(runs=args,pass_num=pass_num)
  # Now parse the result to show
  print "="*80
  for run,contents in result.iteritems():
    if db.debug:
      if contents is not None:
        print "RUN : [{0}]".format(run)
        print '-'*80
        for np,rt in contents.iteritems():
          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,**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]

"""
    assert len(args)>0, "list_obj_run : At least one run must be specified"
    
    if db.verbose:
      print "kwargs : ", kwargs

    result = db.list_objects_type_run(obj_type=obj_type,index=kwargs['index'],runs=args)

    print "="*80
    for run,content in result.iteritems():
      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

"""

    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 result.iteritems():
      print "-"*20
      print "RUN {0}".format(run)
      print "-"*20
      if contents is None:
        print "NOT FOUND"
      else:
        for header in contents:
          #print header
          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,
    }

 
"""
    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 type(obj_type) is StringType,"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 result.iteritems():
      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, **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

    """

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

    result = db.dump(runs=args,**kwargs)
    
    if len(result) == 0:
      print "dump : No tables were found."
    else:
      print '='*80
      for run,content in result.iteritems():
        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,**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
  """
  
  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,version=version,**kwargs)
  
  if not result: # dictionary is empty
    print "Object {0} not found.".format(obj_type)
  print '='*80
  #print "V : {0:15s} {1:65s}".format("Version","Structure")
  #print '-'*80
  for version,structure in result.iteritems():
    print "V {0} : {1}".format(version,','.join(structure))
    print '-'*80
  print '='*80
  
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.
  '''

  result = db.get_db_statistics()
  
  #print result

  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 type,indexes in result.iteritems():
    for index,stats in indexes.iteritems():
      type = str(type).replace('"', '')
      index = str(index).replace('"', '')
      count = stats['count']
      maxVersion = stats['maxVersion']
      print "{0: <35} :  {1: <8}| {2: <8}".format('{0}[{1}]'.format(type,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.

  '''
  assert type(obj_type) is StringType,"get_type_run_range: Type must be a string"
  if index is None:
    index = ''
  result = db.get_type_run_range(obj_type,index,**kwargs)
  
  if len(result) == 0:
    print "Object {0}[{1}] was not found in database".format(obj_type,index)
    return

  print '='*80
  for index,range in result.iteritems():
    print '-'*80
    print 'Index : [{0}]'.format(index)
    print '-'*80
    print '{0}'.format(range)
  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_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():
    parser = optparse.OptionParser()
    parser.add_option('-s', '--server', dest='server_url', help='URL to DB server. 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',
                      default='')
    parser.add_option('-p','--pass',dest="passnum",type="int",help='Specify pass number (only supported in select commands)', default=-1)
    parser.add_option('-d', '--debug', action='store_true',dest='debug', help='Enable debug mode (only supported on select commands)',default=False)
    parser.add_option('-v', '--verbose', action='store_true',dest='verbose', help='Enable verbosity', default=False)
    parser.add_option('-r', '--raw', action='store_true',dest='raw', help='Enable raw object (only used in dump commands)', default=False)


    (options, args) = parser.parse_args()
    if options.server_url == "":
      try:
        options.server_url = os.environ['RATDBSERVER']
      except KeyError:
        options.server_url = "postgres://snoplus@pgsql.snopl.us:5400/ratdb"
    options.server_url = options.server_url.decode('utf-8')
    options.server_url = options.server_url.replace(u'\u201c', '').replace(u'\u201d', '')

    print "Arguments : ", args
    rat = None
    try:
      if len(args) == 0 or (args[0] == 'help' and len(args) == 1):
        main_help()
        return
  
      if args[0] not in commands.keys():
        raise ValueError('Unknown command:', args[0])
        sys.exit(0)
      elif (args[0] == 'help') and (len(args)>1):
        print "Helping with command ", args[1]
        command_help(args[1])
        sys.exit(0)

      # Reset the argument list to split out the command
      command = args[0]
      args = args[1:]
      print ""
      result = urlparse.urlparse(options.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     : ', options.passnum
      print ':: Debug    : ', options.debug
      print ':: Verbose  : ', options.verbose
      print ""

  
      #Start the work proper
      rat = ratdb.RATDBConnector(server = options.server_url,debug=options.debug, verbose=options.verbose)
      # 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
      if 'obj_type' in inspect.getargspec(commands[command]).args:
        #split type and index
        reg = re.search(r"(\w+)\[(\w+)\]",args[0])
        if reg is None:
          # there is no index
          otype = args[0]
          args = args[1:]
        else:
          # There is type and index
          otype = reg.group(1)
          oindex = reg.group(2)
          args = args[1:]
      
      # Now execute the command
      print "ratdb : Executing command..."
      commands[command](rat,pass_num=options.passnum,obj_type=otype,index=oindex,args=args,raw=options.raw)
    except psql.Warning as e:
      print "==>Warning caught"
      # Exception raised for important warnings like data truncations while inserting.
      ratdb.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.
      ratdb.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.
      ratdb.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.
      ratdb.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.
      ratdb.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.
      ratdb.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
      ratdb.print_db_error(dberror=e)
    except psql.InternalError as e:
      print "==> InternalError caught"
      # This is a PyGreSQL specific exception
      ratdb.print_db_error(dberror=e)
    except psql.NotSupportedError:
      print "==> NotSupportedError caught"
      # Exception raised in case a method or database API was used which is not supported by the database
      ratdb.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
      ratdb.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 str:
      print str
    except SystemExit:
      pass
    except:
      e = exc_info()[1]
      print "Unidentified Exception : {0}".format(e)
      
    finally:
      #print "Closing the server."
      # 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()