#!/usr/bin/env python3

import sys

from jack_tokenizer import JackTokenizer
from jack_tokenizer import TOKEN_TYPE_MAP
from symbol_table import SymbolTable
from vm_writer import VMWriter

OPS = ['+', '-', '*', '/', '&', '|', '<', '>', '=']
UNARY_OPS = ['-', '~']
KIND_KEYWORD_MAP = {
    'STATIC': 'STATIC',
    'FIELD': 'THIS',
    'ARG': 'ARG',
    'VAR': 'LOCAL',
}

class CompilationEngine():
    def __init__(self, path, out_f, vm_path):
        self._tokenizer = JackTokenizer(path)
        self._f = out_f
        self._classname = ''
        self._funcname = ''
        self._tbl = SymbolTable()
        self._writer = VMWriter(vm_path)
        self._while_index = 0
        self._if_index = 0

    def _writeLines(self, lines):
        for line in lines:
            self._f.write(line + '\n')

    def _tokenize(self, lines):
        self._tokenizer.advance()
        token_type = self._tokenizer.tokenType()
        tag = TOKEN_TYPE_MAP[token_type]
        tag_begin = '<{}>'.format(tag)
        tag_end = '</{}>'.format(tag)
        tag_word = ''
        if token_type == 'KEYWORD':
            word = self._tokenizer.keyWord().lower()
        elif token_type == 'SYMBOL':
            word = self._tokenizer.symbol()
            if word == '<':
                tag_word = '&lt;'
            elif word == '>':
                tag_word = '&gt;'
            elif word == '&':
                tag_word = '&amp;'
        elif token_type == 'IDENTIFIER':
            word = self._tokenizer.identifier()
        elif token_type == 'INT_CONST':
            word = str(self._tokenizer.intVal())
        elif token_type == 'STRING_CONST':
            word = self._tokenizer.stringVal()
        lines.append(' '.join([tag_begin, tag_word if tag_word else word, tag_end]))
        return word

    def compileClass(self):
        lines = []
        lines.append('<class>')
        while self._tokenizer.hasMoreTokens():
            next_token = self._tokenizer.peekNextToken()
            next_word = next_token[0]
            if next_word == 'class':
                self._tokenize(lines) # class
                self._classname = self._tokenize(lines) # classname
            elif next_word == 'static' or next_word == 'field':
                self.compileClassVarDec(lines)
            elif next_word == 'constructor' or next_word == 'function' or next_word == 'method':
                self.compileSubroutine(lines)
            elif self._tokenize(lines) == '}':
                break
        lines.append('</class>')
        self._writeLines(lines)

    def compileClassVarDec(self, lines):
        lines.append('<classVarDec>')
        kind = self._tokenize(lines).upper()
        _type = self._tokenize(lines)
        while self._tokenizer.hasMoreTokens():
            word = self._tokenize(lines)
            if word == ';':
                break
            elif word != ',':
                if not self._tbl.kindOf(word):
                    self._tbl.define(word, _type, kind)
        lines.append('</classVarDec>')

    def compileDo(self, lines):
        lines.append('<doStatement>')
        self._tokenize(lines)   # do
        wordcat = ''
        funcname = ''
        n_exprs = 0
        is_obj = False
        while self._tokenizer.hasMoreTokens():
            word = self._tokenize(lines)
            if word == '(':
                if '.' in wordcat:
                    funcname = wordcat
                    if is_obj:
                        # outer method (obj.hoge())
                        if kind == 'STATIC':
                            segment = 'STATIC'
                        elif kind == 'FIELD':
                            segment = 'THIS'
                        elif kind == 'ARG':
                            segment = 'ARG'
                        elif kind == 'VAR':
                            segment = 'LOCAL'
                        self._writer.writePush(segment, index)
                        n_exprs = n_exprs + 1
                    else:
                        # function (Class.hoge())
                        pass
                else:
                    # inner method (hoge())
                    funcname = self._classname + '.' + wordcat
                    self._writer.writePush('POINTER', 0)
                    n_exprs = n_exprs + 1
                n_exprs = n_exprs + self.compileExpressionList(lines)
            elif word == ';':
                break
            if self._tbl.typeOf(word):
                kind = self._tbl.kindOf(word)
                index = self._tbl.indexOf(word)
                word = self._tbl.typeOf(word)
                is_obj = True
            wordcat += word
        self._writer.writeCall(funcname, n_exprs)
        self._writer.writePop('TEMP', 0)
        lines.append('</doStatement>')

    def compileLet(self, lines):
        lines.append('<letStatement>')
        self._tokenize(lines)   # let
        varname = self._tokenize(lines) # varName
        is_array = False
        while self._tokenizer.hasMoreTokens():
            word = self._tokenize(lines)
            if word == '=':
                self.compileExpression(lines)
            elif word == '[':
                is_array = True
                self.compileExpression(lines)
                self._writer.writePush(KIND_KEYWORD_MAP[self._tbl.kindOf(varname)], self._tbl.indexOf(varname))
                self._writer.writeArithmetic('add')
            elif word == ';':
                break
        if is_array:
            self._writer.writePop('TEMP', 0)
            self._writer.writePop('POINTER', 1)
            self._writer.writePush('TEMP', 0)
            self._writer.writePop('THAT', 0)
        else:
            self._writer.writePop(KIND_KEYWORD_MAP[self._tbl.kindOf(varname)], self._tbl.indexOf(varname))
        lines.append('</letStatement>')

    def compileWhile(self, lines):
        lines.append('<whileStatement>')
        label_exp = 'WHILE_EXP{}'.format(str(self._while_index))
        label_end = 'WHILE_END{}'.format(str(self._while_index))
        self._while_index += 1
        self._writer.writeLabel(label_exp)
        while self._tokenizer.hasMoreTokens():
            word = self._tokenize(lines)
            if word == '(':
                self.compileExpression(lines)
                self._writer.writeArithmetic('not')
                self._writer.writeIf(label_end)
            elif word == '{':
                self.compileStatements(lines)
                self._tokenize(lines)
                break
        self._writer.writeGoto(label_exp)
        self._writer.writeLabel(label_end)
        lines.append('</whileStatement>')

    def compileIf(self, lines):
        lines.append('<ifStatement>')
        label_true = 'IF_TRUE{}'.format(str(self._if_index))
        label_false = 'IF_FALSE{}'.format(str(self._if_index))
        label_end = 'IF_END{}'.format(str(self._if_index))
        self._if_index += 1
        while self._tokenizer.hasMoreTokens():
            word = self._tokenize(lines)
            if word == '(':
                self.compileExpression(lines)
                self._writer.writeIf(label_true)
                self._writer.writeGoto(label_false)
            elif word == '{':
                self._writer.writeLabel(label_true)
                self.compileStatements(lines)
                self._tokenize(lines)
                break
        next_token = self._tokenizer.peekNextToken()
        next_word = next_token[0]
        if next_word == 'else':
            self._writer.writeGoto(label_end)
            self._tokenize(lines) # else
            self._tokenize(lines) # {
            self._writer.writeLabel(label_false)
            self.compileStatements(lines)
            self._tokenize(lines) # }
            self._writer.writeLabel(label_end)
        else:
            self._writer.writeLabel(label_false)
        lines.append('</ifStatement>')

    def compileReturn(self, lines):
        lines.append('<returnStatement>')
        compiled_expression = False
        while self._tokenizer.hasMoreTokens():
            next_token = self._tokenizer.peekNextToken()
            next_word = next_token[0]
            if next_word == 'return' or next_word == ';':
                word = self._tokenize(lines)
                if word == ';':
                    break
            else:
                self.compileExpression(lines)
                compiled_expression = True
        if not compiled_expression:
            self._writer.writePush('CONST', 0)
        self._writer.writeReturn()
        lines.append('</returnStatement>')

    def compileExpressionList(self, lines):
        lines.append('<expressionList>')
        n_exprs = 0
        while self._tokenizer.hasMoreTokens():
            next_token = self._tokenizer.peekNextToken()
            next_word = next_token[0]
            if next_word == ')':
                break
            elif next_word == ',':
                self._tokenize(lines)
            self.compileExpression(lines)
            n_exprs += 1
        lines.append('</expressionList>')
        return n_exprs

    def _writeOpFunction(self, op=None, unary_op=None):
        if op == '+':
            self._writer.writeArithmetic('add')
        elif op == '-':
            self._writer.writeArithmetic('sub')
        elif op == '*':
            self._writer.writeCall('Math.multiply', 2)
        elif op == '/':
            self._writer.writeCall('Math.divide', 2)
        elif op == '&':
            self._writer.writeArithmetic('and')
        elif op == '|':
            self._writer.writeArithmetic('or')
        elif op == '<':
            self._writer.writeArithmetic('lt')
        elif op == '>':
            self._writer.writeArithmetic('gt')
        elif op == '=':
            self._writer.writeArithmetic('eq')
        elif unary_op == '-':
            self._writer.writeArithmetic('neg')
        elif unary_op == '~':
            self._writer.writeArithmetic('not')

    def compileExpression(self, lines):
        lines.append('<expression>')
        is_next_op = self.compileTerm(lines)
        has_op = is_next_op
        while is_next_op:
            word = self._tokenize(lines) # op
            is_next_op = self.compileTerm(lines)
        if has_op:
            self._writeOpFunction(word)
        lines.append('</expression>')

    def compileTerm(self, lines):
        is_next_op = False
        lines.append('<term>')
        is_first_word = True
        wordcat = ''
        is_call_sub = False
        is_array = False
        n_exprs = 0
        is_obj = False
        while self._tokenizer.hasMoreTokens():
            next_token = self._tokenizer.peekNextToken()
            next_word = next_token[0]
            if is_first_word:
                if next_word == '(':
                    self._tokenize(lines) # (
                    self.compileExpression(lines)
                    self._tokenize(lines) # )
                elif next_word in UNARY_OPS:
                    word = self._tokenize(lines)
                    self.compileTerm(lines)
                    self._writeOpFunction(op=None, unary_op=word)
            else:
                if next_word == ')' or next_word == ';' or next_word == ',' or next_word == ']':
                    break
                elif next_word in OPS:
                    is_next_op = True
                    break
                word = self._tokenize(lines)
                if word == '(':
                    is_call_sub = True
                    if '.' in wordcat:
                        obj = wordcat[:wordcat.find('.')]
                        _type = self._tbl.typeOf(obj)
                        if _type:
                            is_obj = True
                            wordcat = wordcat.replace(obj, _type)
                            n_exprs = n_exprs + 1
                    n_exprs = n_exprs + self.compileExpressionList(lines)
                    self._tokenize(lines) # ')'
                    break
                elif word == '[':
                    self.compileExpression(lines)
                    self._tokenize(lines) # ']'
                    next_token = self._tokenizer.peekNextToken()
                    next_word = next_token[0]
                    if next_word in OPS:
                        is_next_op = True
                    is_array = True
                    break
                wordcat += word
            is_first_word = False
        if wordcat.isdigit():
            self._writer.writePush('CONST', wordcat)
        elif is_call_sub:
            if is_obj:
                self._writer.writePush(KIND_KEYWORD_MAP[self._tbl.kindOf(obj)], self._tbl.indexOf(obj))
            self._writer.writeCall(wordcat, n_exprs)
        elif wordcat != '':
            if wordcat == 'true':
                self._writer.writePush('CONST', '0')
                self._writer.writeArithmetic('not')
            elif wordcat == 'false':
                self._writer.writePush('CONST', '0')
            elif wordcat == 'null':
                self._writer.writePush('CONST', '0')
            elif wordcat == 'this':
                self._writer.writePush('POINTER', '0')
            else:
                if self._tbl.kindOf(wordcat):
                    # symbol
                    self._writer.writePush(KIND_KEYWORD_MAP[self._tbl.kindOf(wordcat)], self._tbl.indexOf(wordcat))
                    if is_array:
                        self._writer.writeArithmetic('add')
                        self._writer.writePop('POINTER', 1)
                        self._writer.writePush('THAT', 0)
                else:
                    # string literal
                    self._writer.writePush('CONST', len(wordcat))
                    self._writer.writeCall('String.new', 1)
                    for c in wordcat:
                        self._writer.writePush('CONST', ord(c))
                        self._writer.writeCall('String.appendChar', 2)
        lines.append('</term>')
        return is_next_op

    def compileStatements(self, lines):
        lines.append('<statements>')
        while self._tokenizer.hasMoreTokens():
            next_token = self._tokenizer.peekNextToken()
            next_word = next_token[0]
            if next_word == 'do':
                self.compileDo(lines)
            elif next_word == 'let':
                self.compileLet(lines)
            elif next_word == 'while':
                self.compileWhile(lines)
            elif next_word == 'if':
                self.compileIf(lines)
            elif next_word == 'return':
                self.compileReturn(lines)
                break
            elif next_word == '}':
                break
            else:
                self._tokenize(lines)
        lines.append('</statements>')

    def _compileVarDec(self, lines):
        lines.append('<varDec>')
        self._tokenize(lines)   # var
        _type = self._tokenize(lines)   # type
        while self._tokenizer.hasMoreTokens():
            word = self._tokenize(lines)
            if word == ';':
                break
            elif word != ',':
                self._tbl.define(word, _type, 'VAR')
        lines.append('</varDec>')

    def compileParameterList(self, lines):
        lines.append('<parameterList>')
        while self._tokenizer.hasMoreTokens():
            next_token = self._tokenizer.peekNextToken()
            next_word = next_token[0]
            if next_word == ')':
                break
            word = self._tokenize(lines)
            if word != ',':
                _type = word
                name = self._tokenize(lines)
                self._tbl.define(name, _type, 'ARG')
        lines.append('</parameterList>')

    def _compileSubroutineBody(self, lines, method_type):
        lines.append('<subroutineBody>')
        self._tokenize(lines)
        while self._tokenizer.hasMoreTokens():
            next_token = self._tokenizer.peekNextToken()
            next_word = next_token[0]
            if next_word == 'var':
                self._compileVarDec(lines)
            elif next_word == '}':
                self._tokenize(lines)
                break
            else:
                name = '{}.{}'.format(self._classname, self._funcname)
                self._writer.writeFunction(name, self._tbl.varCount('VAR'))
                if method_type == 'constructor':
                    self._writer.writePush('CONST', self._tbl.varCount('FIELD'))
                    self._writer.writeCall('Memory.alloc', 1)
                    self._writer.writePop('POINTER', 0)
                elif method_type == 'method':
                    self._writer.writePush('ARG', 0)
                    self._writer.writePop('POINTER', 0)
                elif method_type == 'function':
                    pass
                self.compileStatements(lines)
        lines.append('</subroutineBody>')

    def compileSubroutine(self, lines):
        lines.append('<subroutineDec>')
        self._tbl.startSubroutine()
        self._while_index = 0
        self._if_index = 0
        word = self._tokenize(lines)   # constructor/function/method
        if word == 'constructor' or word == 'method' or word == 'function':
            method_type = word
        if method_type == 'method':
            self._tbl.define('this', self._classname, 'ARG')
        self._tokenize(lines)   # type
        self._funcname = self._tokenize(lines)   # function
        while self._tokenizer.hasMoreTokens():
            next_token = self._tokenizer.peekNextToken()
            next_word = next_token[0]
            if next_word == '(':
                self._tokenize(lines)
                self.compileParameterList(lines)
            elif next_word == '{':
                self._compileSubroutineBody(lines, method_type)
                break
            else:
                self._tokenize(lines)
        lines.append('</subroutineDec>')


def main():
    input = sys.argv[1]
    output = input.replace('.jack', '.xml')
    vm_path = input.replace('.jack', '.vm')
    with open(output, 'w') as f:
        engine = CompilationEngine(input, f, vm_path)
        engine.compileClass()

if __name__ == "__main__":
    main()
