# coding: utf-8
"""Utilities for installing extensions"""

# Copyright (c) Jupyter Development Team.
# Distributed under the terms of the Modified BSD License.

import os
import sys
import importlib
from tornado.log import LogFormatter
from traitlets import Bool, Any
from traitlets.utils.importstring import import_item

from jupyter_core.application import JupyterApp
from jupyter_core.paths import (
    jupyter_config_dir, 
    jupyter_config_path, 
    ENV_CONFIG_PATH, 
    SYSTEM_CONFIG_PATH
)
from jupyter_server._version import __version__
from jupyter_server.config_manager import BaseJSONConfigManager


class ArgumentConflict(ValueError):
    pass

_base_flags = {}
_base_flags.update(JupyterApp.flags)
_base_flags.pop("y", None)
_base_flags.pop("generate-config", None)
_base_flags.update({
    "user" : ({
        "BaseExtensionApp" : {
            "user" : True,
        }}, "Apply the operation only for the given user"
    ),
    "system" : ({
        "BaseExtensionApp" : {
            "user" : False,
            "sys_prefix": False,
        }}, "Apply the operation system-wide"
    ),
    "sys-prefix" : ({
        "BaseExtensionApp" : {
            "sys_prefix" : True,
        }}, "Use sys.prefix as the prefix for installing extensions (for environments, packaging)"
    ),
    "py" : ({
        "BaseExtensionApp" : {
            "python" : True,
        }}, "Install from a Python package"
    )
})
_base_flags['python'] = _base_flags['py']

_base_aliases = {}
_base_aliases.update(JupyterApp.aliases)


class BaseExtensionApp(JupyterApp):
    """Base extension installer app"""
    _log_formatter_cls = LogFormatter
    flags = _base_flags
    aliases = _base_aliases
    version = __version__

    user = Bool(False, config=True, help="Whether to do a user install")
    sys_prefix = Bool(True, config=True, help="Use the sys.prefix as the prefix")
    python = Bool(False, config=True, help="Install from a Python package")

    def _log_format_default(self):
        """A default format for messages"""
        return "%(message)s"

def _get_config_dir(user=False, sys_prefix=False):
    """Get the location of config files for the current context

    Returns the string to the enviornment

    Parameters
    ----------

    user : bool [default: False]
        Get the user's .jupyter config directory
    sys_prefix : bool [default: False]
        Get sys.prefix, i.e. ~/.envs/my-env/etc/jupyter
    """
    user = False if sys_prefix else user
    if user and sys_prefix:
        raise ArgumentConflict("Cannot specify more than one of user or sys_prefix")
    if user:
        extdir = jupyter_config_dir()
    elif sys_prefix:
        extdir = ENV_CONFIG_PATH[0]
    else:
        extdir = SYSTEM_CONFIG_PATH[0]
    return extdir


# Constants for pretty print extension listing function.
# Window doesn't support coloring in the commandline
GREEN_ENABLED = '\033[32menabled\033[0m' if os.name != 'nt' else 'enabled'
RED_DISABLED = '\033[31mdisabled\033[0m' if os.name != 'nt' else 'disabled'
GREEN_OK = '\033[32mOK\033[0m' if os.name != 'nt' else 'ok'
RED_X = '\033[31m X\033[0m' if os.name != 'nt' else ' X'

# ------------------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------------------

class ExtensionValidationError(Exception): pass


def validate_server_extension(import_name):
    """Tries to import the extension module. 
    Raises a validation error if module is not found.
    """
    try:
        mod = importlib.import_module(import_name)
        func = getattr(mod, 'load_jupyter_server_extension')
        version = getattr(mod, '__version__', '')
        return mod, func, version
    # If the extension does not exist, raise an exception
    except ImportError:
        raise ExtensionValidationError('{} is not importable.'.format(import_name))
    # If the extension does not have a `load_jupyter_server_extension` function, raise exception.
    except AttributeError:
        raise ExtensionValidationError('Found module "{}" but cannot load it.'.format(import_name))


def toggle_server_extension_python(import_name, enabled=None, parent=None, user=False, sys_prefix=True):
    """Toggle the boolean setting for a given server extension
    in a Jupyter config file.
    """
    sys_prefix = False if user else sys_prefix
    config_dir = _get_config_dir(user=user, sys_prefix=sys_prefix)
    cm = BaseJSONConfigManager(parent=parent, config_dir=config_dir)
    cfg = cm.get("jupyter_server_config")
    server_extensions = (
        cfg.setdefault("ServerApp", {})
        .setdefault("jpserver_extensions", {})
    )
    old_enabled = server_extensions.get(import_name, None)
    new_enabled = enabled if enabled is not None else not old_enabled
    server_extensions[import_name] = new_enabled
    cm.update("jupyter_server_config", cfg)

# ----------------------------------------------------------------------
# Applications
# ----------------------------------------------------------------------

flags = {}
flags.update(BaseExtensionApp.flags)
flags.pop("y", None)
flags.pop("generate-config", None)
flags.update({
    "user" : ({
        "ToggleServerExtensionApp" : {
            "user" : True,
        }}, "Perform the operation for the current user"
    ),
    "system" : ({
        "ToggleServerExtensionApp" : {
            "user" : False,
            "sys_prefix": False,
        }}, "Perform the operation system-wide"
    ),
    "sys-prefix" : ({
        "ToggleServerExtensionApp" : {
            "sys_prefix" : True,
        }}, "Use sys.prefix as the prefix for installing server extensions"
    ),
    "py" : ({
        "ToggleServerExtensionApp" : {
            "python" : True,
        }}, "Install from a Python package"
    ),
})
flags['python'] = flags['py']


class ToggleServerExtensionApp(BaseExtensionApp):
    """A base class for enabling/disabling extensions"""
    name = "jupyter server extension enable/disable"
    description = "Enable/disable a server extension using frontend configuration files."
    
    flags = flags

    user = Bool(False, config=True, help="Whether to do a user install")
    sys_prefix = Bool(True, config=True, help="Use the sys.prefix as the prefix")
    python = Bool(False, config=True, help="Install from a Python package")
    _toggle_value = Bool()
    _toggle_pre_message = ''
    _toggle_post_message = ''
    
    def toggle_server_extension(self, import_name):
        """Change the status of a named server extension.

        Uses the value of `self._toggle_value`.

        Parameters
        ---------

        import_name : str
            Importable Python module (dotted-notation) exposing the magic-named
            `load_jupyter_server_extension` function
        """
        try:
            self.log.info("{}: {}".format(self._toggle_pre_message.capitalize(), import_name))
            # Validate the server extension.
            self.log.info("    - Validating {}...".format(import_name))
            _, __, version = validate_server_extension(import_name)

            # Toggle the server extension to active.
            toggle_server_extension_python(
                import_name, 
                self._toggle_value, 
                parent=self, 
                user=self.user,
                sys_prefix=self.sys_prefix
            )
            self.log.info("      {} {} {}".format(import_item, version, GREEN_OK))

            # If successful, let's log.
            self.log.info("    - Extension successfully {}.".format(self._toggle_post_message))
        except ExtensionValidationError as err:
            self.log.info("     {} Validation failed: {}".format(RED_X, err))

    def toggle_server_extension_python(self, package):
        """Change the status of some server extensions in a Python package.

        Uses the value of `self._toggle_value`.

        Parameters
        ---------

        package : str
            Importable Python module exposing the
            magic-named `_jupyter_server_extension_paths` function
        """
        _, server_exts = _get_server_extension_metadata(package)
        for server_ext in server_exts:
            module = server_ext['module']
            self.toggle_server_extension(module)

    def start(self):
        """Perform the App's actions as configured"""
        if not self.extra_args:
            sys.exit('Please specify a server extension/package to enable or disable')
        for arg in self.extra_args:
            if self.python:
                self.toggle_server_extension_python(arg)
            else:
                self.toggle_server_extension(arg)


class EnableServerExtensionApp(ToggleServerExtensionApp):
    """An App that enables (and validates) Server Extensions"""
    name = "jupyter server extension enable"
    description = """
    Enable a server extension in configuration.
    
    Usage
        jupyter server extension enable [--system|--sys-prefix]
    """
    _toggle_value = True
    _toggle_pre_message = "enabling"
    _toggle_post_message = "enabled"


class DisableServerExtensionApp(ToggleServerExtensionApp):
    """An App that disables Server Extensions"""
    name = "jupyter server extension disable"
    description = """
    Disable a server extension in configuration.
    
    Usage
        jupyter server extension disable [--system|--sys-prefix]
    """
    _toggle_value = False
    _toggle_pre_message = "disabling"
    _toggle_post_message = "disabled"


class ListServerExtensionsApp(BaseExtensionApp):
    """An App that lists (and validates) Server Extensions"""
    name = "jupyter server extension list"
    version = __version__
    description = "List all server extensions known by the configuration system"

    def list_server_extensions(self):
        """List all enabled and disabled server extensions, by config path

        Enabled extensions are validated, potentially generating warnings.
        """
        config_dirs = jupyter_config_path()
        for config_dir in config_dirs:
            cm = BaseJSONConfigManager(parent=self, config_dir=config_dir)
            data = cm.get("jupyter_server_config")
            server_extensions = (
                data.setdefault("ServerApp", {})
                .setdefault("jpserver_extensions", {})
            )
            if server_extensions:
                self.log.info(u'config dir: {}'.format(config_dir))
            for import_name, enabled in server_extensions.items():
                self.log.info(u'    {} {}'.format(
                              import_name,
                              GREEN_ENABLED if enabled else RED_DISABLED))
                try:
                    self.log.info("    - Validating {}...".format(import_name))
                    _, __, version = validate_server_extension(import_name)
                    self.log.info("      {} {} {}".format(import_name, version, GREEN_OK))

                except ExtensionValidationError as err:
                    self.log.warn("      {} {}".format(RED_X, err))

    def start(self):
        """Perform the App's actions as configured"""
        self.list_server_extensions()


_examples = """
jupyter server extension list                        # list all configured server extensions
jupyter server extension enable --py <packagename>   # enable all server extensions in a Python package
jupyter server extension disable --py <packagename>  # disable all server extensions in a Python package
"""


class ServerExtensionApp(BaseExtensionApp):
    """Root level server extension app"""
    name = "jupyter server extension"
    version = __version__
    description = "Work with Jupyter server extensions"
    examples = _examples

    subcommands = dict(
        enable=(EnableServerExtensionApp, "Enable a server extension"),
        disable=(DisableServerExtensionApp, "Disable a server extension"),
        list=(ListServerExtensionsApp, "List server extensions")
    )

    def start(self):
        """Perform the App's actions as configured"""
        super(ServerExtensionApp, self).start()

        # The above should have called a subcommand and raised NoStart; if we
        # get here, it didn't, so we should self.log.info a message.
        subcmds = ", ".join(sorted(self.subcommands))
        sys.exit("Please supply at least one subcommand: %s" % subcmds)


main = ServerExtensionApp.launch_instance

# ------------------------------------------------------------------------------
# Private API
# ------------------------------------------------------------------------------

def _get_server_extension_metadata(module):
    """Load server extension metadata from a module.

    Returns a tuple of (
        the package as loaded
        a list of server extension specs: [
            {
                "module": "mockextension"
            }
        ]
    )

    Parameters
    ----------

    module : str
        Importable Python module exposing the
        magic-named `_jupyter_server_extension_paths` function
    """
    m = import_item(module)
    if not hasattr(m, '_jupyter_server_extension_paths'):
        raise KeyError(u'The Python module {} does not include any valid server extensions'.format(module))
    return m, m._jupyter_server_extension_paths()

if __name__ == '__main__':
    main()
