# This file is part of Rubber and thus covered by the GPL
# (c) Emmanuel Beffara, 2008
# vim: noet:ts=4
"""
General-purpose classes for reading TeX code.
Classes and functions from this module can be used without Rubber.
"""

import re
from io import StringIO

# The catcodes

EOF = -2
CSEQ = -1
ESCAPE = 0
OPEN = 1
CLOSE = 2
MATH = 3
ALIGN = 4
END_LINE = 5
ARGUMENT = 6
SUPER = 7
SUB = 8
IGNORE = 9
SPACE = 10
LETTER = 11
OTHER = 12
ACTIVE = 13
COMMENT = 14
INVALID = 15

cat_names = {
    -2: 'EOF',
    -1: 'CSEQ',
    0: 'ESCAPE',
    1: 'OPEN',
    2: 'CLOSE',
    3: 'MATH',
    4: 'ALIGN',
    5: 'END_LINE',
    6: 'ARGUMENT',
    7: 'SUPER',
    8: 'SUB',
    9: 'IGNORE',
    10: 'SPACE',
    11: 'LETTER',
    12: 'OTHER',
    13: 'ACTIVE',
    14: 'COMMENT',
    15: 'INVALID'
}

# The default categories

catcodes = {
    '\\' : ESCAPE,
    '{' : OPEN,
    '}' : CLOSE,
    '$' : MATH,
    '&' : ALIGN,
    '\n' : END_LINE,
    '#' : ARGUMENT,
    '^' : SUPER,
    '_' : SUB,
    '\000' : IGNORE,
    ' ' : SPACE, '\t' : SPACE,
    '~' : ACTIVE,
    '%' : COMMENT,
    '\177' : INVALID
}

for i in range(0,26):
    catcodes[chr(ord('A')+i)] = LETTER
    catcodes[chr(ord('a')+i)] = LETTER

class Position:
    """
    A class to represent positions in a source file.
    """
    def __init__ (self, file=None, line=None, char=None):
        self.file = file
        self.line = line
        self.char = char
    def __str__ (self):
        text = ''
        if self.file:
            text = file + ':'
        if self.line is not None:
            if text != '':
                text += ':'
            text += '%d' % self.line
            if self.char is not None:
                text += ':%d' % self.char
        return text

class Token:
    """
    The class used to represent tokens. Objects contain a catcode, a value
    (for control sequences) and the raw text that represents them in the input
    file.
    """
    def __init__ (self, cat, val=None, raw=None, pos=None):
        self.cat = cat
        self.val = val
        self.raw = raw
        self.pos = pos

    def __repr__ (self):
        text = 'Token(' + cat_names[self.cat]
        if self.val is not None:
            text += ', ' + repr(self.val)
        return text + ')'

class TokenList (list):
    """
    This class represents a token list. It behaves as a standard list with
    some extra functionality.
    """
    def __init__ (self, data=[], pos=None):
        super (TokenList, self).__init__(data)
        if pos is None and len(data) > 0:
            self.pos = data[0].pos
        else:
            self.pos = pos

    def raw_text (self):
        """
        Return the textual representation of the token list by concatenating
        the raw text of the tokens.
        """
        text = ''
        for token in self:
            text += token.raw
        return text

class ParserBase (object):
    """
    This is the base class for parsers. It holds state information like
    catcodes, handles the push-back buffer, and leaves it to derived classes
    to actually read tokens, using the "read_token" method. This class also
    provides high-level functionality like getting macro arguments.
    """
    def __init__ (self):
        self.catcodes = catcodes.copy()
        self.next = []
        self.math_mode = 0
        self.last_is_math = 0
        self.pos = None

    def catcode (self, char):
        """
        Return the catcode of a character.
        """
        if char in self.catcodes:
            return self.catcodes[char]
        else:
            return OTHER

    def put_token (self, token):
        """
        Put back a token in the input.
        """
        self.next.append(token)

    def put_list (self, list):
        """
        Put back a token list in the input.
        """
        arg = list[:]
        arg.reverse()
        self.next.extend(arg)

    def peek_token (self):
        """
        Return the next token that will be read without updating the state.
        """
        if len(self.next) > 0:
            return self.next[-1]
        token = self.read_token()
        self.put_token(token)
        return token

    def get_token (self):
        """
        Get the next token from the input and update the math mode.
        """
        if len(self.next) > 0:
            token = self.next.pop()
        else:
            token = self.read_token()

        # skip over comment
        if token.cat == COMMENT:
            assert len(self.next) == 0
            assert self.next_char is None
            self.read_line()
            return self.read_token()

        if token.cat == MATH:
            if self.last_is_math:
                if self.math_mode == 1:
                    self.math_mode = 2
                self.last_is_math = 0
            else:
                if self.math_mode == 0:
                    self.math_mode = 1
                else:
                    self.math_mode = 0
                self.last_is_math = 1
        else:
            self.last_is_math = 0

        return token

    def __iter__ (self):
        """
        Return an iterator over all tokens in the input. The EOF token is not
        returned by this iterator.
        """
        while 1:
            token = self.get_token()
            if token.cat == EOF:
                break
            yield token

    def skip_space (self):
        """
        Skip white space in the input.
        """
        while self.peek_token().cat in (SPACE, END_LINE):
            self.get_token()

    def get_group (self):
        """
        Get the list of tokens up to the next closing brace, and drop the
        closing brace.
        """
        value = TokenList()
        level = 1
        while 1:
            token = self.get_token()
            if token.cat == OPEN:
                level += 1
            elif token.cat == CLOSE:
                level -= 1
                if level == 0:
                    break
            elif token.cat == EOF:
                break
            value.append(token)
        return value

    def get_group_text (self):
        """
        Get the list of tokens up to the next closing brace, and drop the
        closing brace. Return the list as a string.
        """
        value = ""
        level = 1
        while 1:
            token = self.get_token()
            if token.cat == OPEN:
                level += 1
            elif token.cat == CLOSE:
                level -= 1
                if level == 0:
                    break
            elif token.cat == EOF:
                break
            value += token.raw
        return value

    def get_argument (self):
        """
        Get a macro argument from the input text. Returns a token list with
        the value of the argument, with surrounding braces removed if
        relevant.
        """
        self.skip_space()
        token = self.get_token()
        if token.cat == EOF:
            return TokenList()
        if token.cat != OPEN:
            return TokenList(data=[token])
        return self.get_group()

    def get_argument_text (self):
        """
        Get a macro argument from the input text. Returns a string with
        the text of the argument, with surrounding braces removed if
        relevant.
        """
        self.skip_space()
        token = self.get_token()
        if token.cat == EOF:
            return None
        if token.cat != OPEN:
            return token.raw
        return self.get_group_text()

    def get_latex_optional (self):
        """
        Check if a LaTeX-style optional argument is present. If such an
        argument is present, return it as a token list, otherwise return None.
        """
        self.skip_space()
        next = self.get_token()

        if next.cat != OTHER or next.raw != '[':
            self.put_token(next)
            return None

        level = 0
        list = TokenList()
        while True:
            token = self.get_token()
            if token.cat == EOF:
                break
            if token.cat == OTHER and token.raw == ']' and level == 0:
                break
            if token.cat == OPEN:
                level += 1
            elif token.cat == CLOSE:
                if level == 0:
                    break
                level -= 1
            list.append(token)

        return list

    def get_latex_optional_text (self):
        """
        Check if a LaTeX-style optional argument is present. If such an
        argument is present, return it as text, otherwise return None.
        """
        list = self.get_latex_optional()
        if list is None:
            return None
        return list.raw_text()

    def get_latex_star (self):
        """
        Check if the command is a starred one.  If so, eat the star,
        and return True.  Otherwise, return False.
        """
        nextt = self.peek_token()
        if (nextt.cat, nextt.raw) == (OTHER, '*'):
            self.get_token()
            return True
        else:
            return False

def re_set (set, complement=False):
    """
    Returns a string that contains a regular expression matching a given set
    of characters, or its complement if the optional argument is true. The set
    must not be empty.
    """
    if len(set) == 0:
        raise RuntimeError('argument of re_set must not be empty')
    if not complement and len(set) == 1:
        c = set[0]
        if c in '.^$*+?{}\\[]|()':
            return '\\' + c
        else:
            return c
    expr = '['
    if complement:
        expr += '^'
    for c in set:
        if c in ']-\\':
            expr += '\\' + c
        else:
            expr += c
    return expr + ']'

class Parser (ParserBase):
    """
    A parser for TeX code that reads its input from a file object.

    The class also provides a hook feature: the method 'set_hooks' declares a
    set of control sequence names, and the method 'next_hook' parses the input
    until it finds a control sequence from this set, ignoring all other
    tokens. This advantage of this method is that is is much faster than
    reading tokens one by one.
    """
    def __init__ (self, input):
        """
        Initialise the parser with a file as input.
        If 'input' is None, then input can only be provided by
        the 'put_token' and 'put_list' methods.
        """
        super (Parser, self).__init__()
        self.input = input
        self.line = ""
        self.pos_line = 1
        self.pos_char = 1
        self.next_char = None

    def read_line (self):
        """
        Reads a line of input and sets the attribute 'line' with it. Returns
        True if reading succeeded and False if it failed.
        """
        if self.input is None:
            return False
        self.line = self.input.readline()
        if self.line == "":
            return False
        return True

    def read_char (self):
        """
        Get the next character from the input and its catcode (without parsing
        control sequences).
        """
        if self.next_char is not None:
            t = self.next_char
            self.next_char = None
            return t

        while self.line == "":
            if not self.read_line():
                return Token(EOF)
        c = self.line[0]
        self.line = self.line[1:]

        pos = Position(line=self.pos_line, char=self.pos_char)
        if c == '\n':
            self.pos_line += 1
            self.pos_char = 1
        else:
            self.pos_char += 1

        return Token(self.catcode(c), raw=c, pos=pos)

    def read_token (self):
        """
        Get the next token from the input.
        """
        token = self.read_char()
        if token.cat != ESCAPE:
            if token.cat in (LETTER, OTHER):
                token.val = token.raw
            return token
        pos = token.pos
        raw = token.raw
        token = self.read_char()
        if token.cat != LETTER:
            token.cat = CSEQ
            token.val = token.raw
            token.raw = raw + token.raw
            return token
        name = ""
        while token.cat == LETTER:
            raw += token.raw
            name += token.raw
            token = self.read_char()
        while token.cat == SPACE:
            raw += token.raw
            token = self.read_char()
        self.next_char = token
        return Token(CSEQ, name, raw, pos=pos)

    def re_cat (self, *cat):
        """
        Returns a regular expression that maches characters whose category is
        in given list.
        """
        return re_set([char for char,code in self.catcodes.items() if code in cat])

    def re_nocat (self, *cat):
        """
        Returns a regular expression that maches characters whose category is
        not in a given list.
        """
        return re_set([char for char,code in self.catcodes.items() if code in cat], True)

    def set_hooks (self, names):
        """
        Define the set of hooks for 'next_hook'.
        """
        expr = '(' \
            + self.re_nocat(ESCAPE, COMMENT) + '|' \
            + self.re_cat(ESCAPE) + self.re_cat(ESCAPE, COMMENT) + ')*' \
            + '(?P<raw>' + self.re_cat(ESCAPE) \
            + '(?P<val>' + '|'.join(names) + ')' \
            + '(' + self.re_cat(SPACE) + '+|(?=' + self.re_nocat(LETTER) + ')|$))'
        self.regex = re.compile(expr)

    def next_hook (self):
        """
        Ignore input until the next control sequence from the set defined by
        'set_hooks'. Returns the associated token, or the EOF token if no hook
        was found.
        """
        while self.line == "":
            if not self.read_line():
                return Token(EOF)
        while True:
            match = self.regex.match(self.line)
            if match is not None:
                self.pos_char = match.end('raw') + 1
                self.line = self.line[match.end('raw'):]
                return Token(CSEQ, match.group('val'), match.group('raw'))
            if not self.read_line():
                return Token(EOF)
            self.pos_line += 1
            self.pos_char = 1

def parse_string (text):
    """
    Factory function for parsing TeX code from a string.
    """
    return Parser(StringIO(text))
