from machine import Pin, SPI, I2C
import ssd1306
import time

# OLED Setup
print("Setting up OLED...")
i2c = I2C(1, sda=Pin(18), scl=Pin(19), freq=400000)
oled_width, oled_height = 128, 64
oled = ssd1306.SSD1306_I2C(oled_width, oled_height, i2c)

# Button Setup
print("Setting up buttons...")
button_nav = Pin(7, Pin.IN, Pin.PULL_UP)  # GP7 for navigation
button_select = Pin(3, Pin.IN, Pin.PULL_UP)   # GP3 for selecting

# SD Card Setup
print("Setting up SD card...")
SCK_PIN, MOSI_PIN, MISO_PIN, CS_PIN = 10, 11, 12, 13
spi = SPI(1, baudrate=400000, polarity=0, phase=0, sck=Pin(SCK_PIN), mosi=Pin(MOSI_PIN), miso=Pin(MISO_PIN))
cs = Pin(CS_PIN, Pin.OUT)
cs.high()

# Menu Setup
menu_options = ["TEST LOCK", "TEMP LOCK", "PERM LOCK", "TEST RW"]
selected_option = 0
menu_start_index = 0
header_height, menu_start_y, menu_offset_x, menu_item_spacing = 10, 24, 68, 12
menu_option_length = max(len(option) for option in menu_options) + 2
displayable_items = (oled_height - menu_start_y) // menu_item_spacing




# Button debounce
last_button_time = 0
debounce_delay = 200  # milliseconds

# Button state tracking
last_nav_state = button_nav.value()
last_select_state = button_select.value()

# Helper functions for SD card operations
def spi_read_write_byte(data):
    resp = bytearray(1)
    spi.write_readinto(bytearray([data]), resp)
    return resp[0]

def send_cmd(cmd, arg, crc):
    cs.high()
    spi.write(b'\xFF' * 8)
    cs.low()
    spi.write(b'\xFF' * 8)
    spi.write(bytearray([cmd | 0x40, (arg >> 24) & 0xFF, (arg >> 16) & 0xFF, (arg >> 8) & 0xFF, arg & 0xFF, crc]))
    for _ in range(8):
        response = spi.read(1)[0]
        if response != 0xFF:
            return response
    spi.write(b'\xFF' * 8)
    cs.high()
    spi.write(b'\xFF' * 8)
    return 0xFF

def wait_ready():
    for _ in range(0xFFFFFF):
        if spi_read_write_byte(0xFF) == 0xFF:
            return True
    return False

def crc7(data):
    crc = 0
    for byte in data:
        crc ^= byte
        for _ in range(8):
            crc = (crc << 1) ^ 0x12 if crc & 0x80 else crc << 1
    return (crc & 0x7F) | 0x01

def crc16(data):
    crc = 0
    for byte in data:
        crc ^= byte << 8
        for _ in range(8):
            crc = (crc << 1) ^ 0x1021 if crc & 0x8000 else crc << 1
    return crc & 0xFFFF

def send_cmd27(csd):
    cmd_data = bytearray([27 | 0x40, 0, 0, 0, 0])
    cmd_crc = crc7(cmd_data)
    response = send_cmd(27, 0, cmd_crc)
    if response == 0x00:
        while spi_read_write_byte(0xFF) != 0xFE:
            pass
        spi.write(csd)
        csd_crc = crc16(csd)
        spi.write(bytearray([(csd_crc >> 8) & 0xFF, csd_crc & 0xFF]))
        response = spi_read_write_byte(0xFF)
        if (response & 0x1F) == 0x05:
            return True
        while spi_read_write_byte(0xFF) != 0xFF:
            pass
    return False

def init_sd_card():
    print("Initializing SD card...")
    for _ in range(10):
        if send_cmd(0, 0, 0x95) == 0x01:
            break
        time.sleep(0.1)
    else:
        raise RuntimeError("Failed to reset the SD card (CMD0)")

    if send_cmd(8, 0x1AA, 0x87) != 0x01:
        raise RuntimeError("CMD8 failed")

    for _ in range(1000):
        if send_cmd(55, 0, 0x65) > 0x01:
            raise RuntimeError("CMD55 failed")
        if send_cmd(41, 0x40000000, 0x77) == 0x00:
            break
    else:
        for _ in range(1000):
            if send_cmd(1, 0, 0xF9) == 0x00:
                break
        else:
            raise RuntimeError("Failed to initialize the SD card (ACMD41/CMD1)")

    if send_cmd(58, 0, 0x01) == 0x00:
        ocr = spi.read(4)
        if not (ocr[0] & 0x40):
            if send_cmd(16, 512, 0x01) != 0x00:
                raise RuntimeError("Failed to set block length (CMD16)")

def read_csd():
    print("Reading CSD...")
    if send_cmd(9, 0, 0x01) == 0x00:
        while spi.read(1)[0] != 0xFE:
            pass
        csd = spi.read(16)
        spi.read(2)  # Dummy CRC
        return csd
    raise RuntimeError("Failed to read CSD (CMD9)")

def modify_csd(csd, lock_type):
    print(f"Modifying CSD for {lock_type}...")
    csd_reg = CSDRegister(int.from_bytes(csd, 'big'))
    if lock_type == "TEMP":
        csd_reg.write_field('TMP_WRITE_PROTECT', 1)
    elif lock_type == "PERM":
        csd_reg.write_field('PERM_WRITE_PROTECT', 1)
    return csd_reg.csd.to_bytes(16, 'big')

# OLED display functions
def display_header():
    oled.fill(0)
    oled.text("SD LOCKER", 25, 0, 1)
    oled.hline(0, 13, oled_width, 1)

def display_menu():
    global menu_start_index

    if selected_option >= menu_start_index + displayable_items:
        menu_start_index = selected_option - displayable_items + 1
    elif selected_option < menu_start_index:
        menu_start_index = selected_option

    # Clear the menu area only
    oled.fill_rect(0, menu_start_y, oled_width, oled_height - menu_start_y, 0)
    for i in range(displayable_items):
        menu_index = menu_start_index + i
        if menu_index >= len(menu_options):
            break
        option = menu_options[menu_index]
        y_pos = menu_start_y + i * menu_item_spacing
        option_x = menu_offset_x - menu_option_length * 4
        if menu_index == selected_option:
            oled.text(">", option_x - 11, y_pos, 1)
            oled.text(option, option_x, y_pos, 1)
            oled.text("<", option_x + menu_option_length * 8 - 10, y_pos, 1)
        else:
            oled.text(option, option_x, y_pos, 1)
    oled.show()





def display_message(message):
    print(f"Displaying message: {message}")
    oled.fill(0)
    oled.text(message, 0, 32, 1)
    oled.show()
    time.sleep(2)

# Add CSDRegister classes and related functions
def bitslice(val, start, end):
    if start > end:
        (start, end) = (end, start)
        reverse = True
    else:
        reverse = False
    r = 0
    shift = 0
    for i in range(start, end+1):
        b = (val >> i) & 1
        if reverse:
            r |= b << shift
            shift += 1
        else:
            r <<= 1
            r |= b
    return r

class CSDRegisterBase:
    FIELD_MAP = {}
    def __init__(self, csd):
        self.csd = csd

    def read_field(self, field):
        slicer = self.FIELD_MAP[field]
        width = max(slicer) - min(slicer) + 1
        binary = bitslice(self.csd, slicer[0], slicer[1])
        return (binary, width)

    def write_field(self, field, value):
        slicer = self.FIELD_MAP[field]
        mask = ((1 << (slicer[0] - slicer[1] + 1)) - 1) << slicer[1]
        self.csd = (self.csd & ~mask) | (value << slicer[1])

class CSDRegister_V10(CSDRegisterBase):
    FIELD_MAP = {
        'CSD_STRUCTURE': (127, 126),
        'TAAC': (119, 112),
        'NSAC': (111, 104),
        'TRAN_SPEED': (103, 96),
        'CCC': (95, 84),
        'READ_BL_LEN': (83, 80),
        'READ_BL_PARTIAL': (79, 79),
        'WRITE_BLK_MISALIGN': (78, 78),
        'READ_BLK_MISALIGN': (77, 77),
        'DSR_IMP': (76, 76),
        'C_SIZE': (73, 62),
        'VDD_R_CURR_MIN': (61, 59),
        'VDD_R_CURR_MAX': (58, 56),
        'VDD_W_CURR_MIN': (55, 53),
        'VDD_W_CURR_MAX': (52, 50),
        'C_SIZE_MULT': (49, 47),
        'ERASE_BLK_EN': (46, 46),
        'SECTOR_SIZE': (45, 39),
        'WP_GRP_SIZE': (38, 32),
        'WP_GRP_ENABLE': (31, 31),
        'R2W_FACTOR': (28, 26),
        'WRITE_BL_LEN': (25, 22),
        'WRITE_BL_PARTIAL': (21, 21),
        'FILE_FORMAT_GRP': (15, 15),
        'COPY': (14, 14),
        'PERM_WRITE_PROTECT': (13, 13),
        'TMP_WRITE_PROTECT': (12, 12),
        'FILE_FORMAT': (11, 9),
        'CRC': (7, 1)
    }

    def calculate_capacity(self):
        c_size, _ = self.read_field('C_SIZE')
        c_size_mult, _ = self.read_field('C_SIZE_MULT')
        read_bl_len, _ = self.read_field('READ_BL_LEN')
        capacity = (c_size + 1) * (2 ** (c_size_mult + 2)) * (2 ** read_bl_len)
        return capacity

class CSDRegister_V20(CSDRegister_V10):
    FIELD_MAP = {
        'CSD_STRUCTURE': (127, 126),
        'TAAC': (119, 112),
        'NSAC': (111, 104),
        'TRAN_SPEED': (103, 96),
        'CCC': (95, 84),
        'READ_BL_LEN': (83, 80),
        'READ_BL_PARTIAL': (79, 79),
        'WRITE_BLK_MISALIGN': (78, 78),
        'READ_BLK_MISALIGN': (77, 77),
        'DSR_IMP': (76, 76),
        'C_SIZE': (69, 48),
        'ERASE_BLK_EN': (46, 46),
        'SECTOR_SIZE': (45, 39),
        'WP_GRP_SIZE': (38, 32),
        'WP_GRP_ENABLE': (31, 31),
        'R2W_FACTOR': (28, 26),
        'WRITE_BL_LEN': (25, 22),
        'WRITE_BL_PARTIAL': (21, 21),
        'FILE_FORMAT_GRP': (15, 15),
        'COPY': (14, 14),
        'PERM_WRITE_PROTECT': (13, 13),
        'TMP_WRITE_PROTECT': (12, 12),
        'FILE_FORMAT': (11, 9),
        'CRC': (7, 1)
    }

    def calculate_capacity(self):
        c_size, _ = self.read_field('C_SIZE')
        capacity = (c_size + 1) * 524288  # 512KB blocks
        return capacity

def CSDRegister(csd):
    if isinstance(csd, str):
        csd = int(csd, 16)
    assert isinstance(csd, int)
    val = bitslice(csd, 127, 126)
    if val == 0:
        return CSDRegister_V10(csd)
    elif val == 1:
        return CSDRegister_V20(csd)
    else:
        raise CSDRegisterError()

def display_csd_info(csd):
    print("Displaying CSD information...")
    try:
        csd_reg = CSDRegister(int.from_bytes(csd, 'big'))
        capacity = csd_reg.calculate_capacity() / (1024*1024)
        perm_wp = csd_reg.read_field("PERM_WRITE_PROTECT")[0]
        temp_wp = csd_reg.read_field("TMP_WRITE_PROTECT")[0]
        
        oled.fill(0)
        oled.text("CSD Info:", 0, 0, 1)
        oled.text(f"Capacity: {capacity:.2f}MB", 0, 12, 1)
        oled.text(f"PERM WP: {perm_wp}", 0, 28, 1)
        oled.text(f"TEMP WP: {temp_wp}", 0, 44, 1)
        oled.text("Press to exit", 0, 56, 1)
        oled.show()
        
        # Wait for button press or timeout
        start_time = time.ticks_ms()
        while time.ticks_diff(time.ticks_ms(), start_time) < 10000:  # 10 seconds timeout
            if button_select.value() == 0 or button_nav.value() == 0:  # If any button is pressed
                break
            time.sleep_ms(100)
        
    except Exception as e:
        print(f"Error in display_csd_info: {str(e)}")
        display_message(f"CSD Info Error")
        time.sleep(2)

def sd_write_block(block_addr, data):
    if send_cmd(24, block_addr, 0xFF) == 0x00:  # CMD24 to write block
        while spi_read_write_byte(0xFF) != 0xFF:
            pass
        spi.write(b'\xFE')  # Data token
        spi.write(data)
        crc = crc16(data)
        spi.write(bytearray([(crc >> 8) & 0xFF, crc & 0xFF]))
        response = spi.read(1)[0]
        if (response & 0x1F) == 0x05:
            while spi_read_write_byte(0xFF) != 0xFF:
                pass
            return True
    return False

def sd_read_block(block_addr):
    if send_cmd(17, block_addr, 0xFF) == 0x00:  # CMD17 to read block
        while spi_read_write_byte(0xFF) != 0xFE:
            pass
        data = spi.read(512)
        spi.read(2)  # Dummy CRC
        return data
    return None

def test_sd_card_rw():
    block_addr = 0x1000  # Example block address
    test_string = "Hello, SD card!"
    test_data = bytearray(test_string, 'utf-8')
    if len(test_data) < 512:
        test_data.extend(b'\x00' * (512 - len(test_data)))  # Pad to 512 bytes

    # Write to the SD card
    if sd_write_block(block_addr, test_data):
        display_message("Write Success")
    else:
        display_message("Write Failed")

    # Read from the SD card
    read_data = sd_read_block(block_addr)
    if read_data:
        read_string = read_data.decode('utf-8').rstrip('\x00')
        print(f"Read data: {read_string}")
        if read_string == test_string:
            display_message("Read/Write test passed!")
        else:
            display_message("Read/Write test failed!")
    else:
        display_message("Read Failed")


# Main execution
try:
    init_sd_card()
    while True:
        display_header()  # Display the header
        display_menu()  # Display the menu

        current_time = time.ticks_ms()
        if current_time - last_button_time > debounce_delay:
            nav_state = button_nav.value()
            select_state = button_select.value()

            if nav_state != last_nav_state:
                if nav_state == 0:  # Button pressed (transition from 1 to 0)
                    selected_option = (selected_option + 1) % len(menu_options)
                    last_button_time = current_time
                last_nav_state = nav_state

            if select_state != last_select_state:
                if select_state == 0:  # Button pressed (transition from 1 to 0)
                    option = menu_options[selected_option]
                    if option == "TEST LOCK":
                        csd = read_csd()
                        display_message("  DECODING CSD")
                        print("CSD:", ' '.join([f'{b:02X}' for b in csd]))
                        display_csd_info(csd)
                    elif option == "TEST RW":
                        test_sd_card_rw()
                    else:
                        csd = read_csd()
                        modified_csd = modify_csd(csd, option.split()[0])
                        if send_cmd27(modified_csd):
                            display_message(f"{option} Success")
                            display_csd_info(modified_csd)
                        else:
                            display_message(f"{option} Failed")
                    last_button_time = current_time
                last_select_state = select_state

        time.sleep_ms(50)




except Exception as e:
    display_message(f"Error: {str(e)}")

finally:
    oled.fill(0)
    oled.show()
    cs.high()
    spi.write(b'\xFF')