"""
AccountingCLI class implementing command line administrative interface to
DIRAC Accounting DataStore Service
"""

from __future__ import print_function
from __future__ import absolute_import
from __future__ import division
__RCSID__ = "$Id$"

import sys

import six

from DIRAC import gLogger
from DIRAC.Core.Base.CLI import CLI, colorize
from DIRAC.AccountingSystem.Client.DataStoreClient import DataStoreClient


class AccountingCLI(CLI):

  def __init__(self):
    CLI.__init__(self)
    self.do_connect(None)

  def start(self):
    """
    Start the command loop
    """
    if not self.connected:
      gLogger.error("Client is not connected")
    try:
      self.cmdloop()
    except KeyboardInterrupt:
      gLogger.warn("Received a keyboard interrupt.")
      self.do_quit("")

  def do_connect(self, args):
    """
    Tries to connect to the server
        Usage: connect
    """
    gLogger.info("Trying to connect to server")
    self.connected = False
    self.prompt = "(%s)> " % colorize("Not connected", "red")
    acClient = DataStoreClient()
    retVal = acClient.ping()
    if retVal['OK']:
      self.prompt = "(%s)> " % colorize("Connected", "green")
      self.connected = True

  def printComment(self, comment):
    commentList = comment.split("\n")
    for commentLine in commentList[:-1]:
      print("# %s" % commentLine.strip())

  def showTraceback(self):
    import traceback
    type, value = sys.exc_info()[:2]
    print("________________________\n")
    print("Exception", type, ":", value)
    traceback.print_tb(sys.exc_info()[2])
    print("________________________\n")

  def do_registerType(self, args):
    """
    Registers a new accounting type
      Usage : registerType <typeName>
      <DIRACRoot>/DIRAC/AccountingSystem/Client/Types/<typeName>
       should exist and inherit the base type
    """
    try:
      argList = args.split()
      if argList:
        typeName = argList[0].strip()
      else:
        gLogger.error("No type name specified")
        return
      # Try to import the type
      try:
        typeModule = __import__("DIRAC.AccountingSystem.Client.Types.%s" % typeName,
                                globals(),
                                locals(), typeName)
        typeClass = getattr(typeModule, typeName)
      except Exception as e:
        gLogger.error("Can't load type %s: %s" % (typeName, str(e)))
        return
      gLogger.info("Loaded type %s" % typeClass.__name__)
      typeDef = typeClass().getDefinition()
      acClient = DataStoreClient()
      retVal = acClient.registerType(*typeDef)
      if retVal['OK']:
        gLogger.info("Type registered successfully")
      else:
        gLogger.error("Error: %s" % retVal['Message'])
    except BaseException:
      self.showTraceback()

  def do_resetBucketLength(self, args):
    """
    Set the bucket Length. Will trigger a recalculation of buckets. Can take a while.
      Usage : resetBucketLength <typeName>
      <DIRACRoot>/DIRAC/AccountingSystem/Client/Types/<typeName>
       should exist and inherit the base type
    """
    try:
      argList = args.split()
      if argList:
        typeName = argList[0].strip()
      else:
        gLogger.error("No type name specified")
        return
      # Try to import the type
      try:
        typeModule = __import__("DIRAC.AccountingSystem.Client.Types.%s" % typeName,
                                globals(),
                                locals(), typeName)
        typeClass = getattr(typeModule, typeName)
      except Exception as e:
        gLogger.error("Can't load type %s: %s" % (typeName, str(e)))
        return
      gLogger.info("Loaded type %s" % typeClass.__name__)
      typeDef = typeClass().getDefinition()
      acClient = DataStoreClient()
      retVal = acClient.setBucketsLength(typeDef[0], typeDef[3])
      if retVal['OK']:
        gLogger.info("Type registered successfully")
      else:
        gLogger.error("Error: %s" % retVal['Message'])
    except BaseException:
      self.showTraceback()

  def do_regenerateBuckets(self, args):
    """
    Regenerate buckets for type. Can take a while.
      Usage : regenerateBuckets <typeName>
      <DIRACRoot>/DIRAC/AccountingSystem/Client/Types/<typeName>
       should exist and inherit the base type
    """
    try:
      argList = args.split()
      if argList:
        typeName = argList[0].strip()
      else:
        gLogger.error("No type name specified")
        return
      # Try to import the type
      try:
        typeModule = __import__("DIRAC.AccountingSystem.Client.Types.%s" % typeName,
                                globals(),
                                locals(), typeName)
        typeClass = getattr(typeModule, typeName)
      except Exception as e:
        gLogger.error("Can't load type %s: %s" % (typeName, str(e)))
        return
      gLogger.info("Loaded type %s" % typeClass.__name__)
      typeDef = typeClass().getDefinition()
      acClient = DataStoreClient()
      retVal = acClient.regenerateBuckets(typeDef[0])
      if retVal['OK']:
        gLogger.info("Buckets recalculated!")
      else:
        gLogger.error("Error: %s" % retVal['Message'])
    except BaseException:
      self.showTraceback()

  def do_showRegisteredTypes(self, args):
    """
    Get a list of registered types
      Usage : showRegisteredTypes
    """
    try:
      acClient = DataStoreClient()
      retVal = acClient.getRegisteredTypes()

      print(retVal)

      if not retVal['OK']:
        gLogger.error("Error: %s" % retVal['Message'])
        return
      for typeList in retVal['Value']:
        print(typeList[0])
        print(" Key fields:\n  %s" % "\n  ".join(typeList[1]))
        print(" Value fields:\n  %s" % "\n  ".join(typeList[2]))
    except BaseException:
      self.showTraceback()

  def do_deleteType(self, args):
    """
    Delete a registered accounting type.
      Usage : deleteType <typeName>
      WARN! It will delete all data associated to that type! VERY DANGEROUS!
      If you screw it, you'll discover a new dimension of pain and doom! :)
    """
    try:
      argList = args.split()
      if argList:
        typeName = argList[0].strip()
      else:
        gLogger.error("No type name specified")
        return
      while True:
        choice = six.moves.input(
            "Are you completely sure you want to delete type %s and all it's data? yes/no [no]: " %
            typeName)
        choice = choice.lower()
        if choice in ("yes", "y"):
          break
        else:
          print("Delete aborted")
          return
      acClient = DataStoreClient()
      retVal = acClient.deleteType(typeName)
      if not retVal['OK']:
        gLogger.error("Error: %s" % retVal['Message'])
        return
      print("Hope you meant it, because it's done")
    except BaseException:
      self.showTraceback()

  def do_compactBuckets(self, args):
    """
    Compact buckets table
      Usage : compactBuckets
    """
    try:
      acClient = DataStoreClient()
      retVal = acClient.compactDB()
      if not retVal['OK']:
        gLogger.error("Error: %s" % retVal['Message'])
        return
      gLogger.info("Done")
    except BaseException:
      self.showTraceback()
