# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

import json
import logging
import os
import subprocess
import sys
from vscode_datascience_helpers.daemon.daemon_python import (
    error_decorator,
    PythonDaemon as BasePythonDaemon,
    change_exec_context,
)


class PythonDaemon(BasePythonDaemon):
    def __init__(self, rx, tx):
        super().__init__(rx, tx)
        self.log.info("DataScience Daemon init")

    def __getitem__(self, item):
        """Override getitem to ensure we use these methods."""
        self.log.info("Execute rpc method %s in DS Daemon", item)
        return super().__getitem__(item)

    @error_decorator
    def m_exec_module(self, module_name, args=[], cwd=None, env=None):
        self.log.info("Exec in DS Daemon %s with args %s", module_name, args)
        args = [] if args is None else args

        if module_name == "jupyter":
            if args[0] == "kernelspec" and self._is_module_installed(
                "jupyter_client.kernelspec"
            ):
                if args == ["kernelspec", "list", "--json"]:
                    return self._execute_and_capture_output(
                        self._print_kernel_list_json
                    )
                elif args == ["kernelspec", "list"]:
                    return self._execute_and_capture_output(self._print_kernel_list)
                elif args == ["kernelspec", "--version"]:
                    return self._execute_and_capture_output(
                        self._print_kernelspec_version
                    )
            if (
                args[0] == "nbconvert"
                and self._is_module_installed("nbconvert")
                and args[-1] != "--version"
            ):
                return self._execute_and_capture_output(lambda: self._convert(args))
            if args[0] == "notebook" and args[1] == "--version":
                try:
                    from notebook import notebookapp as app

                    return {"stdout": ".".join(list(str(v) for v in app.version_info))}
                except Exception:
                    pass
            # kernelspec, nbconvert are subcommands of jupyter.
            # python -m jupyter kernelspec, python -m jupyter nbconvert,
            # In such cases, even if the modules kernelspec or nbconvert are not installed in the current
            # environment, jupyter will find them in current path.
            # So if we cannot find the corresponding subcommands, lets revert to subprocess.
            self.log.info(
                "Exec in DS Daemon with as subprocess, %s with args %s",
                module_name,
                args,
            )
            return self._exec_with_subprocess(module_name, args, cwd, env)
        else:
            self.log.info("check base class stuff")
            return super().m_exec_module(module_name, args, cwd, env)

    def _exec_with_subprocess(self, module_name, args=[], cwd=None, env=None):
        # # result = subprocess.run([sys.executable, "-m"] + args, stdout=sys.stdout, stderr=sys.stderr)
        # return self._execute_and_capture_output(lambda: subprocess.run([sys.executable, "-m", module_name] + args, stdout=sys.stdout, stderr=sys.stderr))
        result = subprocess.run(
            [sys.executable, "-m", module_name] + args, capture_output=True
        )
        encoding = os.getenv("PYTHONIOENCODING", "utf-8")
        stdout = result.stdout.decode(encoding)
        stderr = result.stderr.decode(encoding)
        self.log.info(
            "subprocess output for, %s with args %s, \nstdout is %s, \nstderr is %s",
            module_name,
            args,
            stdout,
            stderr,
        )
        return {"stdout": stdout, "stderr": stderr}

    @error_decorator
    def m_exec_module_observable(self, module_name, args=None, cwd=None, env=None):
        self.log.info(
            "Exec in DS Daemon (observable) %s with args %s", module_name, args
        )
        args = [] if args is None else args

        # Assumption is that `python -m jupyter notebook` or `python -m notebook` with observable output
        # will only ever be used to start a notebook and nothing else.
        # E.g. `python -m jupyter notebook --version` wouldn't require the use of exec_module_observable,
        # In such cases, we can get the output immediately.
        if (module_name == "jupyter" and args[0] == "notebook") or (
            module_name == "notebook"
        ):
            self._start_notebook(args, cwd, env)
        else:
            return super().m_exec_module_observable(module_name, args, cwd, env)

    def _print_kernelspec_version(self):
        import jupyter_client

        # Check whether kernelspec module exists.
        import jupyter_client.kernelspec

        sys.stdout.write(jupyter_client.__version__)
        sys.stdout.flush()

    def _print_kernel_list(self):
        self.log.info("listing kernels")
        # Get kernel specs.
        import jupyter_client.kernelspec

        specs = jupyter_client.kernelspec.find_kernel_specs()
        sys.stdout.write(
            os.linesep.join(list("{0} {1}".format(k, v) for k, v in specs.items()))
        )
        sys.stdout.flush()

    def _print_kernel_list_json(self):
        self.log.info("listing kernels as json")
        # Get kernel specs.
        import jupyter_client.kernelspec

        specs = jupyter_client.kernelspec.KernelSpecManager().get_all_specs()
        all_specs = {"kernelspecs": specs}
        sys.stdout.write(json.dumps(all_specs))
        sys.stdout.flush()

    def _convert(self, args):
        self.log.info("nbconvert")
        from nbconvert import nbconvertapp as app

        sys.argv = [""] + args
        app.main()

    def _start_notebook(self, args, cwd, env):
        from notebook import notebookapp as app

        # Args must not have ['notebook'] in the begining. Drop the `notebook` subcommand when using `jupyter`
        args = args[1:] if args[0] == "notebook" else args
        self.log.info("Starting notebook with args %s", args)

        # When launching notebook always ensure the first argument is `notebook`.
        with change_exec_context(args, cwd, env):
            app.launch_new_instance()