from IPython.core.inputtransformer2 import (
    EscapedCommand,
    MagicAssign,
    SystemAssign,
    find_end_of_continued_line,
    assemble_continued_line
)

import re

ESC_SHELL  = '!'     # Send line to underlying system shell
ESC_SH_CAP = '!!'    # Send line to system shell and capture output
ESC_HELP   = '?'     # Find information about object
ESC_HELP2  = '??'    # Find extra-detailed information about object
ESC_MAGIC  = '%'     # Call magic function
ESC_MAGIC2 = '%%'    # Call cell-magic function
ESC_QUOTE  = ','     # Split args on whitespace, quote each as string and call
ESC_QUOTE2 = ';'     # Quote all args as a single string, call
ESC_PAREN  = '/'     # Call first argument with rest of line as arguments

ESCAPE_SINGLES = {'!', '?', '%', ',', ';', '/'}
ESCAPE_DOUBLES = {'!!', '??'}  # %% (cell magic) is handled separately

def _make_help_call(target, esc):
    """Prepares a pinfo(2)/psearch call from a target name and the escape
    (i.e. ? or ??)"""
    method  = 'pinfo2' if esc == '??' \
                else 'psearch' if '*' in target \
                else 'pinfo'
    arg = " ".join([method, target])
    #Prepare arguments for get_ipython().run_line_magic(magic_name, magic_args)
    t_magic_name, _, t_magic_arg_s = arg.partition(' ')
    t_magic_name = t_magic_name.lstrip(ESC_MAGIC)
    return "get_ipython().run_line_magic({}, {})".format(crepr(t_magic_name), crepr(t_magic_arg_s))

def crepr(s):
    assert isinstance(s, str)
    return "\"" + repr(s)[1:-1] + "\""


def _tr_help(content):
    """Translate lines escaped with: ?

    A naked help line should fire the intro help screen (shell.show_usage())
    """
    if not content:
        return 'get_ipython().show_usage()'

    return _make_help_call(content, '?')

def _tr_help2(content):
    """Translate lines escaped with: ??

    A naked help line should fire the intro help screen (shell.show_usage())
    """
    if not content:
        return 'get_ipython().show_usage()'

    return _make_help_call(content, '??')

def _tr_magic(content):
    "Translate lines escaped with a percent sign: %"
    name, _, args = content.partition(' ')
    return 'get_ipython().run_line_magic({}, {})'.format(crepr(name), crepr(args))

def _tr_quote(content):
    "Translate lines escaped with a comma: ,"
    name, _, args = content.partition(' ')
    return '%s("%s")' % (name, '", "'.join(args.split()) )

def _tr_quote2(content):
    "Translate lines escaped with a semicolon: ;"
    name, _, args = content.partition(' ')
    return '%s("%s")' % (name, args)

def _tr_paren(content):
    "Translate lines escaped with a slash: /"
    name, _, args = content.partition(" ")
    if name == "":
        raise SyntaxError(f'"{ESC_SHELL}" must be followed by a callable name')

    return '%s(%s)' % (name, ", ".join(args.split()))

tr = { ESC_SHELL  : lambda s: 'get_ipython().system({})'.format(crepr(s)),
       ESC_SH_CAP : lambda s: 'get_ipython().getoutput({})'.format(crepr(s)),
       ESC_HELP   : _tr_help,
       ESC_HELP2  : _tr_help2,
       ESC_MAGIC  : _tr_magic,
       ESC_QUOTE  : _tr_quote,
       ESC_QUOTE2 : _tr_quote2,
       ESC_PAREN  : _tr_paren }


class CoatlEscapedCommand(EscapedCommand):
    def transform(self, lines):
        """Transform an escaped line found by the ``find()`` classmethod.
        """
        start_line, start_col = self.start_line, self.start_col

        indent = lines[start_line][:start_col]
        end_line = find_end_of_continued_line(lines, start_line)
        line = assemble_continued_line(lines, (start_line, start_col), end_line)

        if len(line) > 1 and line[:2] in ESCAPE_DOUBLES:
            escape, content = line[:2], line[2:]
        else:
            escape, content = line[:1], line[1:]

        if escape in tr:
            call = tr[escape](content)
        else:
            call = ''

        lines_before = lines[:start_line]
        new_line = indent + call + '\n'
        lines_after = lines[end_line + 1:]

        return lines_before + [new_line] + lines_after


class CoatlMagicAssign(MagicAssign):
    priority = 9

    def transform(self, lines):
        start_line, start_col = self.start_line, self.start_col
        lhs = lines[start_line][:start_col]
        end_line = find_end_of_continued_line(lines, start_line)
        rhs = assemble_continued_line(lines, (start_line, start_col), end_line)
        assert rhs.startswith('%'), rhs
        magic_name, _, args = rhs[1:].partition(' ')

        lines_before = lines[:start_line]
        call = "get_ipython().run_line_magic({}, {})".format(crepr(magic_name), crepr(args))
        new_line = lhs + call + '\n'
        lines_after = lines[end_line+1:]

        return lines_before + [new_line] + lines_after


class CoatlSystemAssign(SystemAssign):
    priority = 9

    def transform(self, lines):
        start_line, start_col = self.start_line, self.start_col

        lhs = lines[start_line][:start_col]
        end_line = find_end_of_continued_line(lines, start_line)
        rhs = assemble_continued_line(lines, (start_line, start_col), end_line)
        assert rhs.startswith('!'), rhs
        cmd = rhs[1:]

        lines_before = lines[:start_line]
        call = "get_ipython().getoutput({})".format(crepr(cmd))
        new_line = lhs + call + '\n'
        lines_after = lines[end_line + 1:]

        return lines_before + [new_line] + lines_after


def coatl_cell_magic(lines):
    if not lines or not lines[0].startswith('%%'):
        return lines
    if re.match(r'%%\w+\?', lines[0]):
        # This case will be handled by help_end
        return lines
    magic_name, _, first_line = lines[0][2:].rstrip().partition(' ')
    body = ''.join(lines[1:])
    return ['get_ipython().run_cell_magic({}, {}, {})\n'
            .format(crepr(magic_name), crepr(first_line), crepr(body))]