#! python

import asyncio
import argparse
import contextlib
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._aimport import eval_source, register
from concert.ext.cmd import plugins


LOG = logging.getLogger(__name__)


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:
                docstring = cs.get_docstring(session)
                print(docstring)
            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.
        register(paths=[cs.path()])

        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())
        path = filename or cs.path(session)

        if non_interactive:
            with open(path, "rb") as f:
                eval_source(f.read(), {}, filename=path)
        else:
            self.run_shell(path=path, session=session)

    def run_shell(self, path=None, session=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__))

        ip_config = traitlets.config.Config()
        if path and path.endswith('.py'):
            docstring = cs.get_docstring(path)
            if docstring:
                print(docstring)

            if session is None:
                # --filename must have been specified
                session = os.path.splitext(os.path.basename(path))[0]
                sys.path.append(os.path.dirname(path))

            ip_config.InteractiveShellApp.exec_lines = [f'from {session} import *']
        else:
            session_code = 'from concert.quantities import q'
            ip_config.InteractiveShellApp.exec_lines = [session_code]

        ip_config.InteractiveShellApp.gui = 'asyncio'
        # 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)
        docstring = cs.get_docstring(session)

        if not docstring:
            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(docstring.encode('utf-8'))


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()
