#!/usr/bin/python

import unittest, cProfile
from ppci import CompilerError
from asmnodes import AInstruction, ABinop, AUnop, ASymbol, ALabel, ANumber
from asm import tokenize, Assembler
import msp430
import arm_cm3

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']
        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']
        self.assertSequenceEqual([tok.typ for tok in tokenize(asmline)], toks)

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

    def testLex2(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 testParse2(self):
        asmline = 'a: mov rax, [rbx + 2]'
        self.a.parse_line(asmline)
        output = []
        output.append(ALabel('a'))
        output.append(AInstruction('mov', [ASymbol('rax'), AUnop('[]', ASymbol('rbx') + ANumber(2))]))
        self.assertSequenceEqual(output, self.a.output)

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

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

    def testParse5(self):
        # An instruction must be optional:
        asmline = 'lab1:'
        self.a.parse_line(asmline)
        output = []
        output.append(ALabel('lab1'))
        self.assertSequenceEqual(output, self.a.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('')

    @unittest.skip 
    def testX86(self):
        testsrc = """ ; tst
        begin:
        mov rax, rbx ; 0x48, 0x89, 0xd8
        xor rcx, rbx ; 0x48, 0x31, 0xd9
        inc rcx ; 0x48 0xff 0xc1
        """
        a = Assembler()
        a.assemble(testsrc)
        # Compare with nasm output:
        nasmbytes = [0x48, 0x89, 0xd8, 0x48, 0x31, 0xd9, 0x48, 0xff, 0xc1]

class AssemblerMSP430TestCase(unittest.TestCase):
    def setUp(self):
        self.t = msp430.msp430target
        self.a = Assembler(target=self.t)

    def testMapMovInstruction(self):
        i = AInstruction('mov', [ASymbol('r14'), ASymbol('r15')])
        ri = self.t.mapInstruction(i)

    def testMapRetiInstruction(self):
        i = AInstruction('reti', [])
        ri = self.t.mapInstruction(i)

    @unittest.skip
    def testMapOperand(self):
        o = ASymbol('r14')
        mo = self.t.mapOperand(o)
        self.assertEqual(mo, msp430.r14)

    @unittest.skip
    def testMapOperandIndirection(self):
        o = AUnop('[]', ASymbol('r14'))
        mo = self.t.mapOperand(o)

    def testMov(self):
        line1 = "mov r14, r15"
        self.a.assemble_line(line1)
        self.assertEqual(bytes([0x0F, 0x4E]), self.a.binout)

    def testMov1337(self):
        line1 = "mov 0x1337, r12"
        self.a.assemble_line(line1)
        self.assertEqual(bytes.fromhex('3C403713'), self.a.binout)

    def testAdd(self):
        line1 = "add r15, r13"
        self.a.assemble_line(line1)
        self.assertEqual(bytes.fromhex('0D5F'), self.a.binout)

    def testReti(self):
        line1 = "reti"
        self.a.assemble_line(line1)
        self.assertEqual(bytes([0x0, 0x13]), self.a.binout)

    def testMSPinstructionCount(self):
        """ Check that there are 27 instructions """
        self.assertEqual(27, len(self.t.instructions))


class AssemblerARMTestCase(unittest.TestCase):
    def setUp(self):
        self.t = arm_cm3.armtarget
        self.a = Assembler(target=self.t)

    def feed(self, line):
        self.a.assemble(line)

    def check(self, hexstr):
        self.assertEqual(bytes.fromhex(hexstr), self.a.binout)
        
    def testMapOperand(self):
        pass

    def testMovImm8(self):
        self.feed('mov r4, 100')
        self.check('6424')

    def testYield(self):
        self.feed('yield')
        self.check('10bf')

    def testPush(self):
        self.feed('push {r2,r3,lr}')
        self.check('0cb5')

    def testPop(self):
        self.feed('pop {r4-r6, pc}')
        self.check('70bd')

    def testStr5(self):
        self.feed('str r4, [r1 + 0]')
        self.check('0c60')

    def testLdr5(self):
        self.feed('ldr r4, [r0 + 0]')
        self.check('0468')

    def testSequence1(self):
        self.feed('mov r5, 3')
        self.feed('add r4, r5, 0')
        self.feed('loop: add r6, r4, 7')
        self.feed('cmp r6, 5')
        self.check('0325 2c1c e61d 052e')

if __name__ == '__main__':
    #cProfile.run('unittest.main()')
    unittest.main()