import unittest

from h8536.dataflow import analyze_dataflow, state_for_instruction
from h8536.model import Instruction


def reg_after(analysis, address, register):
    return analysis["instructions"][address]["after"]["registers"][register]


def reg_before(analysis, address, register):
    return analysis["instructions"][address]["before"]["registers"][register]


def control_after(analysis, address, register):
    return analysis["instructions"][address]["after"]["control"][register]


class DataflowTest(unittest.TestCase):
    def test_tracks_immediate_load_copy_and_simple_arithmetic(self):
        instructions = {
            0x0100: Instruction(0x0100, b"\x58\x02\x00", "MOV:I.W", "#H'0200, R0"),
            0x0103: Instruction(0x0103, b"\xA0\x81", "MOV:G.W", "R0, R1"),
            0x0105: Instruction(0x0105, b"\xA1\x08", "ADD:Q.W", "#1, R1"),
            0x0107: Instruction(0x0107, b"\x0C\x00\x02\x31", "SUB.W", "#H'0002, R1"),
        }

        analysis = analyze_dataflow(instructions)

        self.assertEqual(reg_after(analysis, 0x0100, "R0")["value"], 0x0200)
        self.assertEqual(reg_after(analysis, 0x0100, "R0")["width"], 16)
        self.assertEqual(reg_before(analysis, 0x0103, "R0")["value"], 0x0200)
        self.assertEqual(reg_after(analysis, 0x0103, "R1")["value"], 0x0200)
        self.assertEqual(reg_after(analysis, 0x0105, "R1")["value"], 0x0201)
        self.assertEqual(reg_after(analysis, 0x0107, "R1")["value"], 0x01FF)

    def test_tracks_byte_immediates_without_promising_word_width(self):
        instructions = {
            0x0200: Instruction(0x0200, b"\x52\x7F", "MOV:E.B", "#H'7F, R2"),
            0x0202: Instruction(0x0202, b"\xA2\x83", "MOV:G.B", "R2, R3"),
            0x0204: Instruction(0x0204, b"\x58\x20\x00", "MOV:I.W", "#H'2000, R0"),
            0x0207: Instruction(0x0207, b"\xD0\x84", "MOV:G.W", "@R0, R4"),
        }

        analysis = analyze_dataflow(instructions)

        self.assertEqual(reg_after(analysis, 0x0200, "R2")["value"], 0x7F)
        self.assertEqual(reg_after(analysis, 0x0200, "R2")["width"], 8)
        self.assertEqual(reg_after(analysis, 0x0202, "R3")["value"], 0x7F)
        self.assertEqual(reg_after(analysis, 0x0202, "R3")["width"], 8)
        self.assertEqual(reg_after(analysis, 0x0207, "R0")["value"], 0x2000)
        self.assertFalse(reg_after(analysis, 0x0207, "R4")["known"])
        self.assertEqual(reg_after(analysis, 0x0207, "R4")["reason"], "memory_load")

    def test_calls_and_ambiguous_branches_do_not_leak_known_state(self):
        instructions = {
            0x0300: Instruction(0x0300, b"\x58\x12\x34", "MOV:I.W", "#H'1234, R0"),
            0x0303: Instruction(0x0303, b"\x26\x03", "BNE", "loc_0308", kind="branch", targets=[0x0308]),
            0x0305: Instruction(0x0305, b"\xA0\x08", "ADD:Q.W", "#1, R0"),
            0x0308: Instruction(0x0308, b"\xA0\x08", "ADD:Q.W", "#1, R0"),
            0x030A: Instruction(0x030A, b"\x18\x04\x00", "JSR", "@loc_0400", kind="call", targets=[0x0400]),
            0x030D: Instruction(0x030D, b"\xA0\x08", "ADD:Q.W", "#1, R0"),
        }

        analysis = analyze_dataflow(instructions)

        self.assertFalse(reg_before(analysis, 0x0305, "R0")["known"])
        self.assertEqual(reg_before(analysis, 0x0305, "R0")["reason"], "block_entry")
        self.assertFalse(reg_before(analysis, 0x0308, "R0")["known"])
        self.assertEqual(reg_before(analysis, 0x0308, "R0")["reason"], "block_entry")
        self.assertFalse(reg_after(analysis, 0x030A, "R0")["known"])
        self.assertEqual(reg_after(analysis, 0x030A, "R0")["reason"], "call")
        self.assertFalse(reg_before(analysis, 0x030D, "R0")["known"])

    def test_tracks_control_register_loads_and_stc_copies(self):
        instructions = {
            0x0400: Instruction(
                0x0400,
                b"\x04\xFE\x89",
                "LDC.B",
                "#H'FE, BR",
                writes_br=True,
                br_value=0xFE,
            ),
            0x0403: Instruction(0x0403, b"\xA0\x99", "STC.B", "BR, R1"),
            0x0405: Instruction(0x0405, b"\x04\x01\x48", "ORC.B", "#H'01, CCR"),
        }

        analysis = analyze_dataflow(instructions)

        self.assertEqual(control_after(analysis, 0x0400, "BR")["value"], 0xFE)
        self.assertEqual(control_after(analysis, 0x0400, "BR")["width"], 8)
        self.assertEqual(reg_after(analysis, 0x0403, "R1")["value"], 0xFE)
        self.assertFalse(control_after(analysis, 0x0405, "CCR")["known"])

    def test_state_lookup_helper_returns_instruction_record(self):
        instructions = {
            0x0500: Instruction(0x0500, b"\x58\x00\x01", "MOV:I.W", "#H'0001, R0"),
        }

        analysis = analyze_dataflow(instructions)

        self.assertEqual(state_for_instruction(analysis, 0x0500)["after"]["registers"]["R0"]["value"], 1)
        self.assertEqual(state_for_instruction(analysis, 0x9999), {})


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