"""
NETPIE Smart Door Lock System — MicroPython (ESP32)
====================================================

System Overview
---------------
This firmware runs on an ESP32 and implements a smart door lock with the
following features:

  1. DHT22 Sensor
     Reads temperature and humidity every 3 seconds. If either value exceeds
     the configured threshold an alert is published to NETPIE (@msg/alert) and
     the OLED shows a warning. A 30-second cooldown prevents alert flooding.
     Alert display is suppressed while the user is typing a PIN so keypad
     input is never interrupted. A compact "!ALERT on sensor" line appears on
     the normal screen so the warning is always visible alongside door status.

  2. 4x4 Matrix Keypad — PIN Entry
     The user types a 4-digit PIN followed by '#' to unlock, or '*' to clear.
     After MAX_ATTEMPTS (3) consecutive wrong PINs the system enters a timed
     lockout (LOCKOUT_SEC = 30 s). The lockout can be reset remotely via
     NETPIE.

  3. Physical Button
     A momentary push-button on BUTTON_PIN provides a local one-press unlock,
     identical in behaviour to a correct PIN entry (auto-locks after
     UNLOCK_SEC).

  4. Servo Motor
     Controls the physical latch: 0 degrees = locked, 90 degrees = unlocked.
     PWM is diablesd after every move to prevent jitter and motor heating.

  5. Auto-Lock Timer
     After any unlock event (keypad, button, or remote command) the door
     relocks automatically after UNLOCK_SEC (5 s). The timer is tracked in
     the main loop so remote unlocks are covered as well.

  6. NETPIE MQTT Integration
     - Publishes device state to @shadow/data/update every 3 seconds.
     - Publishes structured log entries to @msg/log on every state change.
     - Publishes alerts tos @mg/alert on threshold violations or wrong PINs.
     - Subscribes to @msg/control for remote commands:
         {"cmd": "unlock"}         remote unlock
         {"cmd": "lock"}           remote lock
         {"cmd": "reset_lockout"}  clear failed-attempt counter

  7. OLED Display (128x64 SSD1306)
     Shows sensor readings, door status, PIN-entry feedback, and alert
     messages. Switches automatically between normal view and PIN-entry view.

  8. Watchdog Timer
     An 8-second hardware watchdog resets the ESP32 if the main loop stalls.
     The WDT is initialised after WiFi and MQTT setup completes to prevent a
     boot loop caused by slow network connection during startup.

  9. WiFi Retry Logic
     The system attempts to connect up to 5 times before continuing in offline
     mode. Core functions (keypad, servo, OLED, LEDs) remain available without
     network access.

Hardware Pins
-------------
  DHT22        GPIO 4
  OLED SDA     GPIO 5
  OLED SCL     GPIO 22
  Button       GPIO 25  (active-low, internal pull-up)
  Green LED    GPIO 23  (on = unlocked)
  Red LED      GPIO 19  (on = locked)
  Servo PWM    GPIO 15
  Keypad rows  GPIO 13, 12, 14, 27
  Keypad cols  GPIO 26, 33, 32, 35
  Note: GPIO 35 is input-only and has no internal pull-up.
        A 10 kΩ external pull-up resistor to 3.3 V is required on that pin.
"""

import network                        
import time
import json
from machine import Pin, PWM, SoftI2C, WDT
import dht
import ssd1306
from umqtt.simple import MQTTClient

# =============================================================================
# WiFi credentials
# =============================================================================
WIFI_SSID = "Wokwi-GUEST"
WIFI_PASS = ""

# =============================================================================
# NETPIE MQTT configuration
# CLIENT_ID, TOKEN, and SECRET are issued by the NETPIE platform per device.
# =============================================================================
CLIENT_ID = "f543dfc9-c789-449e-af2d-39fb9b696436"
TOKEN     = "HXouaBe9FDLeRNLwyt2xBmEffRGAnT5S"
SECRET    = "YL34CPTXqtnqEYmDSvVQDAgZehwY9QKr"

MQTT_BROKER = "mqtt.netpie.io"
MQTT_PORT   = 1883

# NETPIE topic definitions
TOPIC_SHADOW_UPDATE = "@shadow/data/update"   # publish device state
TOPIC_SHADOW_GET    = "@shadow/data/get"      # request current shadow
TOPIC_MSG_CONTROL   = "@msg/control"          # subscribe for remote commands
TOPIC_MSG_ALERT     = "@msg/alert"            # publish threshold alerts
TOPIC_MSG_LOG       = "@msg/log"              # publish event logs

# =============================================================================
# GPIO pin assignments
# =============================================================================
DHT_PIN    = 4
OLED_SDA   = 5
OLED_SCL   = 22
BUTTON_PIN = 25
LED_GREEN  = 23
LED_RED    = 19
SERVO_PIN  = 15

# Keypad row pins (driven LOW one at a time during scan)
KP_ROWS = [13, 12, 14, 27]

# Keypad column pins (read with pull-up; GPIO 35 has no internal pull-up)
KP_COLS = [26, 33, 32, 35]

# =============================================================================
# Keypad character map — 4 rows x 4 columns
# =============================================================================
KEYMAP = [
    ["1", "2", "3", "A"],
    ["4", "5", "6", "B"],
    ["7", "8", "9", "C"],
    ["*", "0", "#", "D"],
]

# =============================================================================
# Security configuration
# =============================================================================
CORRECT_PIN  = "1234"   # expected 4-digit PIN
MAX_ATTEMPTS = 3        # consecutive wrong attempts before lockout
LOCKOUT_SEC  = 30       # lockout duration in seconds
UNLOCK_SEC   = 5        # seconds before auto-relock after unlock

# =============================================================================
# Sensor threshold and validation configuration
# =============================================================================
TEMP_THRESHOLD = 35.0    # degrees C — trigger alert if exceeded
HUM_THRESHOLD  = 80.0    # percent  — trigger alert if exceeded
TEMP_VALID_MIN = -10.0   # degrees C — discard readings below this value
TEMP_VALID_MAX = 80.0    # degrees C — discard readings above this value
HUM_VALID_MIN  = 0.0     # percent  — discard readings below this value
HUM_VALID_MAX  = 100.0   # percent  — discard readings above this value
ALERT_COOLDOWN = 30000   # ms — minimum interval between repeated alerts (k)

# =============================================================================
# Hardware initialisation (L)
# =============================================================================
dht_sensor = dht.DHT22(Pin(DHT_PIN))
i2c        = SoftI2C(sda=Pin(OLED_SDA), scl=Pin(OLED_SCL))
oled       = ssd1306.SSD1306_I2C(128, 64, i2c)
button     = Pin(BUTTON_PIN, Pin.IN, Pin.PULL_UP)
green_led  = Pin(LED_GREEN, Pin.OUT)
red_led    = Pin(LED_RED,   Pin.OUT)
servo      = PWM(Pin(SERVO_PIN), freq=50)

# Keypad row pins — configured as outputs (active-low scan)
row_pins = [Pin(p, Pin.OUT) for p in KP_ROWS]

# Keypad column pins
# GPIO 35 is input-only on ESP32 and has no internal pull-up resistor.
# An external 10 kΩ pull-up resistor to 3.3 V is required on that line.
col_pins = []
for p in KP_COLS:
    if p == 35:
        col_pins.append(Pin(p, Pin.IN))              # external pull-up only
    else:
        col_pins.append(Pin(p, Pin.IN, Pin.PULL_UP)) # internal pull-up

# =============================================================================
# Global state variables
# =============================================================================
temperature     = 0.0
humidity        = 0.0
door_status     = "LOCKED"
alert_active    = False
client          = None            # MQTTClient instance; None when disconnected
last_alert_time = {}              # {reason: ticks_ms} cooldown tracking dict

# PIN / keypad state
pin_input       = ""              # digits entered so far this session
failed_attempts = 0               # consecutive wrong PIN attempts
lockout_until   = 0               # ticks_ms when lockout expires (0 = none)
last_key_time   = 0               # ticks_ms of the last accepted key press
last_key        = None            # previously scanned key for edge detection

# Auto-lock timer — holds ticks_ms of the unlock event; 0 when idle
unlock_time     = 0

# Button debounce state
last_button_time = 0
last_button_val  = 1              # default HIGH (not pressed, active-low input)

# ============================================================================= (k)
# WiFi connection
# Attempts to connect up to max_retries times before returning False and
# continuing in offline mode. Avoids machine.reset() to prevent a boot loop
# when the watchdog timer is active.
# =============================================================================
def connect_wifi(max_retries=5):
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    if wlan.isconnected():
        return True

    for attempt in range(1, max_retries + 1):
        print("Connecting WiFi (attempt {}/{})".format(attempt, max_retries), end="")
        wlan.connect(WIFI_SSID, WIFI_PASS)
        timeout = 20                            # 20 x 0.5 s = 10 s per attempt
        while not wlan.isconnected() and timeout > 0:
            print(".", end="")
            time.sleep(0.5)
            timeout -= 1
        if wlan.isconnected():
            print("\nWiFi Connected:", wlan.ifconfig())
            return True
        print("\nAttempt {} failed, retrying...".format(attempt))
        wlan.disconnect()
        time.sleep(2)

    print("WiFi FAILED after {} attempts — running offline".format(max_retries))
    return False

# =============================================================================(L)
# MQTT message callback
# Invoked by umqtt whenever a message arrives on a subscribed topic.
# Handles three remote commands: unlock, lock, and reset_lockout.
# =============================================================================
def on_message(topic, msg):
    try:
        topic_str = topic.decode("utf-8")
        data = json.loads(msg)
        print("MSG received on", topic_str, ":", data)
        if topic_str == TOPIC_MSG_CONTROL:
            cmd = data.get("cmd", "")
            if cmd == "unlock":
                unlock_door(source="remote")
                show_oled()
                safe_publish_shadow()
            elif cmd == "lock":
                global unlock_time
                unlock_time = 0              # cancel any pending auto-lock timer
                lock_door(source="remote")
                show_oled()
                safe_publish_shadow()
            elif cmd == "reset_lockout":
                reset_lockout(source="remote")
    except Exception as e:
        print("Callback error:", e)

# =============================================================================
# NETPIE MQTT connection helpers
# =============================================================================
def connect_netpie():
    """Create the MQTT client, connect to the broker, and subscribe to control topic."""
    global client
    try:
        client = MQTTClient(
            CLIENT_ID, MQTT_BROKER,
            port=MQTT_PORT, user=TOKEN,
            password=SECRET, keepalive=60
        )
        client.set_callback(on_message)
        client.connect()
        client.subscribe(TOPIC_MSG_CONTROL)
        print("NETPIE Connected & Subscribed")
        return True
    except Exception as e:
        print("NETPIE connect error:", e)
        client = None
        return False

def ensure_connected():
    """Re-establish WiFi and MQTT if the connection has been lost."""
    global client
    if client is None:
        print("Reconnecting NETPIE...")
        wifi_ok = connect_wifi(max_retries=2)
        if wifi_ok:
            connect_netpie()

# =============================================================================(k)
# Servo control
# Converts an angle (0–180 degrees) to a PWM duty cycle, moves the servo, then
# disables the PWM signal. Disabling PWM after each move prevents jitter and
# reduces motor heating while the latch holds its position.
# =============================================================================
def set_servo(angle):
    duty = int((angle / 180) * 77 + 26)
    servo.duty(duty)
    time.sleep(0.5)   # allow the servo to physically reach the target angle
    servo.duty(0)     # disable PWM output to stop jitter

# =============================================================================(L)
# Door control
# =============================================================================
def lock_door(source="auto"):
    """Lock the door: rotate servo to 0 degrees, activate red LED, log event."""
    global door_status
    door_status = "LOCKED"
    set_servo(0)
    red_led.value(1)
    green_led.value(0)
    print("Door LOCKED by:", source)
    send_log("LOCKED", source)

def unlock_door(source="pin"):
    """Unlock the door: rotate servo to 90 degrees, activate green LED, start auto-lock timer."""
    global door_status, unlock_time
    door_status = "UNLOCKED"
    unlock_time = time.ticks_ms()   # record the unlock timestamp for auto-relock
    set_servo(90)
    red_led.value(0)
    green_led.value(1)
    print("Door UNLOCKED by:", source)
    send_log("UNLOCKED", source)

# =============================================================================    (L)
# Keypad scan
# Drives each row LOW in turn and reads all column inputs. Returns the matched
# character from KEYMAP, or None if no key is currently pressed. 
# =============================================================================
def scan_keypad():
    for r, row_pin in enumerate(row_pins):
        row_pin.value(0)
        for c, col_pin in enumerate(col_pins):
            if col_pin.value() == 0:
                row_pin.value(1)
                return KEYMAP[r][c]
        row_pin.value(1)
    return None

# =============================================================================(L)
# Lockout helpers
# =============================================================================
def is_locked_out():
    """Return True if the lockout period has not yet expired."""
    return time.ticks_diff(lockout_until, time.ticks_ms()) > 0

def lockout_remaining_sec():
    """Return the number of seconds remaining in the current lockout (0 if none)."""
    diff = time.ticks_diff(lockout_until, time.ticks_ms())
    return max(0, diff // 1000)

def reset_lockout(source="auto"):
    """Clear the failed-attempt counter, cancel the lockout, and update outputs."""
    global failed_attempts, lockout_until, pin_input
    failed_attempts = 0
    lockout_until   = 0
    pin_input       = ""
    print("Lockout reset by:", source)
    send_log("LOCKOUT_RESET", source)
    show_oled()
    safe_publish_shadow()

# =============================================================================(L)
# Keypad key handler
# Processes a single key press based on the current PIN-entry state machine.
# '*' clears the buffer, '#' validates, digits accumulate up to PIN length.
# =============================================================================
def handle_key(key):
    global pin_input, failed_attempts, lockout_until

    if is_locked_out():
        show_oled_pin(msg="!! LOCKED OUT !!")
        return

    if key == "*":
        pin_input = ""
        show_oled_pin(msg="Cleared")
        return

    if key == "#":
        if pin_input == CORRECT_PIN:
            failed_attempts = 0
            pin_input       = ""
            print("PIN correct - unlocking")
            send_log("PIN_OK", "keypad")
            unlock_door(source="pin")
            show_oled_pin(msg="** ACCESS OK! **")
            safe_publish_shadow()
        else:
            failed_attempts += 1
            pin_input        = ""
            print("PIN wrong! Attempts:", failed_attempts)
            send_log("PIN_FAIL", "keypad")
            send_alert("WRONG_PIN", failed_attempts)

            if failed_attempts >= MAX_ATTEMPTS:
                lockout_until = time.ticks_add(time.ticks_ms(), LOCKOUT_SEC * 1000)
                print("LOCKOUT! Wait", LOCKOUT_SEC, "sec")
                send_alert("LOCKOUT", failed_attempts)
                send_log("LOCKOUT", "auto")
                show_oled_pin(msg="!!! LOCKOUT !!!")
                safe_publish_shadow()
            else:
                show_oled_pin(msg="WRONG {}/{}".format(failed_attempts, MAX_ATTEMPTS))
        return

    # Accumulate digit into the PIN input buffer
    if key.isdigit() and len(pin_input) < len(CORRECT_PIN):
        pin_input += key
        masked = "*" * len(pin_input) + "_" * (len(CORRECT_PIN) - len(pin_input))
        show_oled_pin(msg="PIN: " + masked)

# =============================================================================(L)
# Physical button handler
# Detects a falling edge on the active-low button input with 300 ms software
# debounce. Triggers an immediate unlock when the system is not in lockout.
# =============================================================================
def handle_button():
    global last_button_time, last_button_val
    now = time.ticks_ms()
    val = button.value()

    if val == 0 and last_button_val == 1:                # falling edge (press)
        if time.ticks_diff(now, last_button_time) > 300: # debounce guard
            if not is_locked_out():
                print("Button pressed - unlocking")
                send_log("BUTTON_PRESS", "button")
                unlock_door(source="button")
                show_oled()
                safe_publish_shadow()
            last_button_time = now

    last_button_val = val

# =============================================================================(L)
# OLED — PIN entry view
# Rendered whenever the user is actively entering a PIN or is locked out.
# Shows countdown and recovery instructions during lockout.
# =============================================================================
def show_oled_pin(msg="Enter PIN:"):
    oled.fill(0)
    oled.text("== PIN ENTRY ==", 0, 0)
    oled.text(msg, 0, 20)
    if is_locked_out():
        oled.text("Wait: {}s".format(lockout_remaining_sec()), 0, 36)
        oled.text("Or reset remotely", 0, 48)
    else:
        oled.text("# Confirm * Clear", 0, 48)
    oled.show()

# =============================================================================(L)
# OLED — normal status view
# Displays sensor readings, door status, and lockout or alert indicators.
# When alert_msg is provided, renders a full-screen alert overlay instead.
# When alert_active is True but no alert overlay is shown, a compact tag is
# added on the bottom line so the warning is always visible.
# =============================================================================
def show_oled(alert_msg=None):
    oled.fill(0)
    if alert_msg:
        oled.text("!! ALERT !!", 16, 0)
        oled.text(alert_msg, 0, 16)
        oled.text("Temp:{:.1f}C".format(temperature), 0, 32)
        oled.text("Hum :{:.1f}%".format(humidity),    0, 44)
    else:
        oled.text("NETPIE Door Lock", 0, 0)
        oled.text("Temp: {:.1f} C".format(temperature), 0, 16)
        oled.text("Hum : {:.1f} %".format(humidity),    0, 28)
        oled.text("Door: " + door_status, 0, 40)
        if is_locked_out():
            oled.text("LOCKOUT {}s".format(lockout_remaining_sec()), 0, 52)
        elif alert_active:
            oled.text("!ALERT on sensor", 0, 52)
    oled.show()

# ============================================================================= (k)
# DHT22 sensor reading
# Calls measure() then validates each value against the configured range before
# updating globals. Out-of-range values are logged and ignored.
# =============================================================================
def read_sensor():
    global temperature, humidity
    try:
        dht_sensor.measure()
        t = dht_sensor.temperature()
        h = dht_sensor.humidity()
        if TEMP_VALID_MIN <= t <= TEMP_VALID_MAX:
            temperature = t
        else:
            print("DHT temp out of range, ignored:", t)
        if HUM_VALID_MIN <= h <= HUM_VALID_MAX:
            humidity = h
        else:
            print("DHT hum out of range, ignored:", h)
    except Exception as e:
        print("DHT Error:", e)

# =============================================================================(k)
# NETPIE log publisher
# Sends a structured event log entry (event, source, timestamp) to @msg/log.
# Silently skips if there is no active MQTT connection.
# =============================================================================
def send_log(event, source):
    if client is None:
        return
    try:
        log = {"event": event, "source": source, "tick": time.ticks_ms()}
        client.publish(TOPIC_MSG_LOG, json.dumps(log))
        print("Log sent:", log)
    except Exception as e:
        print("Log error:", e)

# =============================================================================(k)
# NETPIE alert publisher
# Publishes an alert message to @msg/alert, subject to per-reason cooldown.
# Prunes tracking-dict entries older than 10 minutes to bound memory usage.
# =============================================================================
def send_alert(reason, value):
    global last_alert_time
    if client is None:
        return
    now  = time.ticks_ms()
    last = last_alert_time.get(reason, 0)
    if time.ticks_diff(now, last) < ALERT_COOLDOWN:
        return
    try:
        alert = {"alert": reason, "value": value}
        client.publish(TOPIC_MSG_ALERT, json.dumps(alert))
        last_alert_time[reason] = now
        print("Alert sent:", alert)

        # Prune entries older than 10 minutes to prevent unbounded dict growth
        stale = [k for k, v in last_alert_time.items()
                 if time.ticks_diff(now, v) > 600_000]
        for k in stale:
            del last_alert_time[k]
    except Exception as e:
        print("Alert error:", e)

# =============================================================================(L)
# Threshold check
# Compares temperature and humidity against their thresholds. Sets alert_active
# and publishes to NETPIE when a threshold is exceeded. Skips the OLED update
# while the user is entering a PIN to avoid disrupting the PIN-entry view.
# =============================================================================
def check_thresholds():
    global alert_active
    triggered = False
    alert_msg  = None

    if temperature > TEMP_THRESHOLD:
        triggered = True
        alert_msg = "HIGH TEMP!"
        send_alert("HIGH_TEMP", temperature)

    if humidity > HUM_THRESHOLD:
        triggered = True
        alert_msg = "HIGH HUMID!"
        send_alert("HIGH_HUM", humidity)

    alert_active = triggered

    if pin_input == "":
        if alert_msg:
            show_oled(alert_msg=alert_msg)

# =============================================================================(L)
# Shadow publisher
# Builds a full device-state payload and publishes it to the NETPIE shadow
# endpoint. Sets client to None on failure so the reconnection logic activates
# on the next main-loop iteration.
# =============================================================================
def safe_publish_shadow():
    global client
    if client is None:
        return
    try:
        data = {
            "data": {
                "temperature":     temperature,
                "humidity":        humidity,
                "door_status":     door_status,
                "servo_angle":     90 if door_status == "UNLOCKED" else 0,
                "red_led":         red_led.value(),
                "green_led":       green_led.value(),
                "alert_active":    alert_active,
                "failed_attempts": failed_attempts,
                "is_locked_out":   is_locked_out(),
                "lockout_remain":  lockout_remaining_sec(),
            }
        }
        client.publish(TOPIC_SHADOW_UPDATE, json.dumps(data))
        print("Shadow sent:", data)
    except Exception as e:
        print("Shadow error:", e)
        client = None

# =============================================================================(L)
# Startup sequence
# WiFi and MQTT are established before the watchdog is created so that a slow
# network connection during boot does not trigger a hardware reset.
# =============================================================================
connect_wifi()
connect_netpie()

wdt = WDT(timeout=8000)   # 8-second hardware watchdog, started after setup

lock_door(source="boot")
read_sensor()
show_oled()
safe_publish_shadow()

last_sensor_time = 0
last_shadow_time = 0

# =============================================================================
# Main loop
# =============================================================================
while True:
    now = time.ticks_ms()

    wdt.feed()          # reset the watchdog counter on every iteration

    ensure_connected()  # reconnect MQTT if the connection was lost

    # Poll for incoming MQTT messages (non-blocking check)
    if client is not None:
        try:
            client.check_msg()
        except Exception as e:
            print("check_msg error:", e)
            client = None

    # Scan and handle the physical button with edge detection and debounce
    handle_button()

    # Scan the keypad; process a new key only after the inter-key guard elapses
    key = scan_keypad()
    if key is not None and key != last_key:
        if time.ticks_diff(now, last_key_time) > 200:
            print("Key pressed:", key)
            handle_key(key)
            last_key_time = now
    last_key = key

    # Auto-lock: relock the door after UNLOCK_SEC seconds from any unlock source
    if door_status == "UNLOCKED" and unlock_time != 0:
        if time.ticks_diff(now, unlock_time) >= UNLOCK_SEC * 1000:
            unlock_time = 0
            lock_door(source="auto")
            show_oled()
            safe_publish_shadow()

    # Read sensor and check thresholds every 3 seconds
    if time.ticks_diff(now, last_sensor_time) > 3000:
        read_sensor()
        check_thresholds()
        if not alert_active and pin_input == "":
            show_oled()
        last_sensor_time = now

    # Publish the device shadow every 3 seconds
    if time.ticks_diff(now, last_shadow_time) > 3000:
        safe_publish_shadow()
        last_shadow_time = now

    time.sleep(0.05)    # 50 ms yield keeps CPU load low without missing events