"""
Test script for a Winbond W25X40CL (4 Mbit / 512 KB) SPI flash on a Pico 2.

Wiring: CS=GPIO13, CLK=GPIO10, MOSI(DI)=GPIO11, MISO(DO)=GPIO12 (SPI1).
These SPI pins are shared with the LCD, so the LCD's chip select (GPIO15)
is held high here to keep the display off the bus.

WARNING: this erases and rewrites ONE 4 KB sector (the last one by default).
Change TEST_ADDR if you keep data there.
"""
from machine import Pin, SPI
from time import sleep_ms, ticks_ms, ticks_diff

# ---------------- Settings ----------------
BAUD = 10_000_000
FLASH_SIZE = 512 * 1024
SECTOR_SIZE = 4096
PAGE_SIZE = 256
TEST_ADDR = FLASH_SIZE - SECTOR_SIZE      # last sector, sector-aligned

# ---------------- Commands ----------------
CMD_WRITE_ENABLE = 0x06
CMD_READ_STATUS = 0x05
CMD_WRITE_STATUS = 0x01
CMD_READ = 0x03
CMD_PAGE_PROGRAM = 0x02
CMD_SECTOR_ERASE = 0x20
CMD_MANUF_DEV_ID = 0x90
CMD_JEDEC_ID = 0x9F
CMD_UNIQUE_ID = 0x4B

STATUS_BUSY = 0x01
STATUS_WEL = 0x02
STATUS_PROTECT_MASK = 0x7C   # BP0-2, TB, SEC

# ---------------- Hardware ----------------
Pin(15, Pin.OUT, value=1)                 # LCD chip select: keep deselected
cs = Pin(13, Pin.OUT, value=1)
spi = SPI(1, baudrate=BAUD, polarity=0, phase=0,
          sck=Pin(10), mosi=Pin(11), miso=Pin(12))


# ---------------- Low-level helpers ----------------
def addr_bytes(addr):
    return bytes(((addr >> 16) & 0xFF, (addr >> 8) & 0xFF, addr & 0xFF))


def command(cmd, rx_len=0, prefix=b""):
    """Send cmd (+ optional prefix bytes), then read rx_len bytes."""
    cs(0)
    spi.write(bytes((cmd,)) + prefix)
    data = spi.read(rx_len) if rx_len else b""
    cs(1)
    return data


def read_status():
    return command(CMD_READ_STATUS, 1)[0]


def wait_ready(timeout_ms=1000):
    start = ticks_ms()
    while read_status() & STATUS_BUSY:
        if ticks_diff(ticks_ms(), start) > timeout_ms:
            raise OSError("flash busy timeout")
        sleep_ms(1)


def write_enable():
    command(CMD_WRITE_ENABLE)
    if not (read_status() & STATUS_WEL):
        raise OSError("write enable latch did not set")


def flash_read(addr, n):
    cs(0)
    spi.write(bytes((CMD_READ,)) + addr_bytes(addr))
    data = spi.read(n)
    cs(1)
    return data


def page_program(addr, data):
    if len(data) > PAGE_SIZE or (addr % PAGE_SIZE) + len(data) > PAGE_SIZE:
        raise ValueError("page program must not cross a page boundary")
    write_enable()
    cs(0)
    spi.write(bytes((CMD_PAGE_PROGRAM,)) + addr_bytes(addr))
    spi.write(data)
    cs(1)
    wait_ready(50)


def flash_write(addr, data):
    """Write any length, splitting at page boundaries (area must be erased)."""
    pos = 0
    while pos < len(data):
        room = PAGE_SIZE - ((addr + pos) % PAGE_SIZE)
        chunk = data[pos:pos + room]
        page_program(addr + pos, chunk)
        pos += len(chunk)


def sector_erase(addr):
    write_enable()
    command(CMD_SECTOR_ERASE, prefix=addr_bytes(addr))
    wait_ready(1000)


def pattern(n, seed=0):
    return bytes(((i * 7 + (i >> 8) * 31 + seed) & 0xFF) for i in range(n))


# ---------------- Test framework ----------------
results = []


def check(name, ok, detail=""):
    results.append(ok)
    print("[%s] %s %s" % ("PASS" if ok else "FAIL", name, detail))
    return ok


def hexs(b):
    return " ".join("%02X" % x for x in b)


# ---------------- Tests ----------------
def test_ids():
    jedec = command(CMD_JEDEC_ID, 3)
    # W25X40CL: manufacturer 0xEF (Winbond), type 0x30, capacity 0x13 (4 Mbit)
    check("JEDEC ID", jedec == b"\xEF\x30\x13",
          "got %s, expected EF 30 13" % hexs(jedec))
    if jedec in (b"\x00\x00\x00", b"\xFF\xFF\xFF"):
        print("      -> all 0x00/0xFF: check wiring, power, CS and that the LCD CS is high")

    mdid = command(CMD_MANUF_DEV_ID, 2, prefix=b"\x00\x00\x00")
    check("Manufacturer/Device ID", mdid == b"\xEF\x12",
          "got %s, expected EF 12" % hexs(mdid))

    # Not every part supports this command, so report only
    uid = command(CMD_UNIQUE_ID, 8, prefix=b"\x00\x00\x00\x00")
    print("[INFO] Unique ID: %s" % hexs(uid))


def test_status():
    st = read_status()
    print("[INFO] Status register: 0x%02X" % st)
    if st & STATUS_PROTECT_MASK:
        print("[INFO] Block protection bits set, clearing them")
        write_enable()
        command(CMD_WRITE_STATUS, prefix=b"\x00")
        wait_ready(100)
        st = read_status()
    check("No write protection", (st & STATUS_PROTECT_MASK) == 0,
          "status=0x%02X" % st)


def test_erase_write_read():
    # Remember the first bytes of the chip to prove nothing else gets disturbed
    guard_addr = 0 if TEST_ADDR != 0 else FLASH_SIZE - SECTOR_SIZE
    guard_before = flash_read(guard_addr, 256)

    # 1. Erase and verify blank
    t0 = ticks_ms()
    sector_erase(TEST_ADDR)
    erase_ms = ticks_diff(ticks_ms(), t0)
    blank = flash_read(TEST_ADDR, SECTOR_SIZE)
    check("Sector erase -> all 0xFF", blank == b"\xFF" * SECTOR_SIZE,
          "(%d ms)" % erase_ms)

    # 2. Unaligned write across two page boundaries
    data = pattern(600, seed=1)
    off = TEST_ADDR + 200
    flash_write(off, data)
    check("Write 600 bytes across page boundaries",
          flash_read(off, 600) == data)

    # 3. Bytes around the written region must still be 0xFF
    before = flash_read(TEST_ADDR, 200)
    after = flash_read(off + 600, 64)
    check("Neighbouring bytes untouched",
          before == b"\xFF" * 200 and after == b"\xFF" * 64)

    # 4. Full-sector write and read-back with timing
    sector_erase(TEST_ADDR)
    full = pattern(SECTOR_SIZE, seed=5)
    t0 = ticks_ms()
    flash_write(TEST_ADDR, full)
    write_ms = ticks_diff(ticks_ms(), t0)
    t0 = ticks_ms()
    readback = flash_read(TEST_ADDR, SECTOR_SIZE)
    read_ms = ticks_diff(ticks_ms(), t0)
    check("Full 4 KB sector write/read", readback == full,
          "(write %d ms, read %d ms)" % (write_ms, read_ms))

    # 5. Programming can only clear bits (1 -> 0); writing over data ANDs it
    sector_erase(TEST_ADDR)
    flash_write(TEST_ADDR, b"\xF0")
    flash_write(TEST_ADDR, b"\x3C")
    check("Program ANDs with existing data (F0 & 3C = 30)",
          flash_read(TEST_ADDR, 1) == b"\x30")

    # 6. Leave the sector erased and check the rest of the chip
    sector_erase(TEST_ADDR)
    check("Final erase -> all 0xFF",
          flash_read(TEST_ADDR, SECTOR_SIZE) == b"\xFF" * SECTOR_SIZE)
    check("Other flash data undisturbed",
          flash_read(guard_addr, 256) == guard_before)


def test_address_range():
    # A 4 Mbit part has 19 address bits; reading just past the end wraps to 0
    # if the chip is the size we think it is.
    start = flash_read(0, 16)
    wrapped = flash_read(FLASH_SIZE, 16)
    ok = start == wrapped
    check("Address wraps at 512 KB (chip size)", ok,
          "" if ok else "(unexpected: chip may be larger than 512 KB, or this region is blank)")


# ---------------- Run ----------------
def main():
    print("W25X40CL flash test, SPI1 @ %d Hz" % BAUD)
    print("Test sector: 0x%06X" % TEST_ADDR)
    print("-" * 40)
    try:
        test_ids()
        if not results[0]:
            print("ID mismatch, skipping the erase/write tests.")
        else:
            test_status()
            test_erase_write_read()
            test_address_range()
    except Exception as e:
        check("Unexpected exception", False, repr(e))
    print("-" * 40)
    print("%d of %d checks passed" % (sum(results), len(results)))


main()