#! python

import ast
import asyncio
import argparse
import contextlib
import importlib
import inspect
import logging
import os
import re
import sys
import shutil
import signal
import subprocess
import tempfile
import zipfile
import concert
import concert.config
import concert.session.management as cs
from concert.ext.cmd import plugins


LOG = logging.getLogger(__name__)


# Start of enabling importing python files with "await" outside of "async def" functions


class AsyncLoader(importlib.machinery.SourceFileLoader):
    """
    Normal SourceFileLoader with the capability of dealing with `await' outside of `async def'
    functions. The trick is the flag ast.PyCF_ALLOW_TOP_LEVEL_AWAIT and eval() instead of exec()
    used by Python's :class:`_LoaderBasics` class.

    If the module has `await' outside of a function, eval() will return a coroutine which we execute
    in the session's loop, that is why sys.meta_path must be updated *after* the loop has been
    created (to make sure it's the one IPython also uses).
    """
    def exec_module(self, module):
        cobj = compile(
            self.get_data(self.path),
            self.path,
            'exec',
            dont_inherit=True,
            flags=ast.PyCF_ALLOW_TOP_LEVEL_AWAIT
        )
        result = eval(cobj, module.__dict__)
        if inspect.iscoroutine(result):
            loop = asyncio.get_event_loop()
            try:
                loop.run_until_complete(result)
            except RuntimeError as err:
                raise StartError(
                    f"Module `{module.__name__}' uses `await' outside of `async def' "
                    "functions but at the same time tries to execute coroutines in a loop "
                    "which is not allowed, please remove the loop-running code from the module "
                    "and try again"
                ) from err


class AsyncMetaPathFinder(importlib.abc.MetaPathFinder):
    """
    MetaPathFinder implementation which used :class;`.FileFinder` for dealing with paths and
    instantiated with our AsyncLoader for the handling of `await' outside of `async def' functions.
    """
    def find_spec(self, fullname, path, target=None):
        if path is None:
            path = sys.path

        for entry in path:
            finder = importlib.machinery.FileFinder(entry, (AsyncLoader, ('.py',)))
            spec = finder.find_spec(fullname, target=target)
            if spec:
                return spec


class StartError(Exception):
    """Startup errors."""


# End of enabling importing python files with "await" outside of "async def" functions


def docstring_summary(doc):
    if doc and doc.find('.'):
        return doc[:doc.find('.')]
    return doc


def cmp_versions(v1, v2):
    """Compare two version numbers and return cmp compatible result"""
    def normalize(v):
        return [int(x) for x in re.sub(r'(\.0+)*$', '', v).split(".")]

    n1 = normalize(v1)
    n2 = normalize(v2)
    return (n1 > n2) - (n1 < n2)


def get_prompt_config(path):
    """Customize the prompt for session in *path*."""
    from IPython.terminal.prompts import Prompts, Token

    session_name = os.path.splitext(os.path.basename(path))[0]

    class MyPrompt(Prompts):
        def in_prompt_tokens(self, cli=None):
            return [(Token, session_name), (Token.Prompt, ' > ')]

        def out_prompt_tokens(self):
            return []

    return MyPrompt


class SubCommand(object):

    """Base sub-command class (concert [subcommand])."""

    def __init__(self, name, opts):
        """
        SubCommand objects are loaded at run-time and injected into Concert's
        command parser.

        *name* denotes the name of the sub-command parser, e.g. "mv" for the
        MoveCommand. *opts* must be an argparse-compatible dictionary
        of command options.
        """
        self.name = name
        self.opts = opts

    def run(self, *args, **kwargs):
        """Run the command"""
        raise NotImplementedError


class InitCommand(SubCommand):

    """Create a new session."""

    def __init__(self):
        opts = {'session': {'type': str},
                '--force': {'action': 'store_true',
                            'help': "Overwrite existing sessions"},
                '--imports': {'help': "Pre-import processes",
                              'metavar': 'modules',
                              'default': ''}}
        super(InitCommand, self).__init__('init', opts)

    def run(self, session=None, imports="", force=False):
        if cs.exists(session) and not force:
            message = "Session `{0}' already exists."
            message += " Use --force to create it anyway."
            print(message.format(session))
        else:
            cs.create(session, imports.split())


class EditCommand(SubCommand):

    """Edit a session."""

    def __init__(self):
        opts = {'session': {'type': str}}
        super(EditCommand, self).__init__('edit', opts)

    def run(self, session=None):
        if not cs.exists(session):
            print("Session not found, creating {}.".format(session))
            InitCommand().run(session)

        env = os.environ
        editor = env['EDITOR'] if 'EDITOR' in env else 'vi'
        subprocess.call([editor, cs.path(session)])


class LogCommand(SubCommand):

    """Show session logs."""

    def __init__(self):
        opts = {'session': {'type': str,
                            'nargs': '?'},
                '--follow': {'action': 'store_true',
                             'help': 'Show current log'}}
        super(LogCommand, self).__init__('log', opts)

    def run(self, session=None, follow=False):
        logfile = cs.logfile_path()

        if not os.path.exists(logfile):
            return

        # This is danger zone here because we run subprocess.call with
        # shell=True.  However, the only input that we input is args.session
        # which we check first and the logfile itself.

        if session:
            cs.exit_if_not_exists(session)

            if follow:
                cmd = 'tail -f {} | grep --line-buffered "{}:"'.format(logfile, session)
            else:
                cmd = 'grep "{0}:" {1} | less'.format(session, logfile)
        else:
            if follow:
                cmd = 'tail -f {}'.format(logfile)
            else:
                cmd = 'less {}'.format(logfile)

        try:
            subprocess.call(cmd, shell=True)
        except KeyboardInterrupt:
            # When following we can only leave tail by C-c, hence to avoid
            # spamming the terminal with a stack trace we just ignore the
            # Keyboardinterrupt exception.
            pass


class ShowCommand(SubCommand):

    """Show available sessions or details of a given *session*."""

    def __init__(self):
        opts = {'session': {'type': str,
                            'nargs': '?',
                            'default': None,
                            'help': "Show details"}}
        super(ShowCommand, self).__init__('show', opts)

    def run(self, session=None):
        if session:
            try:
                module = cs.load(session)
                print(module.__doc__)
            except IOError:
                print("Cannot find {0}".format(session))
            except ImportError as exception:
                print("Cannot import {0}: {1}".format(session, exception))
        else:
            sessions = cs.get_existing()
            print("Available sessions:")

            for session in sessions:
                print("  %s" % session)


class MoveCommand(SubCommand):

    """Move session *source* to *target*."""

    def __init__(self):
        opts = {'source': {'type': str,
                           'help': "Name of the source session"},
                'target': {'type': str,
                           'help': "Name of the target session"}}
        super(MoveCommand, self).__init__('mv', opts)

    def run(self, source, target):
        if not cs.exists(source):
            sys.exit("`{}' does not exist".format(source))

        if cs.exists(target):
            sys.exit("`{}' already exists".format(target))

        cs.move(source, target)
        print("Renamed {} -> {}".format(source, target))


class CopyCommand(SubCommand):

    """Copy session *source* to *target*."""

    def __init__(self):
        opts = {'source': {'type': str,
                           'help': "Name of the source session"},
                'target': {'type': str,
                           'help': "Name of the target session"}}
        super(CopyCommand, self).__init__('cp', opts)

    def run(self, source, target):
        if not cs.exists(source):
            sys.exit("`{}' does not exist".format(source))

        if cs.exists(target):
            sys.exit("`{}' already exists".format(target))

        cs.copy(source, target)
        print("Copied {} -> {}".format(source, target))


class RemoveCommand(SubCommand):

    """Remove one or more sessions."""

    def __init__(self):
        opts = {'sessions': {'type': str,
                             'nargs': '+',
                             'metavar': 'session'}}
        super(RemoveCommand, self).__init__('rm', opts)

    def run(self, sessions=[]):
        for session in sessions:
            print("Removing {0}...".format(session))
            cs.remove(session)


class ImportCommand(SubCommand):

    """Import an existing *session*."""

    def __init__(self):
        opts = {'url': {'nargs': '+', 'type': str,
                        'help': "Import a Python module and save as a session."
                        " Note: Server certificates of HTTPS requests"
                        " are NOT verified!"},
                '--force': {'action': 'store_true',
                            'help': "Overwrite existing sessions"},
                '--repo': {'action': 'store_true',
                           'help':
                           "Checkout Git repository and import all files"}}
        super(ImportCommand, self).__init__('import', opts)

    def run(self, url, force=False, repo=False):
        for u in url:
            if repo:
                self._import_repo(u, force)
            else:
                self._import_file(u, force)

    def _import_repo(self, url, force):
        path = tempfile.mkdtemp()
        cmd = 'git clone --quiet {0} {1}'.format(url, path)
        proc = subprocess.Popen(cmd.split(), stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE)
        out, err = proc.communicate()

        if proc.returncode != 0:
            sys.exit("Could not clone {0}.".format(url))

        for filename in (x for x in os.listdir(path) if x.endswith('.py')):
            session_name = os.path.basename(filename[:-3])

            if cs.exists(session_name) and not force:
                print("`{0}' already exists (use --force to install"
                      " anyway)".format(session_name))
            else:
                print("Add session {0} ...".format(filename[:-3]))
                shutil.copy(os.path.join(path, filename),
                            cs.path())

        shutil.rmtree(path)

    def _import_file(self, url, force):
        import urllib.request, urllib.error, urllib.parse

        if not url.endswith('.py'):
            sys.exit("`{0}' is not a Python module".format(url))

        session_name = os.path.basename(url[:-3])

        if cs.exists(session_name) and not force:
            sys.exit("`{0}' already exists".format(session_name))

        print("Add session {0} ...".format(session_name))
        local_url = self._get_url(url)

        with contextlib.closing(urllib.request.urlopen(local_url)) as data:
            with open(cs.path(session_name), 'w') as output:
                output.write(data.read())

    def _get_url(self, path_or_url):
        import urllib.parse

        result = urllib.parse.urlsplit(path_or_url)

        if result.scheme:
            return path_or_url

        if not os.path.exists(path_or_url):
            sys.exit("Cannot find module `{0}'.".format(path_or_url))

        result = ('file', '', os.path.abspath(path_or_url), '', '')
        return urllib.parse.urlunsplit(result)


class ExportCommand(SubCommand):

    """Export all sessions as a Zip archive."""

    def __init__(self):
        opts = {'name': {'type': str,
                         'help': "Name of the archive"}}
        super(ExportCommand, self).__init__('export', opts)

    def run(self, name):
        name = name if name.endswith('.zip') else name + '.zip'

        with zipfile.ZipFile(name, 'w') as archive:
            for path in (cs.path(session) for session in cs.get_existing()):
                archive.writestr(os.path.basename(path), open(path).read())


class StartCommand(SubCommand):

    """Start a session."""

    def __init__(self):
        opts = {'session': {'nargs': '?', 'type': str, 'default': None},
                '--filename':  {'type': str, 'default': None},
                '--logto': {'choices': ['stderr', 'file'],
                            'default': 'file'},
                '--logfile': {'type': str},
                '--loglevel': {'choices': ['perfdebug', 'aiodebug', 'debug', 'info', 'warning',
                                           'error', 'critical'],
                               'default': 'info'},
                '--non-interactive': {'action': 'store_true'}}
        super(StartCommand, self).__init__('start', opts)

    def run(self, session=None, filename=None,
            non_interactive=False,
            logto='file', logfile=None, loglevel=None):
        import IPython

        if IPython.version_info >= (8, 0):
            # Ipython creates a new event loop in get_asyncio_loop function, but only sets it in
            # the prompt_for_code of the
            # IPython.terminal.interactiveshell.TerminalInteractiveShell, so when we use
            # enabel_gui programatically, that loop setting will never be triggered. Thus, we do
            # it ourselves before anything else with any concert part happens.
            from IPython.core.async_helpers import get_asyncio_loop
            asyncio.set_event_loop(get_asyncio_loop())

        # From now on, we can import code with `await' outside of `async def' functions
        # It is crucial to call this here, after setting the asyncio's loop to the one from IPython,
        # otherwise we might be importing with another event loop than the one used in the concert
        # session and thus have devices and Locks and other things bound to a wrong loop.
        sys.meta_path.insert(0, AsyncMetaPathFinder())

        if session:
            cs.exit_if_not_exists(session)

        if logto == 'file':
            logfilename = logfile if logfile else cs.logfile_path()
            handler = logging.FileHandler(logfilename)
        else:
            handler = logging.StreamHandler(sys.stderr)

        logging.addLevelName(concert.config.PERFDEBUG, 'PERFDEBUG')
        logging.addLevelName(concert.config.AIODEBUG, 'AIODEBUG')
        handler.setLevel(loglevel.upper())
        logformat = '[%(asctime)s] %(levelname)s: %(name)s: {}: %(message)s'
        formatter = logging.Formatter(logformat.format(session))
        handler.setFormatter(formatter)

        self.logger = logging.getLogger()
        self.logger.addHandler(handler)
        self.logger.setLevel(loglevel.upper())

        # Add session path, so that sessions can import other sessions
        sys.path.append(cs.path())

        if non_interactive:

            if session:
                code_obj = compile(
                    open(cs.path(session), "rb").read(),
                    cs.path(session),
                    'exec',
                    flags=ast.PyCF_ALLOW_TOP_LEVEL_AWAIT
                )
                eval_result = eval(code_obj, globals())
                if inspect.iscoroutine(eval_result):
                    loop = asyncio.get_event_loop()
                    # More-or less a copy of asyncio.run() but with our loop
                    try:
                        loop.run_until_complete(eval_result)
                    finally:
                        to_cancel = asyncio.all_tasks(loop=loop)
                        for task in to_cancel:
                            task.cancel()
                        loop.run_until_complete(asyncio.gather(*to_cancel, return_exceptions=True))
                        loop.run_until_complete(loop.shutdown_asyncgens())
                        if hasattr(loop, 'shutdown_default_executor'):
                            # Only in Python 3.9+
                            loop.run_until_complete(loop.shutdown_default_executor())
        else:
            self.run_shell(path=filename or cs.path(session) if session else None)

    def run_shell(self, path=None):
        import IPython
        import traitlets.config
        from concert.session.utils import abort_awaiting

        # ctrl-c handling (abort whatever what is running in the foreground (blocks))
        orig_sigint_handler = signal.getsignal(signal.SIGINT)

        # Cancel running tasks in IPython's asyncio event loop
        def concert_sigint_handler(sig, frame):
            LOG.info('KeyboardInterrupt in awaiting mode')
            aborted = abort_awaiting()
            LOG.debug('Aborted: %s', aborted)
            if not aborted:
                # If there is a task being awaited cancel it and swallow the exception so that we
                # don't flood the terminal with the stack trace, otherwise do the usual business
                # TODO: we may still keep hanging if the task being cancelled calls a blocking
                # function not in executor.
                orig_sigint_handler(sig, frame)

        signal.signal(signal.SIGINT, concert_sigint_handler)

        print("Welcome to Concert {0}".format(concert.__version__))

        if path:
            with open(path) as session_file:
                session_code = session_file.read()
            tree = ast.parse(session_code)
            docstring = ast.get_docstring(tree)
            if docstring:
                print(docstring)
        else:
            session_code = 'from concert.quantities import q'

        ip_config = traitlets.config.Config()
        ip_config.InteractiveShellApp.gui = 'asyncio'
        ip_config.InteractiveShellApp.exec_lines = [session_code]
        # This is the most robust way when taking virtualenv into account I have found so far
        ip_config.InteractiveShellApp.exec_files = [os.path.join(concert.__path__[0],
                                                                 '_ipython_setup.py')]
        ip_config.TerminalInteractiveShell.prompts_class = get_prompt_config(path or 'concert')
        ip_config.TerminalInteractiveShell.loop_runner = 'asyncio'
        ip_config.TerminalInteractiveShell.autoawait = True
        ip_config.InteractiveShell.confirm_exit = False
        ip_config.TerminalIPythonApp.display_banner = False
        ip_config.Completer.use_jedi = False

        # Now we start ipython with our configuration
        IPython.start_ipython(argv=[], config=ip_config)


class DocsCommand(SubCommand):

    """Create documentation of *session* docstring."""

    def __init__(self):
        opts = {'session': {'type': str, 'metavar': 'session'}}
        super(DocsCommand, self).__init__('docs', opts)

    def run(self, session):
        import subprocess
        import shlex

        try:
            subprocess.check_output(['pandoc', '-v'])
        except OSError:
            print("Please install pandoc and pdftex to generate docs.")
            sys.exit(1)

        cs.exit_if_not_exists(session)
        module = cs.load(session)

        if not module.__doc__:
            print("No docstring in `{}' found".format(session))

        cmd_line = shlex.split('pandoc -f markdown -t latex -o {}.pdf'.format(session))
        pandoc = subprocess.Popen(cmd_line, stdin=subprocess.PIPE)
        pandoc.communicate(module.__doc__)


def main():
    parser = argparse.ArgumentParser()

    parser.add_argument('--version',
                        action='version',
                        version="Concert v%s " % concert.__version__)

    subparsers = parser.add_subparsers(title="Concert commands",
                                       metavar="")

    commands = [InitCommand(),
                EditCommand(),
                LogCommand(),
                ShowCommand(),
                MoveCommand(),
                CopyCommand(),
                RemoveCommand(),
                ImportCommand(),
                ExportCommand(),
                StartCommand(),
                DocsCommand()]

    commands.extend(plugins)

    for command in commands:
        summary = docstring_summary(command.__doc__)
        cmd_parser = subparsers.add_parser(command.name, help=summary)
        cmd_parser.set_defaults(func=command.run)

        for arg in list(command.opts.keys()):
            cmd_parser.add_argument(arg, **command.opts[arg])

    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(0)

    args = parser.parse_args()
    func = args.func
    del args.func
    func(**vars(args))


if __name__ == '__main__':
    # This must be here for mp.get_context('spawn') in concert.ext.viewers to work
    __spec__ = None
    main()
