import time
import board
import digitalio
import busio
#import neopixel
import supervisor
import sys

def modbusCrc(msg:str) -> int:
    crc = 0xFFFF
    for n in range(len(msg)):
        crc ^= msg[n]
        for i in range(8):
            if crc & 1:
                crc >>= 1
                crc ^= 0xA001
            else:
                crc >>= 1
    return crc


cs = digitalio.DigitalInOut(board.A0)
#cs.direction = digitalio.Direction.OUTPUT
cs.switch_to_output()
cs.value = True     # this should disable the flash chip

STATE_WAIT, STATE_ATTENTION,        \
        STATE_WRITE, STATE_WRITE_ADDRESS, STATE_WRITE_DATA, STATE_CHECK_CRC,        \
        STATE_WRITE_FLASH, STATE_READ, STATE_READ_ADDRESS, STATE_READ_PAGE,         \
        STATE_ERASE = range(11)

# current state
state = STATE_WAIT
address = 0
page = []
block = 0

def read_line () :
    value = ""
    while True :
        n = supervisor.runtime.serial_bytes_available
        if n == 0 : continue
        ch = sys.stdin.read(1)
        if ch == "\r" or ch == "\n" :
            break
        value += ch.upper()
    return value

def state_wait () :
    global state
    line = read_line()
    if line == "AT" :
        state = STATE_ATTENTION
    elif line == "WRITE" :
        state = STATE_WRITE
    elif line == ("READ") :
        state = STATE_READ
    elif line == ("ERASE") :
        state = STATE_ERASE

def state_attention () :
    global state
    print ("OK" + " ATTENTION")            # this sends both a CR and LF
    state = STATE_WAIT

def state_write () :
    global state
    print ("OK")
    state = STATE_WRITE_ADDRESS

def state_write_address () :
    global state,address,page,block
    line = read_line()
    if line.startswith("ADDRESS") :
        address = int(line.split(":")[1])
        print ("OK")
        page = bytearray(256)
        block = 0
        state = STATE_WRITE_DATA
    else :
        print ("NO WRITE ADDRESS")
        state = STATE_WAIT
    
def state_write_data () :
    global state,page,block
    data = read_line()
    page[block:block+16] = bytes([int(h.strip(),16) for h in data.split(",")])
    print ("OK", block)
    if block == 240 :
        state = STATE_CHECK_CRC
    block += 16

def state_check_crc () :
    global state
    text = read_line()
    if text.startswith("CRC") :
        crc = int(text.split(":")[1],16)
        if crc == modbusCrc(page) :
            state = STATE_WRITE_FLASH
        else :
            print ("CRC CHECK FAILED")
            state = STATE_WAIT
    else :
        print ("MALFORMED CRC")
        state = STATE_WAIT
        
def state_write_flash () :
    global state
    # send a write enable
    cs.value = False
    spi.write (bytes([0x06]))
    cs.value = True
    # send the page
    #cmd = bytearray([0x02, (address >> 16) & 0xFF, (address >> 8) & 0xFF, address & 0xFF])
    cs.value = False
    #spi.write (cmd)
    spi.write (bytes([0x02, (address >> 16) & 0xFF, (address >> 8) & 0xFF, address & 0xFF]))
    spi.write (bytes(page))
    cs.value = True
    # wait for the write to finish
    print ("OK")
    """
    for n in range(10):
        cs.value = False
        spi.write (bytes([0x05]))
        status = bytearray(2)
        spi.readinto (buffer=status, write_value=0xFF)
        cs.value = True
        print ("status =",hex(status[0]),":",hex(status[0]))
        time.sleep(0.1)
    """
    state = STATE_WAIT

def state_read () :
    global state
    print ("OK")
    state = STATE_READ_ADDRESS

def state_read_address () :
    global state,address
    line = read_line()
    if line.startswith("ADDRESS") :
        address = int(line.split(":")[1])
        state = STATE_READ_PAGE
    else :
        print ("NO WRITE ADDRESS")
        state = STATE_WAIT

def state_read_page () :
    global state,page
    page = bytearray(256)
    cs.value = False        # the Chip Select
    cmd = bytearray([0x03, (address >> 16) & 0xFF, (address >> 8) & 0xFF, address & 0xFF])
    spi.write (cmd)
    spi.readinto (buffer=page, write_value = 0xFF)
    cs.value = True         # disable the chip
    # send the data 16 bytes at a time (to save memory)
    for n in range(16) :
        def my_hex(num) :
            return hex(num)[2:]
        text = ",".join(map(my_hex, page[(n * 16):(n * 16) + 16]))
        print (text)
    state = STATE_WAIT

def state_erase () :
    global state
    # send a write enable
    cs.value = False
    spi.write (bytes([0x06]))
    cs.value = True
    # send the chip erase 
    cs.value = False
    spi.write (bytes([0xC7]))
    cs.value = True
    print ("OK")
    state = STATE_WAIT
    
spi = busio.SPI(clock=board.SCK ,MISO=board.MISO, MOSI=board.MOSI)
while not spi.try_lock() :      # lock out anything else using the SPI
    pass
spi.configure(baudrate=100000, phase=0, polarity=0)


functions = [state_wait, state_attention,
             state_write, state_write_address, state_write_data, state_check_crc,
             state_write_flash, state_read, state_read_address, state_read_page,
             state_erase]

while True :
    try :
        #print ("state = ",state)
        functions[state]()
    except Exception as err :
        print ("State:",state,"    Exception:",err)
        state = STATE_WAIT
        time.sleep(10)
    finally:
        pass