from machine import Pin, PWM
from time import sleep_ms, ticks_ms, ticks_diff
import ujson
import uhashlib

# -----------------------------
# CONFIGURATION
# -----------------------------
MAX_ATTEMPTS = 3
BASE_LOCKOUT_MS = 5000
AUTO_LOCK_MS = 15000
INPUT_TIMEOUT = 5000

# -----------------------------
# HASH FUNCTION
# -----------------------------
def hash_code(code):
    return uhashlib.sha256(code.encode()).digest()

# -----------------------------
# LOAD / SAVE CODES
# -----------------------------
def load_codes():
    try:
        with open("codes.json", "r") as f:
            return ujson.load(f)
    except:
        return {
            "user": "1234",
            "master": "9876",
            "duress": "4321"
        }

def save_codes(data):
    with open("codes.json", "w") as f:
        ujson.dump(data, f)

codes = load_codes()

# -----------------------------
# DISPLAY ALL CODES AT START
# -----------------------------
print("""
==============================
  MULTI-LEVEL DIGITAL SAFE
==============================
USER CODE   : {}
MASTER CODE : {}
DURESS CODE : {}
==============================
""".format(codes["user"], codes["master"], codes["duress"]))

# -----------------------------
# HASHED VALUES
# -----------------------------
USER_CODE_HASH = hash_code(codes["user"])
MASTER_CODE_HASH = hash_code(codes["master"])
DURESS_CODE_HASH = hash_code(codes["duress"])

# -----------------------------
# OUTPUT DEVICES
# -----------------------------
green_led = Pin(12, Pin.OUT)
red_led = Pin(13, Pin.OUT)
hidden_alert = Pin(14, Pin.OUT)

servo = PWM(Pin(10))
servo.freq(50)

buzzer = PWM(Pin(11))
buzzer.duty_u16(0)

# -----------------------------
# KEYPAD
# -----------------------------
row_pins = [Pin(2, Pin.OUT), Pin(3, Pin.OUT), Pin(4, Pin.OUT), Pin(5, Pin.OUT)]
col_pins = [Pin(6, Pin.IN, Pin.PULL_DOWN),
            Pin(7, Pin.IN, Pin.PULL_DOWN),
            Pin(8, Pin.IN, Pin.PULL_DOWN),
            Pin(9, Pin.IN, Pin.PULL_DOWN)]

keys = [
    ['1','2','3','A'],
    ['4','5','6','B'],
    ['7','8','9','C'],
    ['*','0','#','D']
]

# -----------------------------
# STATE VARIABLES
# -----------------------------
entered = ""
attempts = 0
lockout_until = 0
safe_unlocked = False
master_mode = False
change_user_mode = False

last_input_time = 0
last_unlock_time = 0

# -----------------------------
# HARDWARE FUNCTIONS
# -----------------------------
def set_servo_angle(angle):
    pulse = 500 + int((angle / 180) * 2000)
    duty = int(pulse * 65535 / 20000)
    servo.duty_u16(duty)

def lock_safe():
    global safe_unlocked, master_mode, entered
    set_servo_angle(0)
    safe_unlocked = False
    master_mode = False
    entered = ""
    green_led.off()
    red_led.on()
    hidden_alert.off()
    print("SAFE LOCKED")

def unlock_safe():
    global safe_unlocked, last_unlock_time
    set_servo_angle(90)
    safe_unlocked = True
    green_led.on()
    red_led.off()
    last_unlock_time = ticks_ms()
    print("SAFE UNLOCKED")

def beep(freq=1200, dur=100):
    buzzer.freq(freq)
    buzzer.duty_u16(20000)
    sleep_ms(dur)
    buzzer.duty_u16(0)

def error_tone():
    beep(500, 250)

def success_tone():
    beep(1200, 80)
    sleep_ms(50)
    beep(1600, 100)

# -----------------------------
# KEYPAD SCAN
# -----------------------------
def scan_keypad():
    for row in row_pins:
        row.low()

    for r, row in enumerate(row_pins):
        row.high()
        sleep_ms(2)

        for c, col in enumerate(col_pins):
            if col.value():
                key = keys[r][c]
                while col.value():
                    sleep_ms(10)
                sleep_ms(120)
                return key
        row.low()
    return None

# -----------------------------
# PROCESS CODE
# -----------------------------
def process_code(code):
    global attempts, lockout_until, master_mode, change_user_mode
    global USER_CODE_HASH

    h = hash_code(code)

    if change_user_mode:
        if len(code) == 4 and code.isdigit():
            USER_CODE_HASH = h
            codes["user"] = code
            save_codes(codes)
            print("NEW USER CODE SAVED")
            success_tone()
        else:
            print("INVALID CODE")
            error_tone()
        change_user_mode = False
        return

    if h == USER_CODE_HASH:
        print("USER ACCESS")
        unlock_safe()
        success_tone()
        attempts = 0
        master_mode = False
        hidden_alert.off()

    elif h == MASTER_CODE_HASH:
        print("MASTER ACCESS")
        unlock_safe()
        success_tone()
        attempts = 0
        master_mode = True
        print("PRESS B TO CHANGE USER CODE")

    elif h == DURESS_CODE_HASH:
        print("DURESS ACTIVATED")
        unlock_safe()
        success_tone()
        hidden_alert.on()

    else:
        attempts += 1
        print("WRONG CODE:", attempts)
        error_tone()
        sleep_ms(300)

        if attempts >= MAX_ATTEMPTS:
            lockout_until = ticks_ms() + (BASE_LOCKOUT_MS * attempts)
            print("SYSTEM LOCKED")
            attempts = 0

# -----------------------------
# INIT
# -----------------------------
lock_safe()
print("SYSTEM READY\n")

# -----------------------------
# MAIN LOOP
# -----------------------------
while True:
    now = ticks_ms()

    # Lockout handling
    if ticks_diff(lockout_until, now) > 0:
        red_led.toggle()
        beep(700, 50)
        sleep_ms(100)
        continue

    # Auto-lock
    if safe_unlocked and ticks_diff(now, last_unlock_time) > AUTO_LOCK_MS:
        print("AUTO LOCK")
        lock_safe()

    # Input timeout
    if entered and ticks_diff(now, last_input_time) > INPUT_TIMEOUT:
        entered = ""
        print("INPUT TIMEOUT")

    key = scan_keypad()

    if key:
        last_input_time = ticks_ms()
        print("KEY:", key)

        if key == '*':
            entered = ""
            print("CLEARED")

        elif key == 'A':
            lock_safe()

        elif key == 'B':
            if master_mode and safe_unlocked:
                change_user_mode = True
                entered = ""
                print("ENTER NEW 4-DIGIT USER CODE")
            else:
                error_tone()

        elif key == '#':
            if entered:
                process_code(entered)
                entered = ""
            else:
                error_tone()

        elif key.isdigit():
            if len(entered) < 4:
                entered += key
                print("*" * len(entered))
                beep(1000, 40)
            else:
                error_tone()

    sleep_ms(20)