"""Helpers for working with extensions to DIRAC"""
from __future__ import print_function
from __future__ import absolute_import
from __future__ import division

import argparse
from collections import defaultdict
import fnmatch
import functools
import glob
import importlib
import os
import pkgutil
import sys

import importlib_resources
import six
import DIRAC

try:
  from importlib.machinery import PathFinder
except ImportError:
  # Fallback for Python 2
  import imp

  class ModuleSpec(object):
    def __init__(self, name, submodule_search_locations):
      self.name = name
      self.submodule_search_locations = submodule_search_locations

  class PathFinder(object):
    @classmethod
    def find_spec(cls, name, path=None):
      try:
        _, pathname, _ = imp.find_module(name, path)
      except ImportError:
        return None
      else:
        return ModuleSpec(name, [pathname])


def iterateThenSort(func):
  """Convenience decorator the find* functions

  Converts a function that takes a python module into one which accepts one or
  modules/module names. The returned values are converted to a sorted list of
  unique values.
  """
  @functools.wraps(func)
  def newFunc(modules, *args, **kwargs):
    if not isinstance(modules, (list, tuple)):
      modules = [modules]

    results = set()
    for module in modules:
      if isinstance(module, six.string_types):
        module = importlib.import_module(module)
      results |= set(func(module, *args, **kwargs))
    return sorted(results)
  return newFunc


@iterateThenSort
def findSystems(module):
  """Find the systems for one or more DIRAC extension(s)

  :param list/str/module module: One or more Python modules or Python module names
  :returns: list of system names
  """
  return {x.name for x in _findSystems(module)}


def findAgents(modules):
  """Find the agents for one or more DIRAC extension(s)

  :param list/str/module module: One or more Python modules or Python module names
  :returns: list of tuples of the form (SystemName, AgentName)
  """
  return findModules(modules, "Agent", "*Agent")


def findExecutors(modules):
  """Find the executors for one or more DIRAC extension(s)

  :param list/str/module module: One or more Python modules or Python module names
  :returns: list of tuples of the form (SystemName, ExecutorName)
  """
  return findModules(modules, "Executor")


def findServices(modules):
  """Find the services for one or more DIRAC extension(s)

  :param list/str/module module: One or more Python modules or Python module names
  :returns: list of tuples of the form (SystemName, ServiceName)
  """
  return findModules(modules, "Service", "*Handler")


@iterateThenSort
def findDatabases(module):
  """Find the DB SQL schema defintions for one or more DIRAC extension(s)

  :param list/str/module module: One or more Python modules or Python module names
  :returns: list of tuples of the form (SystemName, dbSchemaFilename)
  """
  # This can be "fn.name" when DIRAC is Python 3 only
  return {
      (system.name, os.path.basename(str(fn)))
      for system, fn in _findFile(module, "DB", "*DB.sql")
  }


@iterateThenSort
def findModules(module, submoduleName, pattern="*"):
  """Find the direct submodules from one or more DIRAC extension(s) that match a pattern

  :param list/str/module module: One or more Python modules or Python module names
  :param str submoduleName: The submodule under ``module`` in which to look
  :param str pattern: A ``fnmatch``-style pattern that the submodule must match
  :returns: list of tuples of the form (SystemName, ServiceName)
  """
  for system in _findSystems(module):
    agentModule = PathFinder.find_spec(submoduleName, path=system.submodule_search_locations)
    if not agentModule:
      continue
    for _, name, _ in pkgutil.iter_modules(agentModule.submodule_search_locations):
      if fnmatch.fnmatch(name, pattern):
        yield system.name, name


def entrypointToExtension(entrypoint):
  """"Get the extension name from an EntryPoint object"""
  # In Python 3.9 this can be "entrypoint.module"
  module = entrypoint.pattern.match(entrypoint.value).groupdict()["module"]
  return module.split(".")[0]


def extensionsByPriority():
  """Get the list of installed extensions"""
  if six.PY3:
    return _extensionsByPriorityPy3()
  else:
    return _extensionsByPriorityPy2()


def _extensionsByPriorityPy2():
  initList = glob.glob(os.path.join(DIRAC.rootPath, '*DIRAC', '__init__.py'))
  extensions = [os.path.basename(os.path.dirname(k)) for k in initList]
  # Return the extensions, sorting such that vanilla DIRAC is always last
  # It's not correct but it's less incorrect than ComponentInstaller.getExtensions
  return sorted(extensions, key=lambda x: (x == "DIRAC", x))


def _extensionsByPriorityPy3():
  """Discover extensions using the setuptools metadata"""
  # This is Python 3 only, Python 2 installations should never try to use this
  from importlib import metadata  # pylint: disable=no-name-in-module

  priorties = defaultdict(list)
  for entrypoint in set(metadata.entry_points()['dirac']):
    extensionName = entrypointToExtension(entrypoint)
    extension_metadata = entrypoint.load()()
    priorties[extension_metadata["priority"]].append(extensionName)

  extensions = []
  for priority, extensionNames in sorted(priorties.items(), reverse=True):
    if len(extensionNames) != 1:
      print(
          "WARNING: Found multiple extensions with priority",
          "{} ({})".format(priority, extensionNames),
          file=sys.stderr,
      )
    # If multiple are passed, sort the extensions so things are deterministic at least
    extensions.extend(sorted(extensionNames))
  return extensions


def getExtensionMetadata(extensionName):
  """Get the metadata for a given extension name"""
  # This is Python 3 only, Python 2 installations should never try to use this
  from importlib import metadata  # pylint: disable=no-name-in-module

  for entrypoint in metadata.entry_points()['dirac']:
    if extensionName == entrypointToExtension(entrypoint):
      return entrypoint.load()()


def recurseImport(modName, parentModule=None, hideExceptions=False):
  from DIRAC import S_OK, S_ERROR, gLogger

  if parentModule is not None:
    raise NotImplementedError(parentModule)
  try:
    return S_OK(importlib.import_module(modName))
  except ImportError as excp:
    if str(excp).startswith("No module named"):
      return S_OK()
    errMsg = "Can't load %s" % modName
    if not hideExceptions:
      gLogger.exception(errMsg)
    return S_ERROR(errMsg)


def _findSystems(module):
  """Implementation of _findSystems that returns a generator of system names"""
  for _, name, _ in pkgutil.iter_modules(module.__path__):
    if name.endswith("System"):
      yield PathFinder.find_spec(name, path=module.__path__)


def _findFile(module, submoduleName, pattern="*"):
  """Implementation of findDatabases"""
  for system in _findSystems(module):
    try:
      dbModule = importlib_resources.files(".".join([module.__name__, system.name, submoduleName]))
    except ImportError:
      continue
    for file in dbModule.iterdir():
      if fnmatch.fnmatch(file.name, pattern):
        yield system, file


def parseArgs():
  """CLI interface for use with the DIRAC integration tests"""
  parser = argparse.ArgumentParser()
  if six.PY3:
    subparsers = parser.add_subparsers(required=True, dest='function')
  else:
    subparsers = parser.add_subparsers()
  defaultExtensions = extensionsByPriority()
  for func in [findSystems, findAgents, findExecutors, findServices, findDatabases]:
    subparser = subparsers.add_parser(func.__name__)
    subparser.add_argument("--extensions", nargs="+", default=defaultExtensions)
    subparser.set_defaults(func=func)
  args = parser.parse_args()
  # Get the result and print it
  extensions = [importlib.import_module(e) for e in args.extensions]
  for result in args.func(extensions):
    if not isinstance(result, str):
      result = " ".join(result)
    print(result)


if __name__ == "__main__":
  parseArgs()