# 1. Import necessary libraries
from machine import Pin, PWM, I2C
from ssd1306 import SSD1306
import time

# 2. Hardware initialization
## 2.1 4x4 Membrane Keypad
row_pins = [
    Pin(26, Pin.IN, Pin.PULL_UP),
    Pin(22, Pin.IN, Pin.PULL_UP),
    Pin(21, Pin.IN, Pin.PULL_UP),
    Pin(20, Pin.IN, Pin.PULL_UP)
]
col_pins = [
    Pin(19, Pin.OUT),
    Pin(18, Pin.OUT),
    Pin(17, Pin.OUT),
    Pin(16, Pin.OUT)
]
keymap = [['1','2','3','A'], ['4','5','6','B'], ['7','8','9','C'], ['*','0','#','D']]

## 2.2 SSD1306 OLED (I2C Address: 0x3c)
WIDTH = 128
HEIGHT = 64
i2c = I2C(1, scl=Pin(3), sda=Pin(2), freq=400000)
oled = SSD1306(WIDTH, HEIGHT, i2c, addr=0x3C)
oled.fill(0)
oled.text("System Init...", 0, 0)
oled.show()
time.sleep(0.5)

## 2.3 Servo Motor (GP15)
servo = PWM(Pin(15), freq=50)
def set_servo(locked):
    # Locked = 90° (duty cycle: 7.5% → 4915), Unlocked = 0° (duty cycle: 2.5% → 1638)
    servo.duty_u16(4915 if locked else 1638)

## 2.4 LED Indicators
green_led = Pin(13, Pin.OUT, value=0)  # Green LED: Unlocked status
red_led = Pin(14, Pin.OUT, value=0)    # Red LED: Error status

# 3. Keypad Scanning Function
def scan_key():
    for col_idx in range(4):
        col_pins[col_idx].value(0)  # Pull down current column to detect keys
        time.sleep_ms(2)            # Debounce delay (2ms to avoid false triggers)
        for row_idx in range(4):
            if row_pins[row_idx].value() == 0:  # Key press detected (row pulled low)
                while row_pins[row_idx].value() == 0:  # Wait for key release
                    time.sleep_ms(1)
                col_pins[col_idx].value(1)  # Restore column to high level
                return keymap[row_idx][col_idx]  # Return pressed key character
        col_pins[col_idx].value(1)  # Restore column to high level if no key press
    return None  # Return None if no key is pressed

# 4. Main Logic
def main():
    CORRECT_PW = "21906229"  # University ID (8-digit password)
    current_input = ""        # Buffer to store user's password input
    is_locked = True          # Initial state: Locked

    # Display Update Function (refresh OLED based on current state)
    def update_display():
        oled.fill(0)  # Clear OLED screen (black background)
        if is_locked:
            oled.text("Status: LOCKED", 0, 0)
            # Display masked password (e.g., "****" for 4-digit input)
            oled.text(f"PW: {len(current_input)*'*'}", 0, 20)
            oled.text("Press # to Enter", 0, 40)
        else:
            oled.text("Status: UNLOCKED", 0, 0)
            oled.text("Press * to LOCK", 0, 25)  # Clear lock prompt
            oled.text("Keep Unlocked...", 0, 45)
        oled.show()  # Update OLED with new content

    # Initialize system to locked state
    set_servo(locked=True)
    update_display()

    while True:
        key = scan_key()
        
        if key:  # Execute only when a key is pressed
            
            # --- * Key Pressed ---
            if key == "*":
                if not is_locked:
                    # If in unlocked state: Lock immediately
                    print("Executing Manual Lock...")
                    is_locked = True
                    set_servo(locked=True)
                    green_led.value(0)  # Turn off green LED (unlocked indicator)
                    current_input = ""  # Reset password input buffer
                    update_display()    # Refresh OLED to locked state
                else:
                    # If already locked: * key acts as input clear (Backspace)
                    current_input = ""
                    update_display()    # Refresh OLED to show empty input

            # --- # Key Pressed (Only active in locked state) ---
            elif key == "#":
                if is_locked:
                    if current_input == CORRECT_PW:
                        # Correct password → Unlock the system
                        print("Unlock Success")
                        is_locked = False
                        set_servo(locked=False)  # Rotate servo to unlocked position
                        green_led.value(1)       # Turn on green LED (unlocked indicator)
                        current_input = ""        # Reset input buffer
                        update_display()          # Refresh OLED to unlocked state
                    else:
                        # Incorrect password → Show error message
                        oled.fill(0)
                        oled.text("WRONG PASSWORD!", 0, 20)
                        oled.show()
                        red_led.value(1)         # Turn on red LED (error indicator)
                        time.sleep(1)            # Error message stays for 1 second only
                        red_led.value(0)         # Turn off red LED
                        current_input = ""        # Reset input buffer
                        update_display()          # Restore locked state display

            # --- Numeric Key Input (Only active in locked state) ---
            elif is_locked and key.isdigit():
                if len(current_input) < 8:  # Limit input to 8 digits (matches ID length)
                    current_input += key    # Append pressed digit to input buffer
                    update_display()        # Refresh OLED to show updated masked input

if __name__ == "__main__":
    main()