import unittest
import xml.etree.ElementTree as ET

from ast_generator.ast_generator import *
from ast_generator.gazprea_ast_grammar import *


class TestGeneration(unittest.TestCase):

    def setUp(self):
        self.ast_gen = AstGenerator(GAZPREA_TOP_LEVEL, json.loads("{}"))

    def test_generate_literal(self):
        out: ET.Element = self.ast_gen.generate_literal()

        self.assertIsNotNone(out.get("type"))

    def test_generate_variable(self):
        out: ET.Element = self.ast_gen.generate_variable()

        self.assertIsNotNone(out.get("type"))
        self.assertIsNotNone(out.get("name"))

    def test_generate_declaration(self):
        out: ET.Element = self.ast_gen.generate_declaration()

        self.assertIsNotNone(out.find("variable"))
        self.assertIsNotNone(out.find("rhs"))

    def test_generate_operation(self):
        out: ET.Element = self.ast_gen.generate_operator()

        self.assertIsNotNone(out.find("lhs"))
        self.assertIsNotNone(out.find("rhs"))

    def test_generate_stream(self):
        out: ET.Element = self.ast_gen.generate_stream()

        self.assertIsNotNone(out.get("type"))

    def test_generate_routine(self):
        out: ET.Element = self.ast_gen.generate_routine()

        self.assertIsNotNone(out.find("block"))

    def test_generate_block(self):
        out: ET.Element = self.ast_gen.generate_block()

        self.assertIsNotNone(out)

    def test_generate_main(self):
        out: ET.Element = self.ast_gen.generate_main_routine()

        self.assertIsNotNone(out)
        self.assertEqual("main", out.get("name"))
        self.assertEqual("int", out.get("return_type"))

    def test_generate_ast(self):
        out: ET.Element = self.ast_gen.generate_ast()

        self.assertIsNotNone(out)

        has_main = False
        for child in out:
            if child.get("name") == "main":
                has_main = True
                break

        self.assertTrue(has_main)





class TestRoutines(unittest.TestCase):

    def setUp(self):
        self.ast_gen = AstGenerator(GAZPREA_TOP_LEVEL, json.loads("{}"))

    def test_populate_function(self):
        with open("xml/routine_1.xml", 'r') as i:
            input_elemnet = ET.fromstring(i.read())
            out: ET.Element = self.ast_gen.populate_routine(input_elemnet, 'int', 'blahaj')

            self.assertEqual('int', out.get("return_type"))
            self.assertEqual('int', out.find("block").find("return").get("type"))

    def test_generate_return(self):
        out: ET.Element = self.ast_gen.generate_return("int")

        self.assertEqual('int', out.get("type"))

    def test_populate_typed_procedure(self):
        with open("xml/routine_2.xml", 'r') as i:
            input_elemnet = ET.fromstring(i.read())
            out: ET.Element = self.ast_gen.populate_routine(input_elemnet, 'int', 'blahaj')

            self.assertEqual('int', out.get("return_type"))
            self.assertEqual('int', out.find("block").find("return").get("type"))

    def test_populate_void_procedure(self):
        with open("xml/routine_3.xml", 'r') as i:
            input_elemnet = ET.fromstring(i.read())
            out: ET.Element = self.ast_gen.populate_routine(input_elemnet, 'void', 'blahaj')

            self.assertEqual('void', out.get("return_type"))
            self.assertIsNone(out.find("block").find("return"))

    def test_populate_void_procedure_with_return(self):
        with open("xml/routine_2.xml", 'r') as i:
            input_elemnet = ET.fromstring(i.read())
            out: ET.Element = self.ast_gen.populate_routine(input_elemnet, 'void', 'blahaj')

            self.assertEqual('void', out.get("return_type"))
            self.assertIsNotNone(out.find("block").find("return").get("type"))
            self.assertEqual('void', out.find("block").find("return").get("type"))


    def test_populate_nonvoid_procedure_without_return(self):
        with open("xml/routine_3.xml", 'r') as i:
            input_elemnet = ET.fromstring(i.read())
            out: ET.Element = self.ast_gen.populate_routine(input_elemnet, 'int', 'blahaj')

            self.assertEqual('int', out.get("return_type"))
            self.assertIsNotNone(out.find("block").find("return"))
            self.assertEqual('int', out.find("block").find("return").get("type"))




    # def test_populate_assignment(self):
    #     pass
    #
    # def test_populate_top_block(self):
    #     pass


    # def test_populate_call(self):
    #     pass




class TestValues(unittest.TestCase):

    def setUp(self):
        self.ast_gen = AstGenerator(GAZPREA_TOP_LEVEL ,json.loads("{}"))

    def test_populate_int(self):
        with open("xml/literal_1.xml", 'r') as i:
            input = i.read()
            input_element = ET.fromstring(input)
            out: ET.Element = self.ast_gen.populate_literal(input_element, "int", 42)

            self.assertEqual("int", out.get("type"))
            self.assertEqual("42", str(out.get("value")))

    def test_populate_var(self):
        with open("xml/variable_1.xml", 'r') as i:
            input = i.read()
            input_element = ET.fromstring(input)
            out: ET.Element = self.ast_gen.populate_variable(input_element, "var", "float", "bean")

            self.assertEqual("var", out.get("mut"))
            self.assertEqual("float", out.get("type"))
            self.assertEqual("bean", out.get("name"))

    def test_populate_operator(self):
        with open("xml/operator_1.xml", 'r') as i:
            input = i.read()
            input_element = ET.fromstring(input)
            out: ET.Element = self.ast_gen.populate_operator(input_element, '*', 'int')
            inner_var = out.find("type")

            self.assertEqual("int", out.get("type"))
            self.assertEqual("*", out.get("op"))
            self.assertEqual("int", out.find("lhs").find("variable").get("type"))
            self.assertEqual("int", out.find("rhs").find("literal").get("type"))

    def test_populate_stream(self):
        with open("xml/stream_1.xml", 'r') as i:
            input_element = ET.fromstring(i.read())
            out: ET.Element = self.ast_gen.populate_stream(input_element, "std_output")

            self.assertEqual("std_output", out.get("type"))






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