# -*- coding: utf-8 -*-
import os
from random import random
import base64

from pyparsing import nestedExpr, Optional, Word, alphanums, alphas,\
                      originalTextFor, Literal, SkipTo, Empty, Or, ZeroOrMore, \
                      delimitedList  # ,restOfLine
from pyparsing import _MAX_INT as pyparsing_MAX_INT

from .simplekernel import SimpleKernel


def skipToMatching(opener, closer):
    """

    :param opener: opening token
    :param closer: closing token

    """
    # https://github.com/sagemath/sagetex/issues/6#issuecomment-734968972
    nest = nestedExpr(opener, closer)
    return originalTextFor(nest)


class TexElement(object):
    """
    :param src: ancestral tex document, to which istart and iend refer
    :param istart: opening index
    :param iend: closing index
    :param parent: TexElement inmmediately containing this
    """

    # iend = None since src[slice(0, None)] is the full src
    # curiousity: mystring[0:None] gives an error
    def __init__(self, src=None, istart=0, iend=None, parent=None, *args, **kwds):
        self.src = src
        self.istart = istart
        self.iend = iend
        self.parent = parent

    def __str__(self):
        return self.src[slice(self.istart, self.iend)]


class TexSource(TexElement):
    """Some bulk tex code, as opposed to environments or commands.
    The whole document or the content of an environment are examples.
    """

    def __init__(self, *args, **kwds):
        super(TexSource, self).__init__(*args, **kwds)


class Arguments(object):
    """docstring for Arguments."""

    def __init__(self, *args, **kwds):
        self.args = args + tuple(kwds.keys())
        self.named_args = {}
        for k, v in kwds.items():
            self.named_args[k] = v

    def __getitem__(self, k):
        return self.args[k]

    def __getattr__(self, attr):
        if (attr == 'named_args') or attr not in self.named_args:
            return super(Arguments, self).__getattr__(attr)
        return self.named_args[attr]


class EnvOrCommand(TexElement):
    """
    :param name: name of the Environment or Command
    :param optional_args: a list of TexElement
    """

    def __init__(self, name='', options=None, arguments=None, *args, **kwds):
        ''''''
        super(EnvOrCommand, self).__init__(*args, **kwds)
        self.name = name
        self.options = options or Arguments()
        self.arguments = arguments or Arguments()

    def __repr__(self):
        return repr(self.name)

    @staticmethod
    def newEnvOrCommand(name, *args, **kwds):
        if name[0] == '\\':
            return Command(name, *args, **kwds)
        return Environment(name, *args, **kwds)


class Command(EnvOrCommand):
    """docstring for Command."""

    def __init__(self, *args, **kwds):
        super(Command, self).__init__(*args, **kwds)
        # When it comes to searching nested commands (which is what icontent is about),
        # a command starts with '\commandname{...' (length is 2+lname)
        # finishes with '...}'
        # and 'content' is everything in between
        lname = len(self.name)
        self.icontent = self.istart + 2 + lname

    def __repr__(self):
        return "'" + '\\' + self.name + "'"

    def long_repr(self):
        if self.arguments.named_args:
            return repr((self.name, self.arguments.named_args))
        elif self.arguments.args:
            return repr((self.name, self.arguments.args))
        else:
            return repr(self.name)


class Environment(EnvOrCommand):
    """docstring for Environment."""

    def __init__(self, *args, **kwds):
        super(Environment, self).__init__(*args, **kwds)
        lname = len(self.name)
        # The environment starts with \begin{envname} (length is 8+lname)
        # finishes with \end{envname}  (length is 6+lname)
        # content is everything in between
        self.icontent = self.istart + 8 + lname
        self.content = kwds.get('content', '')
        if not self.content and self.iend is not None:
            # TODO: lazy property ?
            self.content = self.src[slice(self.icontent, self.iend - (6 + lname))]

    def long_repr(self):
        if self.arguments.named_args:
            return repr((self.name, self.arguments.named_args, self.content))
        return repr((self.name, self.content))


class Match(tuple):
    """A call to findall() return a ResultSet of Matches

    parent : TexElement
        first element of the selector
    children: ResultSet of Match
        matches for the rest of the selector that are nested within parent

    displays as a tuple, if possible compatible with current unit tests
    """
    def __new__(cls, parent, children=None):
        return super(Match, cls).__new__(cls, (parent, children))

    def __init__(self, parent, children=None):
        self.parent = parent
        self.children = children

    def flatten(self):
        '''return a FlatMatch, keeping only the first child, and also recursively
        '''
        elements = self._recursive_first_child()
        return FlatMatch(elements)

    def _recursive_first_child(self):
        return (self.parent,) + (self.children[0]._recursive_first_child()
                                 if self.children else ())

    def __repr__(self):
        if self.children:
            return super(Match, self).__repr__()
#            return repr((self.parent, self.children))
        elif isinstance(self.parent, EnvOrCommand):
            return self.parent.long_repr()
        else:
            return repr(self.parent)


class FlatMatch(tuple):
    """A call to find() returns a FlatMatch, which contains the elements that correspond
    to each part of the selecto.

    parent : TexElement
        first element of the selector
    child: FlatMatch
        match for the rest of the selector that is nested within parent
    """
    def __new__(cls, elements):
        return super(FlatMatch, cls).__new__(cls, elements)

    def __repr__(self):
        if len(self)>1:
            return ('('
                    + ', '.join(repr(o) for o in self[:len(self)-1])
                    + ', '
                    + self[-1].long_repr()
                    + ')')
        else:
            return self[0].long_repr()


class ResultSet(list):
    """A ResultSet is just a list that keeps track of the TexSurgery
    that created it."""
    # Borrowed from BeautifulSoup
    # http://www.crummy.com/software/BeautifulSoup/

    def __init__(self, ts, contentlist):
        super(ResultSet, self).__init__(contentlist)
        self.ts = ts

    def __getattr__(self, key):
        """Raise a helpful exception to explain a common code fix."""
        raise AttributeError(
            "ResultSet object has no attribute '%s'. You're probably treating a list of elements "
            "like a single element. Did you call findall() when you meant to call find()?" % key
        )


class TexSurgery(TexElement):
    """TexSurgery allows to make some replacements in LaTeX code"""

    # TODO: logging level
    def __init__(self, tex_source, path='.', verbose=True):
        super(TexSurgery, self).__init__(tex_source)
        self.original_src = tex_source
        self.src = tex_source
        self.path = path
        # self.kernel is a lazy property
        self._kernels = dict()
        self.kernel_names = []
        self._auxfiles = 0
        # A random number to distinguish different concurrent jobs
        self._id = str(random())[2:]
        # self.codeparser is a lazy property
        self._codeparser = None
        self.verbose = verbose

    def __del__(self):
        """
        Destructor. Shuts down kernel safely.
        """
        self.shutdown()

    def shutdown(self):
        if self._kernels:
            for kernel in self._kernels.values():
                kernel.kernel_manager.shutdown_kernel()
            self._kernels = dict()

    @property
    def kernels(self):
        if not self._kernels:
            self._kernels = {
                kernelname:SimpleKernel(kernelname, verbose=self.verbose)
                for kernelname in self.kernel_names}
        return self._kernels

    @property
    def codeparser(self):
        if not self._codeparser:
            self._build_codeparser()
        return self._codeparser

    def _add_import_action(self, packagename, options):
        def action(l, s, t):
            return '\\documentclass' + t.restofline + '\n\\usepackage%s{%s}' % (
                '[%s]' % options if options else '',
                packagename
            )
        return action

    def add_import(self, packagename, options=''):
        documentclass = (
            '\\documentclass' + SkipTo('\n')('restofline')
        )
        documentclass.setParseAction(
            self._add_import_action(packagename, options)
        )
        self.src = documentclass.transformString(self.src)
        return self

    def data_surgery(self, replacements):
        # TODO: use pyparsing instead of regex, for the sake of uniformity
        src = self.src
        import re
        revars = re.compile('|'.join(r'\\'+key for key in replacements))
        pos, pieces = 0, []
        m = revars.search(src)
        while m:
            start, end = m.span()
            pieces.append(src[pos:start])
            # start+1, since the backslash \ is not part of the key
            name = src[start+1:end]
            pieces.append(replacements[name])
            pos = end
            m = revars.search(src, pos=pos)
        pieces.append(src[pos:])
        self.src = ''.join(map(str, pieces))
        return self

    def _latexify(self, results):
        # TODO do something special with 'text/html'?
        # TODO error -> texttt
        result = ''
        for r in results:
            hasimage = r.get('image/png')
            if hasimage:
                images_folder = 'images'
                images_path = os.path.join(self.path, images_folder)
                filename = 'texsurgery_image_{}_{}.png'.format(
                    self._id, self._auxfiles
                )
                fullpath = os.path.join(images_path, filename)
                if not os.path.exists(images_path):
                    os.mkdir(images_path)
                with open(fullpath, 'wb') as fd:
                    fd.write(base64.b64decode(hasimage))
                result += '\n\\includegraphics{%s}\n' % os.path.join(images_folder, filename)
                self._auxfiles += 1
            else:
                if r.get('text/latex'):
                    result += r.get('text/latex')[1:-1]
                else:
                    result += r.get('text/plain') or r.get('text/html') or r.get('error')
        return result

    def _select_kernel(self, t):
        if 'options' in t:
            kernel = self.kernels[t['options']]
        else:
            kernel = self.kernels[self.kernel_names[0]]
        return kernel

    def _runsilent(self, l, s, t):
        self._select_kernel(t).executesilent(t.content)
        return ''

    def _run(self, l, s, t):
        return self._latexify(self._select_kernel(t).execute(t.content, allow_errors=True))

    def _eval(self, l, s, t):
        # TODO: optional format code here
        code = t.code[1:-1]
        return self._latexify(self._select_kernel(t).execute(code))

    def _srepl(self, l, s, t):
        r"""
        Use for a block of code that should reflect what happens in an interactive sage session
        (both input and output)
        """
        code = t.content.strip()
        lines = code.split('\n')+['']
        if not lines:
            return ''
        result = '\\begin{verbatim}\n'
        partialblock = lines.pop(0) + '\n'
        result += 'sage: ' + partialblock
        while lines:
            line = lines.pop(0)
            if line and line[0] in [' ', '\t']:
                partialblock += line + '\n'
                result += '....: ' + line + '\n'
            else:
                answer = self._latexify(self.kernels['sagemath'].execute(partialblock))
                if len(answer)>0:
                    result += answer + '\n'
                partialblock = line + '\n'
                if line:
                    result += 'sage: ' + line + '\n'
        return result + '\\end{verbatim}'

    def _latex_escape(self, text):
        """
            :param text: a plain text message
            :return: the message escaped to appear correctly in LaTeX
        """
        # TODO: use pyparsing, not regex
        import re
        conv = {
            '&': r'\&',
            '%': r'\%',
            '$': r'\$',
            '#': r'\#',
            '_': r'\_',
            '{': r'\{',
            '}': r'\}',
            '~': r'\textasciitilde{}',
            '^': r'\^{}',
            '\\': r'\textbackslash{}',
            '<': r'\textless{}',
            '>': r'\textgreater{}',
        }
        regex = re.compile('|'.join(
            re.escape(key)
            for key in sorted(conv.keys(), key=lambda item: - len(item))))
        text_wo_latex_special = regex.sub(lambda match: conv[match.group()], text)
        return text_wo_latex_special.replace('\xc2\xa0', '').replace('€', '\\geneuro')

    def _escape_string(self, s):
        if isinstance(s, str) and (s[0] == s[-1] == "'" or s[0] == s[-1] == '"'):
            return s[1:-1].replace(r'\\', '\\').replace('\\n', '\n')
        return s

    def _evalstr(self, l, s, t):
        return self._latex_escape(self._evaltex(l, s, t))

    def _evaltex(self, l, s, t):
        code = t.code[1:-1]
        kernel = self.kernels[self.kernel_names[0]]
        results = kernel.execute(code)
        return '\n'.join(
            r.get('text/tex') or
            self._escape_string(r.get('text/plain')) or
            r.get('text/html') or
            r.get('error') or ''
            for r in results
        )

    def _sage(self, l, s, t):
        code = t.code[1:-1]
        return self._latexify(self.kernels['sagemath'].execute('latex(%s)'%code))

    def _sinput(self, l, s, t):
        filename = t.filename[1:-1]
        with open(filename, 'r') as codefile:
            code = codefile.read()
        return self._latexify(self._select_kernel(t).execute(code))

    def _truish(self, s):
        '''Return True if the string correspond to the True value
        in the current kernel.'''
        if self.kernel_names[0] in ('python2', 'python3', 'sagemath'):
            # TODO: non exhaustive (but just a helper for the user!)
            return s not in ('False', '', '[]', '0', '0.0')
        else:
            return s in ('true', 'True')

    def _sif(self, l, s, t):
        r'''\sif{condition}{texif}{texelse}
        Uses only the first kernel.
        The strings texif and texelse are not code to be executed, but tex strings,
        but those strings may contain \eval{code}.'''
        kernel = self.kernels[self.kernel_names[0]]
        code = t.condition[1:-1]
        results = kernel.execute(code)
        if (len(results) == 1 and self._truish(results[0].get('text/plain'))):
            return self.codeparser.transformString(t.texif[1:-1])
        else:
            return self.codeparser.transformString(t.texelse[1:-1])

    def _build_codeparser(self):
        # Look for usepackage[kernel]{surgery} markup to choose sage, python, R, julia
        #  or whatever interactive command line application
        # Use pyparsing as in student_surgery to go through sage|sagestr|sagesilent|sif|schoose|etc
        # in order
        # Use SimpleKernel to comunicate with the kernel

        usepackage = ('\\usepackage' +
                      Optional('[' + delimitedList(Word(alphanums))('kernels') + ']') +
                      '{texsurgery}')
        self.kernel_names = list(usepackage.searchString(self.src, maxMatches=1)[0]['kernels'])
        usepackage.setParseAction(lambda l, s, t: '')
        run = self._parserFor('run')
        run.setParseAction(self._run)
        runsilent = self._parserFor('runsilent')
        runsilent.setParseAction(self._runsilent)
        eval = self._parserFor('\\eval{code}', options=False)
        eval.setParseAction(self._eval)
        evalstr = self._parserFor('\\evalstr{code}', options=False)
        evalstr.setParseAction(self._evalstr)
        evaltex = self._parserFor('\\evaltex{code}', options=False)
        evaltex.setParseAction(self._evaltex)
        sage = self._parserFor('\\sage{code}', options=False)
        sage.setParseAction(self._sage)
        sinput = self._parserFor('\\sinput{filename}', options=False)
        sinput.setParseAction(self._sinput)
        sif = self._parserFor(
            '\\sif{condition}{texif}{texelse}', options=False
        )
        sif.setParseAction(self._sif)
        srepl = self._parserFor('srepl')
        srepl.setParseAction(self._srepl)
        self._codeparser = (usepackage | run | runsilent | eval | evalstr |
                            evaltex | sage | sif | sinput | srepl)
        # Do not ignore latex comments, because % is a useful python operator
        # Although there is an uncommon situation where this causes trouble:
        # if a comment contains an \eval{...some code...} and evaluation of the code
        # gives a result with a line break, compilation of the tex file will probably fail
#        self._codeparser.ignore('%' + restOfLine)

    def code_surgery(self):
        self.src = self.codeparser.transformString(self.src)
        return self

    def _opts_parser(self, str_opts):
        '''parse a string of the form "key1=val1,key2=val2..."'''
        return dict(map((lambda s: s.strip()), str_pair.strip().split('='))
                    for str_pair in str_opts.split(','))

    def _parserFor(self, selector, options=True):
        parts, args, restrictions = self._parse_selector(selector)
        name = parts.name
        if args:
            args_parser = sum(
                (Literal('{%s}' % restrictions[arg])(arg) if (arg in restrictions)
                 else skipToMatching('{', '}')(arg)
                 for arg in args),
                Empty()
            )
        # Drop this, since a command may have no arguments
#        elif name[0]=='\\':
#            args_parser = skipToMatching('{','}')('content')
        else:
            args_parser = Empty()
        if options:
            args_parser = Optional('[' + Word(alphanums)('options') + ']') + args_parser
        if name[0] == '\\':
            return (Literal(name)('name') + args_parser).leaveWhitespace()
        else:
            return ('\\begin{' + Literal(name)('name') + '}' +
                    args_parser +
                    SkipTo('\\end{'+name+'}')('content') +
                    '\\end{' + name + '}').leaveWhitespace()

    def _wholeEnvParserFor(self, env):
        return originalTextFor(
                ('\\begin{' + Literal(env) + '}')
                + SkipTo('\\end{' + env + '}')
                + ('\\end{' + env + '}')
            )('all')

    command_parser = (
        originalTextFor(Optional('\\') + Word(alphas))('name')
        + originalTextFor(Optional(nestedExpr('[', ']')))('options')
        + (ZeroOrMore(nestedExpr('{', '}')))('namedargs')
    )

    def _parse_selector(self, selector):
        parts = self.command_parser.searchString(selector)[0]
        args = []
        if parts.namedargs:
            args += [m[0] for m in parts.namedargs]
        if parts.options:
            options = self._opts_parser(parts.options[1:-1])
            if '_nargs' in options:
                nargs = int(options['_nargs'])
                args += ['arg%d' % k for k in range(nargs)]
                del options['_nargs']
            restrictions = options
        else:
            restrictions = {}
        return parts, args, restrictions

    def insertAfter(self, selector, text):
        match = self.find(selector)
        iend = match[-1].iend
        self.src = self.src[:iend] + text + self.src[iend:]
        return self

    def find(self, selector):
        # quick solution for issue #7
        # TODO: more efficient solution (see last proposal)
        # res = self.findall(selector, maxMatches=1)
        res = self.findall(selector)
        return res[0].flatten() if res else None

    def findall(self, selector, tex=None, maxMatches=pyparsing_MAX_INT, ibegin=0):
        r'''Finds all occurrences of a given selector

        currently it is not possible to look for commands nested inside commands,
        the parent can only be an environment.

        :param str selector: a string with the CSS-style selector
        :param str tex: string to search, usually None except for recursive calls
        :param int maxMatches: maximum number of matches, usually either 1
                               or a very big number

        :returns ResultSet of Match

        >>> from texsurgery.texsurgery import TexSurgery
        >>> tex = open('tests/test_find.tex').read()
        >>> TexSurgery(tex).findall('question,questionmultx runsilent')
        [('questionmultx', [('runsilent', 'a = randint(1,10)\n')]),
         ('question',
           [('runsilent', 'a = randint(2,10)\nf = sin(a*x)\nfd = f.derivative(x)\n')])]
        >>> TexSurgery(tex).findall('question,questionmultx choices \correctchoice')
        [('question', [('choices', [('\correctchoice', '$\sage{fd}$')])])]
        >>> TexSurgery(tex).findall('questionmultx \AMCnumericChoices[_nargs=2]')
        [('questionmultx', [('\\AMCnumericChoices',
          ['\\eval{8+a}', 'digits=2,sign=false,scoreexact=3'])]
        )]
        '''
        if tex == '':
            return ResultSet(self, [])
        elif tex is None:
            tex = self.src
        # First, if there is a ", " at the top level, we split there
        if ', ' in selector:
            partials = sum((self.findall(subselector) for subselector in selector.split(', ')), [])
            return ResultSet(self, partials)
        parent, *rest = selector.split(' ', 1)
        if rest:
            # If parent is a command, we want to capture its first argument
            # to look for the rest of the selector inside
            alternatives = [env_or_command+'{content}' if env_or_command[0] == '\\'
                            else env_or_command
                            for env_or_command in parent.split(',')]
            names = Or([self._parserFor(env_or_command)
                        for env_or_command in alternatives])
            # if rest, then parent has no attributes :-/
            # issue #7: find calls findall with optional argument maxMatches=1
            #  => find does not find some nested commands that findall finds
            # if self.findall(rest[0], match.content) is None, the only match is discarded
            elements = [
                (EnvOrCommand.newEnvOrCommand(name=match.name,
                                              src=self.src,
                                              istart=ibegin + istart, iend=ibegin + iend),
                 match)
                for match, istart, iend in names.scanString(tex, maxMatches=maxMatches)
            ]
            matches = [
                (element, self.findall(rest[0], match.content, ibegin=element.icontent))
                for (element, match) in elements
            ]
            return ResultSet(
                self,
                [Match(element, nest) for (element, nest) in matches if nest]
            )

        names = Or([self._parserFor(env_or_command)
                    for env_or_command in parent.split(',')])
        # The tail part of the selector may have named arguments and restrictions
        parts, args, restrictions = self._parse_selector(selector)
        if selector[0] == '\\':
            return ResultSet(self, [
                Match(Command(name=match.name,
                              src=self.src, istart=ibegin + istart, iend=ibegin + iend,
                              options=match.options,
                              arguments=(Arguments(**{arg: match[arg][1:-1] for arg in args}))
                ))
                for match, istart, iend in names.scanString(tex, maxMatches=maxMatches)
            ])
        else:
            return ResultSet(self, [
                Match(Environment(name=match.name,
                                  src=self.src, istart=ibegin + istart, iend=ibegin + iend,
                                  options=match.options,
                                  arguments=Arguments(**{arg: match[arg][1:-1] for arg in args}),
                                  content=match.content))
                for match, istart, iend in names.scanString(tex, maxMatches=maxMatches)
            ])

    def shuffle(self, parentselector, childrenselector, seed=1):
        pass
