import cocotb
from cocotb.handle import Immediate

from cocotb.clock import Clock
from cocotb.triggers import Timer, RisingEdge, FallingEdge

from collections import defaultdict

import struct
import random

CLK_PERIOD = 5

memory = defaultdict(int)

def write_dword(addr: int, data: int):
    memory[addr + 0] = (data >>  0) & 0xff
    memory[addr + 1] = (data >>  8) & 0xff
    memory[addr + 2] = (data >> 16) & 0xff
    memory[addr + 3] = (data >> 24) & 0xff

def write_byte(addr: int, data: int):
    memory[addr] = data & 0xff

def write_bytes(addr: int, data: bytes| list[int]):
    for i, val in enumerate(data):
        memory[addr + i] = int(val)

async def handle_memory(dut):
    while True:
        await RisingEdge(dut.clk)
        addr = int(dut.AB.value)
        we = bool(dut.WE.value)

        dut.DI.value = memory[addr]

        if we:
            memory[addr] = int(dut.DO.value)

async def check_instruction_sequence(dut, instruction_sequence):
    for expected_output in instruction_sequence:
        await RisingEdge(dut.clk)

        if expected_output:
            expected_addr, expected_we, expected_do = expected_output
            dut_addr = int(dut.AB.value)
            dut_we = bool(dut.WE.value)
            dut_do = int(dut.DO.value)

            assert dut_addr == expected_addr
            assert dut_we == expected_we

            if dut_we:
                assert dut_do == expected_do

@cocotb.test
async def test_reset(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff4, 0x12345678)

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        (0x00000100, True,  (int(dut.PC.value) >> 24) & 0xff),  # High addr
        (0x000001ff, True,  (int(dut.PC.value) >> 16) & 0xff),  # Mid high addr
        (0x000001fe, True,  (int(dut.PC.value) >>  8) & 0xff),  # Mid low addr
        (0x000001fd, True,  (int(dut.PC.value) >>  0) & 0xff),  # Low addr
        (0x000001fc, True,  int(dut.P.value)),                  # Status
        (0xfffffff4, False, int(dut.regfile.value)),            # read vector byte 0
        (0xfffffff5, False, int(dut.regfile.value)),            # read vector byte 1
        (0xfffffff6, False, int(dut.regfile.value)),            # read vector byte 2
        (0xfffffff7, False, int(dut.regfile.value)),            # read vector byte 3
        (0x12345678, False, int(dut.regfile.value)),            # Read first instruction
        (0x12345679, False, int(dut.regfile.value)),            # Read second byte
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)

@cocotb.test
async def test_absolute(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff4, 0x200)

    # lda $abcd1234
    # sta $50515253
    # wai
    write_bytes(0x200, [0xad, 0x34, 0x12, 0xcd, 0xab])
    write_bytes(0x205, [0x8d, 0x53, 0x52, 0x51, 0x50])
    write_byte(0x20a, 0xcb)
    write_byte(0xabcd1234, 0x55)

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00000200, False, None),            # Read first instruction
        (0x00000201, False, None),            # Read address byte 0
        (0x00000202, False, None),            # Read address byte 1
        (0x00000203, False, None),            # Read address byte 2
        (0x00000204, False, None),            # Read address byte 3
        (0xabcd1234, False, None),            # Read from absolute address
        (0x00000205, False, None),            # Read second instruction
        (0x00000206, False, None),            # Read address byte 0
        (0x00000207, False, None),            # Read address byte 1
        (0x00000208, False, None),            # Read address byte 2
        (0x00000209, False, None),            # Read address byte 3
        (0x50515253, True, 0x55),             # Write to absolute address
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)


@cocotb.test
async def test_absolute_x(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff4, 0x200)

    # ldx #1
    # lda $abcd1234,x
    # inx
    # sta $01020304,x
    # wai
    write_bytes(0x200, [0xa2, 0x01])
    write_bytes(0x202, [0xbd, 0x34, 0x12, 0xcd, 0xab])
    write_bytes(0x207, [0xe8])
    write_bytes(0x208, [0x9d, 0x04, 0x03, 0x02, 0x01])

    write_byte(0x20d, 0xcb)
    write_byte(0xabcd1235, 0xaa)

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00000200, False, None),            # ldx #1
        (0x00000201, False, None),            # Immediate
        (0x00000202, False, None),            # ldx $abcd1234,x
        (0x00000203, False, None),            # addr 0
        (0x00000204, False, None),            # addr 1
        (0x00000205, False, None),            # addr 2
        (0x00000206, False, None),            # addr 3
        (0xabcd1235, False, None),            # Read from address
        (0x00000207, False, None),            # inx
        (0x00000208, False, None),            # sta $01020304,x
        (0x00000208, False, None),            # store reg
        (0x00000209, False, None),            # addr 0
        (0x0000020a, False, None),            # addr 1
        (0x0000020b, False, None),            # addr 2
        (0x0000020c, False, None),            # addr 3
        (0x01020306, False, None),            # Write to address
        (0x01020306, True, 0xaa),             # Write to address
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)


@cocotb.test
async def test_absolute_y(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff4, 0x200)

    # ldy #1
    # lda $abcd1234,y
    # iny
    # sta $01020304,y
    # wai
    write_bytes(0x200, [0xa0, 0x01])
    write_bytes(0x202, [0xb9, 0x34, 0x12, 0xcd, 0xab])
    write_bytes(0x207, [0xc8])
    write_bytes(0x208, [0x99, 0x04, 0x03, 0x02, 0x01])

    write_byte(0x20d, 0xcb)
    write_byte(0xabcd1235, 0xaa)

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00000200, False, None),            # ldx #1
        (0x00000201, False, None),            # Immediate
        (0x00000202, False, None),            # ldx $abcd1234,x
        (0x00000203, False, None),            # addr 0
        (0x00000204, False, None),            # addr 1
        (0x00000205, False, None),            # addr 2
        (0x00000206, False, None),            # addr 3
        (0xabcd1235, False, None),            # Read from address
        (0x00000207, False, None),            # inx
        (0x00000208, False, None),            # sta $01020304,x
        (0x00000208, False, None),            # store reg
        (0x00000209, False, None),            # addr 0
        (0x0000020a, False, None),            # addr 1
        (0x0000020b, False, None),            # addr 2
        (0x0000020c, False, None),            # addr 3
        (0x01020306, False, None),            # Write to address
        (0x01020306, True, 0xaa),             # Write to address
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)


@cocotb.test
async def test_absolute_x_indirect(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff4, 0x200)

    # ldx #1
    # jmp ($deadbeef,x)
    write_bytes(0x200, [0xa2, 0x01])
    write_bytes(0x202, [0x7c, 0xef, 0xbe, 0xad, 0xde])
    write_byte(0xbeefb055, 0xcb)

    write_dword(0xdeadbeef + 1, 0xbeefb055)

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00000200, False, None),            # ldx #1
        (0x00000201, False, None),            # Immediate
        (0x00000202, False, None),            # jmp ($deadbeef,x)
        (0x00000203, False, None),            # addr 0
        (0x00000204, False, None),            # addr 1
        (0x00000205, False, None),            # addr 2
        (0x00000206, False, None),            # addr 3
        (0xdeadbef0, False, None),            # addr 0
        (0xdeadbef1, False, None),            # addr 1
        (0xdeadbef2, False, None),            # addr 2
        (0xdeadbef3, False, None),            # addr 3
        (0xbeefb055, False, None),            # target
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)


@cocotb.test
async def test_indirect(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff4, 0x200)

    # jmp ($deadbeef)
    write_bytes(0x200, [0x6c, 0xef, 0xbe, 0xad, 0xde])
    write_byte(0xbeefb055, 0xcb)

    write_dword(0xdeadbeef, 0xbeefb055)

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00000200, False, None),            # jmp ($deadbeef)
        (0x00000201, False, None),            # addr 0
        (0x00000202, False, None),            # addr 1
        (0x00000203, False, None),            # addr 2
        (0x00000204, False, None),            # addr 3
        (0xdeadbeef, False, None),            # addr 0
        (0xdeadbef0, False, None),            # addr 1
        (0xdeadbef1, False, None),            # addr 2
        (0xdeadbef2, False, None),            # addr 3
        (0xbeefb055, False, None),            # target
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)

@cocotb.test
async def test_indexed_indirect(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff4, 0x200)

    # ldy #1
    # lda ($04),y
    # iny
    # inc
    # sta ($04),y
    write_bytes(0x200, [0xa0, 0x02])
    write_bytes(0x202, [0xb1, 0x04])
    write_bytes(0x204, [0xc8])
    write_bytes(0x205, [0x1a])
    write_bytes(0x206, [0x91, 0x04])
    write_byte(0x208, 0xcb)

    write_dword(0x04, 0xabcdbeef)
    write_dword(0xabcdbeef+2, 0xaa)

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00000200, False, None),            # ldy #0
        (0x00000201, False, None),            # Immediate
        (0x00000202, False, None),            # lda ($04),y
        (0x00000203, False, None),            # ZP index
        (0x00000004, False, None),            # zp addr 0
        (0x00000005, False, None),            # zp addr 1
        (0x00000006, False, None),            # zp addr 2
        (0x00000007, False, None),            # zp addr 3
        (0xabcdbef1, False, None),            # fetch data
        (0x00000204, False, None),            # iny
        (0x00000205, False, None),            # iny
        (0x00000205, False, None),            # inc
        (0x00000206, False, None),            # inc
        (0x00000206, False, None),            # sta ($04),y
        (0x00000207, False, None),            # ZP index
        (0x00000004, False, None),            # zp addr 0
        (0x00000005, False, None),            # zp addr 1
        (0x00000006, False, None),            # zp addr 2
        (0x00000007, False, None),            # zp addr 3
        (0xabcdbef2, False, None),            # store takes extra cycle
        (0xabcdbef2, True, 0xAB),             # write data
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)

@cocotb.test
async def test_indirect_indexed(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff4, 0x200)


    # ldx #4
    # lda ($04,x)
    # inc
    # ldx #8
    # sta ($04,x)
    write_bytes(0x200, [0xa2, 0x04])
    write_bytes(0x202, [0xa1, 0x04])
    write_bytes(0x204, [0x1a])
    write_bytes(0x205, [0xa2, 0x08])
    write_bytes(0x207, [0x81, 0x04])
    write_byte(0x209, 0xcb)

    write_dword(0x08, 0xfeedf00d)
    write_dword(0x0c, 0xf00d600d)

    write_byte(0xfeedf00d, 0x69)

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00000200, False, None),            # ldx #4
        (0x00000201, False, None),            # Immediate
        (0x00000202, False, None),            # lda ($04,x)
        (0x00000203, False, None),            # ZP index
        (0x00000004, False, None),            # Compute ZP index
        (0x00000008, False, None),            # zp addr 0
        (0x00000009, False, None),            # zp addr 1
        (0x0000000a, False, None),            # zp addr 2
        (0x0000000b, False, None),            # zp addr 3
        (0xfeedf00d, False, None),            # fetch data
        (0x00000204, False, None),            # iny
        (0x00000205, False, None),            # iny
        (0x00000205, False, None),            # inc
        (0x00000206, False, None),            # inc
        (0x00000207, False, None),            # sta ($04),y
        (0x00000208, False, None),            # ZP index
        (0x00000004, False, None),            # Compute ZP index
        (0x0000000c, False, None),            # zp addr 0
        (0x0000000d, False, None),            # zp addr 1
        (0x0000000e, False, None),            # zp addr 2
        (0x0000000f, False, None),            # zp addr 3
        (0xf00d600d, True, 0x6a),             # write data
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)

@cocotb.test
async def test_jsr(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff4, 0x200)

    # @0x200
    # ldx #$0
    # txs
    # jsr $12345678
    # wai
    #
    # @0x1234
    # rts
    write_bytes(0x200, [0xa2, 0xff, 0x9a, 0x20, 0x78, 0x56, 0x34, 0x12, 0xcb])
    write_bytes(0x12345678, [0x60])

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00000200, False, None),  # ldx #$00
        (0x00000201, False, None),  # Immediate
        (0x00000202, False, None),  # txs
        (0x00000203, False, None),  # second cycle of txs
        (0x00000203, False, None),  # jsr $12345678
        (0x00000204, False, None),  # first byte of address
        (0x000001ff, True,  0x00),  # 24-31
        (0x000001fe, True,  0x00),  # 16-23
        (0x000001fd, True,  0x02),  # 8-15
        (0x000001fc, True,  0x05),  # 7-0
        (0x00000205, False, None),  # second byte of address
        (0x00000206, False, None),  # third byte of address
        (0x00000207, False, None),  # fourth byte of address
        (0x00000208, False, None),  # receive last byte of address
        (0x12345678, False, None),  # rts
        (0x12345679, False, None),  # rts
        (0x000001fb, False, None),  # current stack while we add 1 to it
        (0x000001fc, False, None),  # 7-0
        (0x000001fd, False, None),  # 15-8
        (0x000001fe, False, None),  # 23-16
        (0x000001ff, False, None),  # 31-24
        (0x1234567c, False, None),  # Updating PC before jump
        (0x00000208, False, None),  # WAI
        (0x00000209, False, None),  # second wai
        (0x0000020a, False, None),  # third wai
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)

# @cocotb.test
async def test_adc(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff4, 0x200)


    def zp_indirect():
        value = random.randint(0,255)
        address = random.randint(0x300, 0xffffffff)

        address_le = struct.pack("<I", address)

        # lda #$0
        # adc ($00)
        # wai
        write_bytes(0x200, [0xa9, 0x00, 0x72, 0x00, 0xcb])

        write_bytes(0x00, address_le)
        write_byte(address, value)

        return value
    
    def absolute():
        value = random.randint(0,255)
        address = random.randint(0x300, 0xffffffff)

        address_le = struct.pack("<I", address)

        # lda $#0
        # adc $address
        # wai

        write_bytes(0x200, [0xa9, 0x00, 0x6d])
        write_bytes(0x203, address_le)
        write_bytes(0x207, [0xcb])

        write_byte(address, value)

        return value

    def indirect_x():
        value = random.randint(0,255)
        address = random.randint(0x300, 0xffffff00)
        x_offset = random.randint(0,250)

        address_le = struct.pack("<I", address)

        # lda #$0
        # ldx #x_offset
        # adc ($00,x)
        # wai

        write_bytes(0x200, [0xa9, 0x00, 0xa2, x_offset, 0x61, 0x00, 0xcb])
        write_bytes(x_offset, address_le)
        write_byte(address, value)

        return value

    tests = [zp_indirect, absolute, indirect_x]

    for test in tests:
        value = test()

        dut.RDY.value = Immediate(1)

        dut.reset.value = Immediate(1)
        for _ in range(10):
            await RisingEdge(dut.clk)
        dut.reset.value = 0

        await FallingEdge(dut.RDY_O)

        assert value == int(dut.A.value)