import os
import time
import board
import digitalio
import struct

# 4k ram tester
def make_dios(*args):
    return [digitalio.DigitalInOut(a) for a in args]


address_pins = make_dios(*(getattr(board, f"GP{i}") for i in range(11, 11+12)))
for a in address_pins:
    a.switch_to_output()

data_pins = make_dios(*(getattr(board, f"GP{i}") for i in range(8)))

not_write_enable = digitalio.DigitalInOut(board.GP10)
not_write_enable.switch_to_output(True)

not_output_enable = digitalio.DigitalInOut(board.GP8)
not_output_enable.switch_to_output(True)

not_chip_enable = digitalio.DigitalInOut(board.GP9)
not_chip_enable.switch_to_output(True)

random_data = os.urandom(16)

def set_pins(value, pins):
    for p in pins:
        p.switch_to_output(value & 1)
        value >>= 1

def get_pins(pins):
    value = 0
    for i, p in enumerate(pins):
        value |= p.value << i
    return value

def read_out(*args):
    for p in data_pins:
        p.switch_to_input(digitalio.Pull.UP)

    not_chip_enable.value = False

    if not args: args = (0x800,)
    for addr in range(*args):
        set_pins(addr, address_pins) # There's no way the 150ns setup is violated at CP speeds!
        not_output_enable.value = False
        yield get_pins(data_pins)
        not_output_enable.value = True

    not_chip_enable.value = True

def _write(data, start_addr=0):
    for p in data_pins:
        p.switch_to_output(False)

    not_output_enable.value = True
    
    print(end=f"{start_addr:04x}")
    for offset, byte in enumerate(data):
        addr = start_addr + offset
        set_pins(addr, address_pins)
        set_pins(byte, data_pins)
        not_chip_enable.value = False
        assert get_pins(data_pins) == byte
        not_write_enable.value = False
        not_write_enable.value = True
        not_chip_enable.value = True
        print(end='.')

    not_chip_enable.value = True

def write(data, start_addr=0):
    _write(data, start_addr)
    print()

def verify(data, start_addr = 0):
    for p in data_pins:
        p.switch_to_input(digitalio.Pull.UP)

    not_output_enable.value = False
    not_chip_enable.value = False

    result = True

    for offset, byte in enumerate(data):
        addr = start_addr + offset
        set_pins(addr, address_pins)

        readback = get_pins(data_pins)
        if readback != byte:
            print(f"0x{addr:04x} {byte:02x} {readback:02x}")
            result = False

    if not result:
        raise ValueError("Verify failed")

    print("Verified")


def write_verify(data, start_addr = 0):
    _write(data, start_addr)
    verify(data, start_addr)

## Uncomment if you want to pre-erase the eeprom (not needed)
#write_verify(b"\xff" * 2048)

def readn(start, count=1):
    return bytes(read_out(start, start+count))

def print_ihex(address, rectype, data):
    data = list(data)
    data_str = "".join(f"{b:02x}" for b in data)
    data_len = len(data)
    assert data_len < 256
    assert 0 <= rectype < 256
    assert 0 <= address < 65536
    
    checksum = (-(data_len + sum(data) + rectype + address + address // 256)) & 0xff

    print(f":{data_len:02x}{address:04x}{rectype:02x}{data_str}{checksum:02x}")

if __name__ == '__main__':
    command = input("[P]rogram, [V]erify or [R]eadout? ").lower() 
    if command == "p":
        write(random_data, 0)
    if command in ("p", "v"):
        verify(random_data, 0)
    else:
        stride = 32
        end = 4096
        if end % stride != 0:
            raise ValueError("{stride=} incompatible with {end=}")

        for start_address in range(0, 2048, stride):
            print_ihex(start_address, 0, read_out(start_address, start_address + stride))
        print_ihex(0, 1, [])