# -*- coding: utf-8 -*-
import inspect
import os
import sys
import traceback
import warnings

from django.db import connections
from django.conf import settings
from django.core.management.base import BaseCommand, CommandError
from django.utils.datastructures import OrderedSet

from django_extensions.management.shells import import_objects
from django_extensions.management.utils import signalcommand
from django_extensions.management.debug_cursor import monkey_patch_cursordebugwrapper


def use_vi_mode():
    editor = os.environ.get("EDITOR")
    if not editor:
        return False
    editor = os.path.basename(editor)
    return editor.startswith("vi") or editor.endswith("vim")


def shell_runner(flags, name, help=None):
    """
    Decorates methods with information about the application they are starting

    :param flags: The flags used to start this runner via the ArgumentParser.
    :param name: The name of this runner for the help text for the ArgumentParser.
    :param help: The optional help for the ArgumentParser if the dynamically generated help is not sufficient.
    """  # noqa: E501

    def decorator(fn):
        fn.runner_flags = flags
        fn.runner_name = name
        fn.runner_help = help

        return fn

    return decorator


class Command(BaseCommand):
    help = "Like the 'shell' command but autoloads the models of all installed Django apps."  # noqa: E501
    extra_args = None
    tests_mode = False

    def __init__(self):
        super().__init__()
        self.runners = [
            member
            for name, member in inspect.getmembers(self)
            if hasattr(member, "runner_flags")
        ]

    def add_arguments(self, parser):
        super().add_arguments(parser)

        group = parser.add_mutually_exclusive_group()
        for runner in self.runners:
            if runner.runner_help:
                help = runner.runner_help
            else:
                help = "Tells Django to use %s." % runner.runner_name

            group.add_argument(
                *runner.runner_flags,
                action="store_const",
                dest="runner",
                const=runner,
                help=help,
            )

        parser.add_argument(
            "--connection-file",
            action="store",
            dest="connection_file",
            help="Specifies the connection file to use if using the --kernel option",
        )
        parser.add_argument(
            "--no-startup",
            action="store_true",
            dest="no_startup",
            default=False,
            help=(
                "When using plain Python, ignore the PYTHONSTARTUP environment "
                "variable and ~/.pythonrc.py script."
            ),
        )
        parser.add_argument(
            "--use-pythonrc",
            action="store_true",
            dest="use_pythonrc",
            default=False,
            help=(
                "When using plain Python, load the PYTHONSTARTUP environment variable "
                "and ~/.pythonrc.py script."
            ),
        )
        parser.add_argument(
            "--print-sql",
            action="store_true",
            default=False,
            help="Print SQL queries as they're executed",
        )
        parser.add_argument(
            "--truncate-sql",
            action="store",
            type=int,
            help="Truncate SQL queries to a number of characters.",
        )
        parser.add_argument(
            "--print-sql-location",
            action="store_true",
            default=False,
            help="Show location in code where SQL query generated from",
        )
        parser.add_argument(
            "--dont-load",
            action="append",
            dest="dont_load",
            default=[],
            help="Ignore autoloading of some apps/models. Can be used several times.",
        )
        parser.add_argument(
            "--quiet-load",
            action="store_true",
            default=False,
            dest="quiet_load",
            help="Do not display loaded models messages",
        )
        parser.add_argument(
            "--vi",
            action="store_true",
            default=use_vi_mode(),
            dest="vi_mode",
            help="Load Vi key bindings (for --ptpython and --ptipython)",
        )
        parser.add_argument(
            "--no-browser",
            action="store_true",
            default=False,
            dest="no_browser",
            help="Don't open the notebook in a browser after startup.",
        )
        parser.add_argument(
            "-c",
            "--command",
            help=(
                "Instead of opening an interactive shell, "
                "run a command as Django and exit."
            ),
        )

    def run_from_argv(self, argv):
        if "--" in argv[2:]:
            idx = argv.index("--")
            self.extra_args = argv[idx + 1 :]
            argv = argv[:idx]
        return super().run_from_argv(argv)

    def get_ipython_arguments(self, options):
        ipython_args = "IPYTHON_ARGUMENTS"
        arguments = getattr(settings, ipython_args, [])
        if not arguments:
            arguments = os.environ.get(ipython_args, "").split()
        return arguments

    def get_notebook_arguments(self, options):
        notebook_args = "NOTEBOOK_ARGUMENTS"
        arguments = getattr(settings, notebook_args, [])
        if not arguments:
            arguments = os.environ.get(notebook_args, "").split()
        return arguments

    def get_imported_objects(self, options):
        imported_objects = import_objects(options, self.style)
        if self.tests_mode:
            # save imported objects so we can run tests against it later
            self.tests_imported_objects = imported_objects
        return imported_objects

    @shell_runner(flags=["--kernel"], name="IPython Kernel")
    def get_kernel(self, options):
        try:
            from IPython import release

            if release.version_info[0] < 2:
                print(
                    self.style.ERROR("--kernel requires at least IPython version 2.0")
                )
                return
            from IPython import start_kernel
        except ImportError:
            return traceback.format_exc()

        def run_kernel():
            imported_objects = self.get_imported_objects(options)
            kwargs = dict(
                argv=[],
                user_ns=imported_objects,
            )
            connection_file = options["connection_file"]
            if connection_file:
                kwargs["connection_file"] = connection_file
            start_kernel(**kwargs)

        return run_kernel

    def load_base_kernel_spec(self, app):
        """Finds and returns the base Python kernelspec to extend from."""
        ksm = app.kernel_spec_manager
        try_spec_names = getattr(
            settings,
            "NOTEBOOK_KERNEL_SPEC_NAMES",
            [
                "python3",
                "python",
            ],
        )

        if isinstance(try_spec_names, str):
            try_spec_names = [try_spec_names]

        ks = None
        for spec_name in try_spec_names:
            try:
                ks = ksm.get_kernel_spec(spec_name)
                break
            except Exception:
                continue
        if not ks:
            raise CommandError(
                "No notebook (Python) kernel specs found. Tried %r" % try_spec_names
            )

        return ks

    def generate_kernel_specs(self, app, ipython_arguments):
        """Generate an IPython >= 3.0 kernelspec that loads django extensions"""
        ks = self.load_base_kernel_spec(app)
        ks.argv.extend(ipython_arguments)
        ks.display_name = getattr(
            settings, "IPYTHON_KERNEL_DISPLAY_NAME", "Django Shell-Plus"
        )

        manage_py_dir, manage_py = os.path.split(os.path.realpath(sys.argv[0]))
        if manage_py == "manage.py" and os.path.isdir(manage_py_dir):
            pythonpath = ks.env.get("PYTHONPATH", os.environ.get("PYTHONPATH", ""))
            pythonpath = pythonpath.split(os.pathsep)
            if manage_py_dir not in pythonpath:
                pythonpath.append(manage_py_dir)

            ks.env["PYTHONPATH"] = os.pathsep.join(filter(None, pythonpath))

        return {"django_extensions": ks}

    def run_notebookapp(self, app_init, options, use_kernel_specs=True, history=True):
        no_browser = options["no_browser"]

        if self.extra_args:
            # if another '--' is found split the arguments notebook, ipython
            if "--" in self.extra_args:
                idx = self.extra_args.index("--")
                notebook_arguments = self.extra_args[:idx]
                ipython_arguments = self.extra_args[idx + 1 :]
            # otherwise pass the arguments to the notebook
            else:
                notebook_arguments = self.extra_args
                ipython_arguments = []
        else:
            notebook_arguments = self.get_notebook_arguments(options)
            ipython_arguments = self.get_ipython_arguments(options)

        # Treat IPYTHON_ARGUMENTS from settings
        if "django_extensions.management.notebook_extension" not in ipython_arguments:
            ipython_arguments.extend(
                ["--ext", "django_extensions.management.notebook_extension"]
            )

        # Treat NOTEBOOK_ARGUMENTS from settings
        if no_browser and "--no-browser" not in notebook_arguments:
            notebook_arguments.append("--no-browser")
        if "--notebook-dir" not in notebook_arguments and not any(
            e.startswith("--notebook-dir=") for e in notebook_arguments
        ):
            notebook_arguments.extend(["--notebook-dir", "."])

        # IPython < 3 passes through kernel args from notebook CLI
        if not use_kernel_specs:
            notebook_arguments.extend(ipython_arguments)

        # disable history if not already configured in some other way
        if not history and not any(
            arg.startswith("--HistoryManager") for arg in ipython_arguments
        ):
            ipython_arguments.append("--HistoryManager.enabled=False")

        if not callable(app_init):
            app = app_init
            warnings.warn(
                "Initialize should be a callable not an app instance",
                DeprecationWarning,
            )
            app.initialize(notebook_arguments)
        else:
            app = app_init(notebook_arguments)

        # IPython >= 3 uses kernelspecs to specify kernel CLI args
        if use_kernel_specs:
            ksm = app.kernel_spec_manager
            for kid, ks in self.generate_kernel_specs(app, ipython_arguments).items():
                roots = [os.path.dirname(ks.resource_dir), ksm.user_kernel_dir]

                for root in roots:
                    kernel_dir = os.path.join(root, kid)
                    try:
                        if not os.path.exists(kernel_dir):
                            os.makedirs(kernel_dir)
                        with open(os.path.join(kernel_dir, "kernel.json"), "w") as f:
                            f.write(ks.to_json())
                        break
                    except OSError:
                        continue
                else:
                    raise CommandError(
                        "Could not write kernel %r in directories %r" % (kid, roots)
                    )

        app.start()

    @shell_runner(flags=["--notebook"], name="IPython Notebook")
    def get_notebook(self, options):
        try:
            from IPython import release
        except ImportError:
            return traceback.format_exc()
        try:
            from notebook.notebookapp import NotebookApp
        except ImportError:
            if release.version_info[0] >= 7:
                return traceback.format_exc()
            try:
                from IPython.html.notebookapp import NotebookApp
            except ImportError:
                if release.version_info[0] >= 3:
                    return traceback.format_exc()
                try:
                    from IPython.frontend.html.notebook import notebookapp

                    NotebookApp = notebookapp.NotebookApp
                except ImportError:
                    return traceback.format_exc()

        use_kernel_specs = release.version_info[0] >= 3

        def app_init(*args, **kwargs):
            app = NotebookApp.instance()
            app.initialize(*args, **kwargs)
            return app

        def run_notebook():
            self.run_notebookapp(app_init, options, use_kernel_specs)

        return run_notebook

    @shell_runner(flags=["--lab"], name="JupyterLab Notebook")
    def get_jupyterlab(self, options):
        try:
            from jupyterlab.labapp import LabApp
        except ImportError:
            return traceback.format_exc()

        # check for JupyterLab 3.0
        try:
            from notebook.notebookapp import NotebookApp
        except ImportError:
            NotebookApp = None

        if not NotebookApp or not issubclass(LabApp, NotebookApp):
            app_init = LabApp.initialize_server
        else:

            def app_init(*args, **kwargs):
                app = LabApp.instance()
                app.initialize(*args, **kwargs)
                return app

        def run_jupyterlab():
            self.run_notebookapp(app_init, options, history=False)

        return run_jupyterlab

    @shell_runner(flags=["--plain"], name="plain Python")
    def get_plain(self, options):
        # Using normal Python shell
        import code

        # Set up a dictionary to serve as the environment for the shell.
        imported_objects = self.get_imported_objects(options)

        use_pythonrc = options["use_pythonrc"]
        no_startup = options["no_startup"]

        # We want to honor both $PYTHONSTARTUP and .pythonrc.py, so follow system
        # conventions and get $PYTHONSTARTUP first then .pythonrc.py.
        if use_pythonrc or not no_startup:
            for pythonrc in OrderedSet(
                [os.environ.get("PYTHONSTARTUP"), os.path.expanduser("~/.pythonrc.py")]
            ):
                if not pythonrc:
                    continue
                if not os.path.isfile(pythonrc):
                    continue
                with open(pythonrc) as handle:
                    pythonrc_code = handle.read()
                # Match the behavior of the cpython shell where an error in
                # PYTHONSTARTUP prints an exception and continues.
                try:
                    exec(compile(pythonrc_code, pythonrc, "exec"), imported_objects)
                except Exception:
                    traceback.print_exc()
                    if self.tests_mode:
                        raise

        # By default, this will set up readline to do tab completion and to read and
        # write history to the .python_history file, but this can be overridden by
        # $PYTHONSTARTUP or ~/.pythonrc.py.
        try:
            hook = sys.__interactivehook__
        except AttributeError:
            # Match the behavior of the cpython shell where a missing
            # sys.__interactivehook__ is ignored.
            pass
        else:
            try:
                hook()
            except Exception:
                # Match the behavior of the cpython shell where an error in
                # sys.__interactivehook__ prints a warning and the exception
                # and continues.
                print("Failed calling sys.__interactivehook__")
                traceback.print_exc()

        try:
            # Try activating rlcompleter, because it's handy.
            import readline
        except ImportError:
            pass
        else:
            # We don't have to wrap the following import in a 'try', because
            # we already know 'readline' was imported successfully.
            import rlcompleter

            readline.set_completer(rlcompleter.Completer(imported_objects).complete)
            # Enable tab completion on systems using libedit (e.g. macOS).
            # These lines are copied from Lib/site.py on Python 3.4.
            readline_doc = getattr(readline, "__doc__", "")
            if readline_doc is not None and "libedit" in readline_doc:
                readline.parse_and_bind("bind ^I rl_complete")
            else:
                readline.parse_and_bind("tab:complete")

        def run_plain():
            code.interact(local=imported_objects)

        return run_plain

    @shell_runner(flags=["--bpython"], name="BPython")
    def get_bpython(self, options):
        try:
            from bpython import embed
        except ImportError:
            return traceback.format_exc()

        def run_bpython():
            imported_objects = self.get_imported_objects(options)
            kwargs = {}
            if self.extra_args:
                kwargs["args"] = self.extra_args
            embed(imported_objects, **kwargs)

        return run_bpython

    @shell_runner(flags=["--ipython"], name="IPython")
    def get_ipython(self, options):
        try:
            from IPython import start_ipython

            def run_ipython():
                imported_objects = self.get_imported_objects(options)
                ipython_arguments = self.extra_args or self.get_ipython_arguments(
                    options
                )
                start_ipython(argv=ipython_arguments, user_ns=imported_objects)

            return run_ipython
        except ImportError:
            str_exc = traceback.format_exc()
            # IPython < 0.11
            # Explicitly pass an empty list as arguments, because otherwise
            # IPython would use sys.argv from this script.
            # Notebook not supported for IPython < 0.11.
            try:
                from IPython.Shell import IPShell
            except ImportError:
                return str_exc + "\n" + traceback.format_exc()

            def run_ipython():
                imported_objects = self.get_imported_objects(options)
                shell = IPShell(argv=[], user_ns=imported_objects)
                shell.mainloop()

            return run_ipython

    @shell_runner(flags=["--ptpython"], name="PTPython")
    def get_ptpython(self, options):
        try:
            from ptpython.repl import embed, run_config
        except ImportError:
            tb = traceback.format_exc()
            try:  # prompt_toolkit < v0.27
                from prompt_toolkit.contrib.repl import embed, run_config
            except ImportError:
                return tb

        def run_ptpython():
            imported_objects = self.get_imported_objects(options)
            history_filename = os.path.expanduser("~/.ptpython_history")
            embed(
                globals=imported_objects,
                history_filename=history_filename,
                vi_mode=options["vi_mode"],
                configure=run_config,
            )

        return run_ptpython

    @shell_runner(flags=["--ptipython"], name="PT-IPython")
    def get_ptipython(self, options):
        try:
            from ptpython.repl import run_config
            from ptpython.ipython import embed
        except ImportError:
            tb = traceback.format_exc()
            try:  # prompt_toolkit < v0.27
                from prompt_toolkit.contrib.repl import run_config
                from prompt_toolkit.contrib.ipython import embed
            except ImportError:
                return tb

        def run_ptipython():
            imported_objects = self.get_imported_objects(options)
            history_filename = os.path.expanduser("~/.ptpython_history")
            embed(
                user_ns=imported_objects,
                history_filename=history_filename,
                vi_mode=options["vi_mode"],
                configure=run_config,
            )

        return run_ptipython

    @shell_runner(flags=["--idle"], name="Idle")
    def get_idle(self, options):
        from idlelib.pyshell import main

        def run_idle():
            sys.argv = [
                sys.argv[0],
                "-c",
                """
from django_extensions.management import shells
from django.core.management.color import no_style
for k, m in shells.import_objects({}, no_style()).items():
    globals()[k] = m
""",
            ]
            main()

        return run_idle

    def set_application_name(self, options):
        """
        Set the application_name on PostgreSQL connection

        Use the fallback_application_name to let the user override
        it with PGAPPNAME env variable

        https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-PARAMKEYWORDS
        """  # noqa: E501
        supported_backends = (
            "django.db.backends.postgresql",
            "django.db.backends.postgresql_psycopg2",
        )
        opt_name = "fallback_application_name"
        default_app_name = "django_shell"
        dbs = getattr(settings, "DATABASES", [])

        for connection in connections.all():
            alias = connection.alias
            mro = inspect.getmro(connection.__class__)
            if any(klass.__module__.startswith(supported_backends) for klass in mro):
                if "OPTIONS" not in dbs[alias] or opt_name not in dbs[alias]["OPTIONS"]:
                    dbs[alias].setdefault("OPTIONS", {}).update(
                        {opt_name: default_app_name}
                    )

    @signalcommand
    def handle(self, *args, **options):
        verbosity = options["verbosity"]
        get_runner = options["runner"]
        print_sql = getattr(settings, "SHELL_PLUS_PRINT_SQL", False)
        runner = None
        runner_name = None
        truncate = None if options["truncate_sql"] == 0 else options["truncate_sql"]

        with monkey_patch_cursordebugwrapper(
            print_sql=options["print_sql"] or print_sql,
            truncate=truncate,
            print_sql_location=options["print_sql_location"],
            confprefix="SHELL_PLUS",
        ):
            SETTINGS_SHELL_PLUS = getattr(settings, "SHELL_PLUS", None)

            def get_runner_by_flag(flag):
                for runner in self.runners:
                    if flag in runner.runner_flags:
                        return runner
                return None

            self.set_application_name(options)

            if not get_runner and SETTINGS_SHELL_PLUS:
                get_runner = get_runner_by_flag("--%s" % SETTINGS_SHELL_PLUS)
                if not get_runner:
                    runner = None
                    runner_name = SETTINGS_SHELL_PLUS

            if get_runner:
                runner = get_runner(options)
                runner_name = get_runner.runner_name
            else:

                def try_runner(get_runner):
                    runner_name = get_runner.runner_name
                    if verbosity > 2:
                        print(self.style.NOTICE("Trying: %s" % runner_name))

                    runner = get_runner(options)
                    if callable(runner):
                        if verbosity > 1:
                            print(self.style.NOTICE("Using: %s" % runner_name))
                        return runner
                    return None

                tried_runners = set()

                # try the runners that are least unexpected (normal shell runners)
                preferred_runners = [
                    "ptipython",
                    "ptpython",
                    "bpython",
                    "ipython",
                    "plain",
                ]
                for flag_suffix in preferred_runners:
                    get_runner = get_runner_by_flag("--%s" % flag_suffix)
                    tried_runners.add(get_runner)
                    runner = try_runner(get_runner)
                    if runner:
                        runner_name = get_runner.runner_name
                        break

                # try any remaining runners if needed
                if not runner:
                    for get_runner in self.runners:
                        if get_runner not in tried_runners:
                            runner = try_runner(get_runner)
                            if runner:
                                runner_name = get_runner.runner_name
                                break

            if not callable(runner):
                if runner:
                    print(runner)
                if not runner_name:
                    raise CommandError("No shell runner could be found.")
                raise CommandError("Could not load shell runner: '%s'." % runner_name)

            if self.tests_mode:
                return 130

            if options["command"]:
                imported_objects = self.get_imported_objects(options)
                exec(options["command"], imported_objects)
                return None

            runner()