from string import ascii_uppercase
from pyparsing import *


class Node(object):
    def __init__(self, s, loc, toks):
        self.children = toks

    def __repr__(self):
        return '%s(%s)' % (self.__class__.__name__, ', '.join(map(str, self.children)))

    @classmethod
    def grammar(cls, expression):
        return expression.setParseAction(cls)


class Module(Node):
    def __init__(self, s, loc, toks):
        self.body = toks
        super().__init__(s, loc, toks)
class Expression(Node):
    def __init__(self, s, loc, toks):
        self.left, self.op, self.right = toks[0]
        super().__init__(s, loc, toks)
class Main(Node):
    def __init__(self, s, loc, toks):
        self.body = toks
        super().__init__(s, loc, toks)
class AliasDefinition(Node):
    def __init__(self, s, loc, toks):
        self.alias, self.value = toks
        super().__init__(s, loc, toks)
class Alias(Node):
    def __init__(self, s, loc, toks):
        self.alias = toks[0]
        super().__init__(s, loc, toks)
class Assignment(Node):
    def __init__(self, s, loc, toks):
        self.left, self.right = toks
        super().__init__(s, loc, toks)
class MemoryAssignment(Node):
    def __init__(self, s, loc, toks):
        self.left, self.right = toks
        super().__init__(s, loc, toks)
class AugmentedAssignment(Node):
    def __init__(self, s, loc, toks):
        self.left, self.operator, self.right = toks
        super().__init__(s, loc, toks)
class BuiltinCall(Node):
    def __init__(self, s, loc, toks):
        self.name = toks[0]
        self.args = toks[1:]
        super().__init__(s, loc, toks)
class FunctionCall(Node):
    def __init__(self, s, loc, toks):
        self.name = toks[0]
        super().__init__(s, loc, toks)
class Register(Node):
    def __init__(self, s, loc, toks):
        self.r = toks[0]
        super().__init__(s, loc, toks)
class Name(Node):
    def __init__(self, s, loc, toks):
        self.n = toks[0]
        super().__init__(s, loc, toks)
class Variable(Node):
    def __init__(self, s, loc, toks):
        self.v = toks[0]
        super().__init__(s, loc, toks)
class String(Node):
    def __init__(self, s, loc, toks):
        self.s = toks[0][1:-1]
        super().__init__(s, loc, toks)
class Hex(Node):
    def __init__(self, s, loc, toks):
        self.h = toks[0]
        super().__init__(s, loc, toks)
    @property
    def i(self):
        return int(self.h, 16)
class Integer(Node):
    def __int__(self, s, loc, toks):
        super().__init__(s, loc, toks)
    @property
    def i(self):
        return int(self.children[0])
class MemoryLocation(Node):
    def __init__(self, s, loc, toks):
        if len(toks) == 1:
            self.location = toks[0]
            self.simple = True
        else:
            self.left, self.op, self.right = toks
            self.simple = False
        super().__init__(s, loc, toks)
class ForLoop(Node):
    def __init__(self, s, loc, toks):
        self.targets, self.iterator, self.body = toks
        super().__init__(s, loc, toks)
class FunctionDefinition(Node):
    def __init__(self, s, loc, toks):
        self.alias = toks[0]
        self.body = toks[1:]
        super().__init__(s, loc, toks)
class Import(Node):
    def __init__(self, s, loc, toks):
        self.module = '.'.join(toks[0])
        super().__init__(s, loc, toks)
class Data(Node):
    def __init__(self, s, loc, toks):
        self.alias = toks[0]
        self.values = toks[1:]
        super().__init__(s, loc, toks)


indentStack = [1]

def checkPeerIndent(s,l,t):
    curCol = col(l,s)
    if curCol != indentStack[-1]:
        if (not indentStack) or curCol > indentStack[-1]:
            raise ParseFatalException(s,l,"illegal nesting")
        raise ParseException(s,l,"not a peer entry")

def checkSubIndent(s,l,t):
    curCol = col(l,s)
    if curCol > indentStack[-1]:
        indentStack.append( curCol )
    else:
        raise ParseException(s,l,"not a subentry")

def checkUnindent(s,l,t):
    if l >= len(s): return
    curCol = col(l,s)
    if not(curCol < indentStack[-1] and curCol <= indentStack[-2]):
        raise ParseException(s,l,"not an unindent")

def doUnindent():
    indentStack.pop()

INDENT = lineEnd.suppress() + empty + empty.copy().setParseAction(checkSubIndent)
UNDENT = FollowedBy(empty).setParseAction(checkUnindent)
UNDENT.setParseAction(doUnindent)

stmt = Forward()
suite = OneOrMore(empty + stmt.setParseAction(checkPeerIndent))

register = Register.grammar(Word(ascii_uppercase, exact=1))
_nameish = Word(alphas + '_', alphanums + '_')
name = Name.grammar(delimitedList(_nameish, '.', True))
func_decl = Suppress("def") + _nameish + Suppress(":")
func_def = FunctionDefinition.grammar(func_decl + INDENT + suite + UNDENT)
alias_identifier = Alias.grammar(_nameish)

maindef = Main.grammar(Suppress('main:') + INDENT + suite + UNDENT)

string = String.grammar(dblQuotedString) | String.grammar(sglQuotedString)

numeric = Hex.grammar(Combine('0x' + Word(hexnums, exact=4))) | Integer.grammar(Word(nums))

constant = string | numeric

memloc = MemoryLocation.grammar(Suppress('[') + (register | numeric) + Optional('+' + (register | numeric)) + Suppress(']'))

func_call = FunctionCall.grammar(name + Suppress("()"))
builtin_names = Forward()
for key in BUILTINS:
    builtin_names |= Word(key)
builtin_call = BuiltinCall.grammar(builtin_names + Suppress("(") + Optional(delimitedList(constant)) + Suppress(")"))

operand = builtin_call | func_call | constant

expr = Expression.grammar(operatorPrecedence(operand, [
    (Literal('<<'), 2, opAssoc.LEFT),
    (Literal('|'), 2, opAssoc.LEFT),
]))

for_decl = Suppress('for') + Group(delimitedList(register)) + Suppress('in') + (string | alias_identifier) + Suppress(':')
for_loop = ForLoop.grammar(for_decl + INDENT + Group(suite) + UNDENT)

math_operator = Literal('+') | Literal('-') | Literal('*') | Literal('/') | Literal('|') | Literal('&') | Literal('<<') | Literal('>>')

rvalue = (alias_identifier | memloc | numeric | expr)
assignment = Assignment.grammar(register + Suppress("=") + rvalue)
aug_assignment = AugmentedAssignment.grammar(register + math_operator + Suppress('=') + rvalue)
mem_assignment = MemoryAssignment.grammar(memloc + Suppress('=') + (register | numeric))
stmt << (assignment | aug_assignment | mem_assignment | for_loop | func_call)

alias_rvalue = (builtin_call | name | numeric | expr)

aliasdef = AliasDefinition.grammar(alias_identifier + Suppress("=") + alias_rvalue)

import_def = Import.grammar(Suppress('import') + Group(delimitedList(Word(alphas, alphanums), '.')))

data_rvalue = delimitedList(alias_rvalue | string, ',')

datadef = Data.grammar(alias_identifier + Suppress("= data") + data_rvalue)

grammar = Module.grammar(ZeroOrMore(import_def) + ZeroOrMore(datadef | aliasdef) + Optional(maindef) + ZeroOrMore(func_def))