"""Parsing mechanisms should not be directly invoked publicly, as they are
subject to change."""

from TexSoup.utils import Token, Buffer, MixedBuffer, CharToLineOffset
from TexSoup.data import *
from TexSoup.data import arg_type
from TexSoup.tokens import (
    TC,
    tokenize,
    SKIP_ENVS,
)
import string
import sys


MATH_ENVS = (
    TexDisplayMathModeEnv,
    TexMathModeEnv,
    TexDisplayMathEnv,
    TexMathEnv
)
MATH_TOKEN_TO_ENV = {env.token_begin: env for env in MATH_ENVS}
ARG_BEGIN_TO_ENV = {arg.token_begin: arg for arg in arg_type}

SIGNATURES = {
    'def': (2, 0),
    'textbf': (1, 0),
    'section': (1, 0)
}


__all__ = ['read_expr', 'read_tex']


def read_tex(buf, skip_envs=(), tolerance=0):
    r"""Parse all expressions in buffer

    :param Buffer buf: a buffer of tokens
    :param Tuple[str] skip_envs: environments to skip parsing
    :param int tolerance: error tolerance level (only supports 0 or 1)
    :return: iterable over parsed expressions
    :rtype: Iterable[TexExpr]
    """
    while buf.hasNext():
        yield read_expr(buf,
            skip_envs=SKIP_ENVS + skip_envs,
            tolerance=tolerance)


# TODO: skip envs nested in args or even just brace group won't be skipped
# (skip_env kwarg needs to be propagated)
def read_expr(src, skip_envs=(), tolerance=0):
    r"""Read next expression from buffer

    :param Buffer src: a buffer of tokens
    :param Tuple[str] skip_envs: environments to skip parsing
    :param int tolerance: error tolerance level (only supports 0 or 1)
    :return: parsed expression
    :rtype: [TexExpr, Token]
    """
    c = next(src)
    # TODO: assemble and use groups
    if c.category in MATH_TOKEN_TO_ENV.keys():
        expr = MATH_TOKEN_TO_ENV[c.category]([])
        return read_math_env(src, expr)
    elif c.category == TC.Escape:
        # TODO: reduce to command-parsing only -- assemble envs in 2nd pass
        name, args, steps = peek_command(src)
        if name == 'item':
            contents, arg = read_item(src)
            expr = TexCmd(name, contents, arg)
        elif name == 'begin':
            assert args, 'Begin command must be followed by an env name.'
            expr = TexNamedEnv(args[0].string)
            expr.args = args[1:]
            src.forward(steps)

            if expr.name in skip_envs:
                read_skip_env(src, expr)
            else:
                read_env(src, expr, tolerance=tolerance)
        else:
            src.forward(1)
            expr = TexCmd(name)
            expr.args = read_args(src, args=expr.args)
        return expr
    if c.category == TC.GroupStart:
        return read_arg(src, c, tolerance=tolerance)

    assert isinstance(c, Token)
    return TexText(c)


################
# ENVIRONMENTS #
################


def read_item(src, tolerance=0):
    r"""Read the item content. Assumes escape has just been parsed.

    There can be any number of whitespace characters between \item and the
    first non-whitespace character. Any amount of whitespace between subsequent
    characters is also allowed.

    \item can also take an argument.

    :param Buffer src: a buffer of tokens
    :param int tolerance: error tolerance level (only supports 0 or 1)
    :return: contents of the item and any item arguments

    >>> from TexSoup.category import categorize
    >>> from TexSoup.tokens import tokenize
    >>> def read_item_from(string, skip=1):
    ...     buf = tokenize(categorize(string))
    ...     _ = buf.forward(skip)
    ...     return read_item(buf)
    >>> read_item_from(r'\item aaa {bbb} ccc\end{itemize}')
    (['aaa ', BraceGroup('bbb'), ' ccc'], [])
    >>> read_item_from(r'\item aaa \textbf{itemize}\item no')
    (['aaa ', TexCmd('textbf', [BraceGroup('itemize')])], [])
    >>> read_item_from(r'\item[aaa] yith\item no')
    (['yith'], [BracketGroup('aaa')])
    >>> read_item_from('\\item\n[aaa] yith\\item no')
    (['yith'], [BracketGroup('aaa')])
    >>> read_item_from(r'\item WITCH [nuuu] DOCTORRRR 👩🏻‍⚕️')
    (['WITCH ', '[', 'nuuu', ']', ' DOCTORRRR 👩🏻‍⚕️'], [])
    >>> read_item_from(r'''\begin{itemize}
    ... \item
    ... \item first item
    ... \end{itemize}''', skip=7)
    ([], [])
    >>> read_item_from(r'''\def\itemeqn{\item}''', skip=6)
    ([], [])
    """
    assert next(src) == 'item', src.backward(1)
    args, extras = [], []

    # TODO: fix when spacer tokenization updated
    spacer, rest = read_spacer(src)
    if rest:
        extras.append(rest)

    # TODO: use peek_command instead of manually parsing optional arg
    if not rest and src.hasNext() and src.peek().category == TC.OpenBracket:
        c = next(src)
        args.append(read_arg(src, c, tolerance=tolerance))

        # remove leading spacer after arguments, due to quirk in Item repr
        # which adds space after args.
        spacer, rest = read_spacer(src)
        if rest:
            extras.append(rest)

    while src.hasNext():
        if src.peek().category == TC.Escape:
            cmd_name, cmd_args, steps = peek_command(src, 1, skip=1)
            if cmd_name in ('end', 'item'):
                return extras, args
        elif src.peek().category == TC.GroupEnd:
            break
        extras.append(read_expr(src, tolerance=tolerance))
    return extras, args


def unclosed_env_handler(src, expr, end):
    """Handle unclosed environments.

    Currently raises an end-of-file error. In the future, this can be the hub
    for unclosed-environment fault tolerance.

    :param Buffer src: a buffer of tokens
    :param TexExpr expr: expression for the environment
    :param end str: Actual end token (as opposed to expected)
    """
    clo = CharToLineOffset(str(src))
    explanation = 'Instead got %s' % end if end else 'Reached end of file.'
    line, offset = clo(src.position)
    raise EOFError('[Line: %d, Offset: %d] "%s" env expecting %s. %s' % (
        line, offset, expr.name, expr.end, explanation))


def read_math_env(src, expr):
    r"""Read the environment from buffer.

    Advances the buffer until right after the end of the environment. Adds
    parsed content to the expression automatically.

    :param Buffer src: a buffer of tokens
    :param TexExpr expr: expression for the environment
    :rtype: TexExpr

    >>> from TexSoup.category import categorize
    >>> from TexSoup.tokens import tokenize
    >>> buf = tokenize(categorize(r'\min_x \|Xw-y\|_2^2'))
    >>> read_math_env(buf, TexMathModeEnv())
    Traceback (most recent call last):
        ...
    EOFError: [Line: 0, Offset: 7] "$" env expecting $. Reached end of file.
    """
    content = src.forward_until(lambda c: c.category == expr.token_end)
    if not src.hasNext() or src.peek().category != expr.token_end:
        unclosed_env_handler(src, expr, src.peek())
    next(src)
    expr.append(content)
    return expr


def read_skip_env(src, expr):
    r"""Read the environment from buffer, WITHOUT parsing contents

    Advances the buffer until right after the end of the environment. Adds
    UNparsed content to the expression automatically.

    :param Buffer src: a buffer of tokens
    :param TexExpr expr: expression for the environment
    :rtype: TexExpr

    >>> from TexSoup.category import categorize
    >>> from TexSoup.tokens import tokenize
    >>> buf = tokenize(categorize(r' \textbf{aa \end{foobar}ha'))
    >>> read_skip_env(buf, TexNamedEnv('foobar'))
    TexNamedEnv('foobar', [' \\textbf{aa '], [])
    >>> buf = tokenize(categorize(r' \textbf{aa ha'))
    >>> read_skip_env(buf, TexNamedEnv('foobar'))
    Traceback (most recent call last):
        ...
    EOFError: [Line: 0, Offset: 5] "foobar" env expecting \end{foobar}. Reached end of file.
    """
    condition = lambda s: s.startswith('\\end{%s}' % expr.name)
    contents = [src.forward_until(condition, peek=False)]
    if not src.startswith('\\end{%s}' % expr.name):
        unclosed_env_handler(src, expr, src.peek((0, 6)))
    src.forward(5)
    expr.append(*contents)
    return expr


def read_env(src, expr, tolerance=0):
    r"""Read the environment from buffer.

    Advances the buffer until right after the end of the environment. Adds
    parsed content to the expression automatically.

    :param Buffer src: a buffer of tokens
    :param TexExpr expr: expression for the environment
    :param int tolerance: error tolerance level (only supports 0 or 1)
    :rtype: TexExpr

    >>> from TexSoup.category import categorize
    >>> from TexSoup.tokens import tokenize
    >>> buf = tokenize(categorize(' tingtang \\end\n{foobar}walla'))
    >>> read_env(buf, TexNamedEnv('foobar'))
    TexNamedEnv('foobar', [' tingtang '], [])
    >>> buf = tokenize(categorize(' tingtang \\end\n\n{foobar}walla'))
    >>> read_env(buf, TexNamedEnv('foobar')) #doctest: +ELLIPSIS
    Traceback (most recent call last):
        ...
    EOFError: [Line: 0, Offset: 1] ...
    >>> buf = tokenize(categorize(' tingtang \\end\n\n{nope}walla'))
    >>> read_env(buf, TexNamedEnv('foobar'), tolerance=1)  # error tolerance
    TexNamedEnv('foobar', [' tingtang '], [])
    """
    contents = []
    while src.hasNext():
        if src.peek().category == TC.Escape:
            name, args, steps = peek_command(src, n_required_args=1, skip=1)
            if name == 'end':
                break
        contents.append(read_expr(src, tolerance=tolerance))
    error = not src.hasNext() or not args or args[0].string != expr.name
    if error and tolerance == 0:
        unclosed_env_handler(src, expr, src.peek((0, 6)))
    elif not error:
        src.forward(5)
    expr.append(*contents)
    return expr


############
# COMMANDS #
############


# TODO: handle macro-weirdness e.g., \def\blah[#1][[[[[[[[#2{"#1 . #2"}
# TODO: add newcommand macro
def read_args(src, n_required=-1, n_optional=-1, args=None, tolerance=0):
    r"""Read all arguments from buffer.

    This function assumes that the command name has already been parsed. By
    default, LaTeX allows only up to 9 arguments of both types, optional
    and required. If `n_optional` is not set, all valid bracket groups are
    captured. If `n_required` is not set, all valid brace groups are
    captured.

    :param Buffer src: a buffer of tokens
    :param TexArgs args: existing arguments to extend
    :param int n_required: Number of required arguments. If < 0, all valid
                           brace groups will be captured.
    :param int n_optional: Number of optional arguments. If < 0, all valid
                           bracket groups will be captured.
    :param int tolerance: error tolerance level (only supports 0 or 1)
    :return: parsed arguments
    :rtype: TexArgs

    >>> from TexSoup.category import categorize
    >>> from TexSoup.tokens import tokenize
    >>> test = lambda s, *a, **k: read_args(tokenize(categorize(s)), *a, **k)
    >>> test('[walla]{walla}{ba]ng}')  # 'regular' arg parse
    [BracketGroup('walla'), BraceGroup('walla'), BraceGroup('ba', ']', 'ng')]
    >>> test('\t[wa]\n{lla}\n\n{b[ing}')  # interspersed spacers + 2 newlines
    [BracketGroup('wa'), BraceGroup('lla')]
    >>> test('\t[\t{a]}bs', 2, 0)  # use char as arg, since no opt args
    [BraceGroup('['), BraceGroup('a', ']')]
    >>> test('\n[hue]\t[\t{a]}', 2, 1)  # check stop opt arg capture
    [BracketGroup('hue'), BraceGroup('['), BraceGroup('a', ']')]
    >>> test('\t\\item')
    []
    >>> test('   \t    \n\t \n{bingbang}')
    []
    >>> test('[tempt]{ing}[WITCH]{doctorrrr}', 0, 0)
    []
    """
    args = args or TexArgs()
    if n_required == 0 and n_optional == 0:
        return args

    n_optional = read_arg_optional(src, args, n_optional, tolerance)
    n_required = read_arg_required(src, args, n_required, tolerance)

    if src.hasNext() and src.peek().category == TC.OpenBracket:
        n_optional = read_arg_optional(src, args, n_optional, tolerance)
    if src.hasNext() and src.peek().category == TC.GroupStart:
        n_required = read_arg_required(src, args, n_required, tolerance)
    return args


def read_arg_optional(src, args, n_optional=-1, tolerance=0):
    """Read next optional argument from buffer.

    If the command has remaining optional arguments, look for:
       a. A spacer. Skip the spacer if it exists.
       b. A bracket delimiter. If the optional argument is bracket-delimited,
          the contents of the bracket group are used as the argument.

    :param Buffer src: a buffer of tokens
    :param TexArgs args: existing arguments to extend
    :param int n_optional: Number of optional arguments. If < 0, all valid
                           bracket groups will be captured.
    :param int tolerance: error tolerance level (only supports 0 or 1)
    :return: number of remaining optional arguments
    :rtype: int
    """
    while n_optional != 0:
        spacer, no_optional_args = read_spacer(src)
        if no_optional_args:
            src.backward(1)
            break
        if not (src.hasNext() and src.peek().category == TC.OpenBracket):
            break
        args.append(read_arg(src, next(src), tolerance=tolerance))
        n_optional -= 1
    return n_optional


def read_arg_required(src, args, n_required=-1, tolerance=0):
    """Read next required argument from buffer.

    If the command has remaining required arguments, look for:
       a. A spacer. Skip the spacer if it exists.
       b. A curly-brace delimiter. If the required argument is brace-delimited,
          the contents of the brace group are used as the argument.
       c. Spacer or not, if a brace group is not found, simply use the next
          character.

    :param Buffer src: a buffer of tokens
    :param TexArgs args: existing arguments to extend
    :param int n_required: Number of required arguments. If < 0, all valid
                           brace groups will be captured.
    :param int tolerance: error tolerance level (only supports 0 or 1)
    :return: number of remaining optional arguments
    :rtype: int

    >>> from TexSoup.category import categorize
    >>> from TexSoup.tokens import tokenize
    >>> buf = tokenize(categorize('{wal]la}{ba ng}'))
    >>> args = TexArgs()
    >>> read_arg_required(buf, args)  # 'regular' arg parse
    -3
    >>> args
    [BraceGroup('wal', ']', 'la'), BraceGroup('ba ng')]
    """
    while n_required != 0:
        spacer, ungrouped_required_chars = read_spacer(src)

        # TODO: technically may drop text from buffer that is not used. Fix
        # after fixing generic string tokeinzation
        while n_required > 0 and ungrouped_required_chars:
            args.append('{%s}' % ungrouped_required_chars[0])
            spacer, ungrouped_required_chars = read_spacer(
                Buffer([TexText(ungrouped_required_chars[1:])]))
            n_required -= 1

        if ungrouped_required_chars:
            src.backward()

        if n_required == 0 or (n_required < 0 and ungrouped_required_chars):
            break

        if src.hasNext() and src.peek().category == TC.GroupStart:
            args.append(read_arg(src, next(src), tolerance=tolerance))
            n_required -= 1
            continue
        elif src.hasNext() and n_required > 0:
            args.append('{%s}' % next(src))
            n_required -= 1
            continue

        if spacer:
            src.backward(1)
        break
    return n_required


def read_arg(src, c, tolerance=0):
    r"""Read the argument from buffer.

    Advances buffer until right before the end of the argument.

    :param Buffer src: a buffer of tokens
    :param str c: argument token (starting token)
    :param int tolerance: error tolerance level (only supports 0 or 1)
    :return: the parsed argument
    :rtype: TexGroup

    >>> from TexSoup.category import categorize
    >>> from TexSoup.tokens import tokenize
    >>> s = r'''{\item\abovedisplayskip=2pt\abovedisplayshortskip=0pt~\vspace*{-\baselineskip}}'''
    >>> buf = tokenize(categorize(s))
    >>> read_arg(buf, next(buf))
    BraceGroup(TexCmd('item'))
    """
    content = [c]
    arg = ARG_BEGIN_TO_ENV[c.category]
    while src.hasNext():
        if src.peek().category == arg.token_end:
            src.forward()
            return arg(*content[1:])
        else:
            content.append(read_expr(src, tolerance=tolerance))

    clo = CharToLineOffset(str(src))
    line, offset = clo(c.position)
    raise TypeError(
        '[Line: %d, Offset %d] Malformed argument. First and last elements '
        'must match a valid argument format. In this case, TexSoup'
        ' could not find matching punctuation for: %s.\n'
        'Just finished parsing: %s' %
        (line, offset, c, content))


# TODO: move spacer tokenization to tokenizer
# WARNING: This method is flawed: Spacer detection only works for first
# instance of spacer in text. Will need to refactor generic string tokenization
# which can only occur after refactoring item, arg, and env readers above,
# which in turn, rely on the skip_env reader.
def read_spacer(buf):
    r"""Extracts the next spacer, if there is one, before non-whitespace

    Define a spacer to be a contiguous string of only whitespace, with at most
    one line break.

    >>> from TexSoup.category import categorize
    >>> from TexSoup.tokens import tokenize
    >>> read_spacer(Buffer(tokenize(categorize('   \t    \n'))))
    ('   \t    \n', '')
    >>> read_spacer(Buffer(tokenize(categorize('   \t    \n\t \n  \t\n'))))
    ('   \t    \n\t ', '\n  \t\n')
    >>> read_spacer(Buffer(tokenize(categorize('{'))))
    ('', '')
    >>> read_spacer(Buffer(tokenize(categorize('   \t    \na'))))
    ('   \t    \n', 'a')
    >>> read_spacer(Buffer(tokenize(categorize('   \t    \n\t \n  \t\na'))))
    ('   \t    \n\t ', '\n  \t\na')
    """
    if not buf.hasNext() or not buf.peek().category == TC.Text:
        return '', ''

    text, lines = next(buf), 1
    spacer, rest, is_spacer = '', '', True
    for c in text:
        if c == '\n':  # TODO: change to token code
            lines += 1
        if is_spacer and lines > 2 or not c.isspace():
            is_spacer = False
        if is_spacer:
            spacer += c
        else:
            rest += c

    return spacer, rest


# TODO: refactor after generic string tokenizer fixed
# TODO: hard-coded to 1 required arg
# TODO: make this a reader, with a generic peek decorator or wrapper
def peek_command(buf, n_required_args=-1, n_optional_args=-1, skip=0):
    r"""Parses command and all arguments. Assumes escape has just been parsed.

    Here are the rules for command name parsing:

    No whitespace is allowed between escape and command name. e.g., \ textbf
    is the "command" "\" and the text "textbf". Only \textbf is the bold
    command.

    >>> from TexSoup.category import categorize
    >>> from TexSoup.tokens import tokenize
    >>> buf = Buffer(tokenize(categorize('\\sect  \t    \n\t{wallawalla}')))
    >>> next(buf)
    '\\'
    >>> peek_command(buf)
    ('sect', [BraceGroup('wallawalla')], 5)
    >>> buf = Buffer(tokenize(categorize('\\sect  \t   \n\t \n{bingbang}')))
    >>> _ = next(buf)
    >>> peek_command(buf)
    ('sect', [], 1)
    >>> buf = Buffer(tokenize(categorize('\\sect{ooheeeee}')))
    >>> _ = next(buf)
    >>> peek_command(buf)
    ('sect', [BraceGroup('ooheeeee')], 4)

    # Broken because abcd is incorrectly tokenized with leading space
    # >>> buf = Buffer(tokenize(categorize('\\sect abcd')))
    # >>> _ = next(buf)
    # >>> peek_command(buf)
    # ('sect', ('a',), 2)
    """
    position = buf.position
    for _ in range(skip):
        next(buf)

    name = next(buf)
    token = Token('', buf.position)
    if n_required_args < 0 and n_optional_args < 0:
        n_required_args, n_optional_args = SIGNATURES.get(name, (-1, -1))
    args = read_args(buf, n_required_args, n_optional_args)

    steps = buf.position - position
    buf.backward(steps)
    return name, args, steps
