#!/usr/bin/python

import unittest, cProfile
from ppci import CompilerError
from ppci.asmnodes import AInstruction, ABinop, AUnop, ASymbol, ALabel, ANumber
from ppci.assembler import tokenize, Assembler
from ppci.objectfile import ObjectFile
from ppci.linker import Linker
import outstream
from target import Label


class AssemblerLexingCase(unittest.TestCase):
    """ Tests the assemblers lexer """

    def testLex0(self):
        """ Check if the lexer is OK """
        asmline, toks = 'mov rax, rbx ', ['ID', 'ID', ',', 'ID', 'EOF']
        self.assertSequenceEqual([tok.typ for tok in tokenize(asmline)], toks)

    def testLex1(self):
        """ Test if lexer correctly maps some tokens """
        asmline, toks = 'lab1: mov rax, rbx ', ['ID', ':', 'ID', 'ID', ',', 'ID', 'EOF']
        self.assertSequenceEqual([tok.typ for tok in tokenize(asmline)], toks)

    def testLex2(self):
        """ Test if lexer correctly maps some tokens """
        asmline, toks = 'mov 3.13 0xC 13', ['ID', 'REAL', 'NUMBER', 'NUMBER', 'EOF']
        self.assertSequenceEqual([tok.typ for tok in tokenize(asmline)], toks)

    def testLex3(self):
        """ Test if lexer fails on a token that is invalid """
        asmline = '0z4: mov rax, rbx $ '
        with self.assertRaises(CompilerError):
            list(tokenize(asmline))


class AssemblerParsingTestCase(unittest.TestCase):
    """ 
        Tests the assembler parts
    """
    def setUp(self):
        self.a = Assembler()

    def testParse(self):
        asmline = 'lab1: mov rax, rbx'
        self.a.parse_line(asmline)

    def expectTree(self, asmline, stack):
        self.a.parse_line(asmline)
        self.assertSequenceEqual(stack, self.a.stack)

    def testParse2(self):
        asmline = 'a: mov rax, [rbx + 2]'
        output = []
        output.append(ALabel('a'))
        output.append(AInstruction('mov', [ASymbol('rax'), AUnop('[]', ASymbol('rbx') + ANumber(2))]))
        self.expectTree(asmline, output)

    def testParse3(self):
        # A label must be optional:
        asmline = 'mov rax, 1'
        output = [AInstruction('mov', [ASymbol('rax'), ANumber(1)])]
        self.expectTree(asmline, output)

    def testParse4(self):
        # Test 3 operands:
        asmline = 'add rax, [4*rbx + 22], rcx'
        ops = []
        ops.append(ASymbol('rax'))
        ops.append(AUnop('[]', ANumber(4) * ASymbol('rbx') + ANumber(22)))
        ops.append(ASymbol('rcx'))
        output = [AInstruction('add', ops)]
        self.expectTree(asmline, output)

    def testParse5(self):
        # An instruction must be optional:
        asmline = 'lab1:'
        output = []
        output.append(ALabel('lab1'))
        self.expectTree(asmline, output)

    def testParse6(self):
        # A line can be empty
        self.a.parse_line('')


class AssemblerOtherTestCase(unittest.TestCase):
    def testWithoutTarget(self):
        a = Assembler()
        with self.assertRaises(CompilerError):
            a.assemble_line('')


class OustreamTestCase(unittest.TestCase):
    def test1(self):
        obj = ObjectFile()
        o = outstream.BinaryOutputStream(obj)
        o.selectSection('.text')
        o.emit(Label('a'))
        self.assertSequenceEqual(bytes(), obj.get_section('.text').data)


class AsmTestCaseBase(unittest.TestCase):
    """ Base testcase for assembly """
    def feed(self, line):
        self.a.assemble(line)

    def check(self, hexstr):
        l = Linker()
        self.obj = l.link([self.obj])
        data = bytes(self.obj.get_section('.text').data)
        self.assertSequenceEqual(bytes.fromhex(hexstr), data)


if __name__ == '__main__':
    unittest.main()