from machine import I2C, Pin, ADC, UART, PWM
from time import sleep, ticks_ms, ticks_diff
import dht
from umqtt_simple import MQTTClient
import network

WIFI_SSID = "Wokwi-GUEST"
WIFI_PASS = ""
MQTT_BROKER = "broker.hivemq.com"
MQTT_CLIENT_ID = "smart_helmet_pico"

TOPIC_FIRE = "helmet/fire_warning"
TOPIC_GAS = "helmet/gas_warning"
TOPIC_FALL = "helmet/fall_warning"
TOPIC_STATUS = "helmet/status"

def connect_wifi_mqtt():
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    wlan.connect(WIFI_SSID, WIFI_PASS)
    while not wlan.isconnected():
        sleep(0.5)
    
    client = MQTTClient(MQTT_CLIENT_ID, MQTT_BROKER)
    client.connect()
    return client

try:
    mqtt_client = connect_wifi_mqtt()
except:
    mqtt_client = None

i2c = I2C(0, sda=Pin(8), scl=Pin(9), freq=400000)
lcd_addr = 0x27

ir_sensor = Pin(15, Pin.IN)
button = Pin(14, Pin.IN, Pin.PULL_UP)

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

def sound_buzzer(freq=1000, duty=32768):
    buzzer.freq(freq)
    buzzer.duty_u16(duty)

def stop_buzzer():
    buzzer.duty_u16(0)

temp_sensor = dht.DHT22(Pin(26))
gas_adc = ADC(Pin(28))

MPU6050_ADDR = 0x68
try:
    i2c.writeto_mem(MPU6050_ADDR, 0x6B, b'\x00')
except:
    pass

gps_uart = UART(1, baudrate=9600, tx=Pin(4), rx=Pin(5))

system_active = False

def read_mpu6050():
    try:
        data = i2c.readfrom_mem(MPU6050_ADDR, 0x3B, 6)
        ax = int.from_bytes(data[0:2], 'big', signed=True) / 16384.0
        ay = int.from_bytes(data[2:4], 'big', signed=True) / 16384.0
        az = int.from_bytes(data[4:6], 'big', signed=True) / 16384.0
        return (ax**2 + ay**2 + az**2)**0.5
    except:
        return 1.0

def get_gps_location():
    try:
        if gps_uart.any():
            line = gps_uart.readline()
            if line and line.startswith(b'$GPRMC'):
                parts = line.decode('ascii', errors='ignore').split(',')
                if parts[2] == 'A':
                    return f"Lat:{parts[3]} Lon:{parts[5]}"
    except:
        pass
    return "Lat:30.788 Lat:31.002"

lcd = I2cLcd(i2c, lcd_addr, 2, 16)

lcd.clear()
lcd.putstr("Smart Helmet\nWaiting Helmet")
stop_buzzer()

while True:
    if not system_active:
        if ir_sensor.value() == 1:
            system_active = True
            lcd.clear()
            lcd.putstr("Helmet Worn!\nSystem Started")
            sound_buzzer(2000, 32768)
            sleep(0.2)
            stop_buzzer()
            if mqtt_client:
                try:
                    mqtt_client.publish(TOPIC_STATUS, "Helmet Worn - System Active")
                except:
                    pass
            sleep(1)
        else:
            sleep(1)
            continue

    try:
        temp_sensor.measure()
        temperature = temp_sensor.temperature()
        humidity = temp_sensor.humidity()
        
        gas_val = gas_adc.read_u16()
        spo2_percentage = int(100 - (gas_val / 65535.0) * 20)

        total_accel = read_mpu6050()
        gps_loc = get_gps_location()

        if temperature > 45:
            lcd.clear()
            lcd.putstr("WARNING!\nFire Predicted")
            sound_buzzer(2500, 40000)
            if mqtt_client:
                try:
                    mqtt_client.publish(TOPIC_FIRE, f"warning ⚠️ fire predicted at {gps_loc}")
                except:
                    pass
            while button.value() == 1:
                sleep(0.1)
            stop_buzzer()
            sleep(1)
            continue

        if spo2_percentage < 90:
            lcd.clear()
            lcd.putstr(f"Gas Leak SpO2\n{spo2_percentage}% < 90%")
            sound_buzzer(1500, 40000)
            if mqtt_client:
                try:
                    mqtt_client.publish(TOPIC_GAS, f"warning ⚠️ gas leakage (SpO2: {spo2_percentage}%) at {gps_loc}")
                except:
                    pass
            while button.value() == 1:
                sleep(0.1)
            stop_buzzer()
            sleep(1)
            continue

        if total_accel > 2.5:
            lcd.clear()
            lcd.putstr("FALL DETECTED!\nCancel in 5s")
            sound_buzzer(3000, 40000)
            
            start_time = ticks_ms()
            cancelled = False
            while ticks_diff(ticks_ms(), start_time) < 5000:
                if button.value() == 0:
                    cancelled = True
                    break
                sleep(0.1)
                
            if cancelled:
                lcd.clear()
                lcd.putstr("Fall Cancelled\nSystem Resumed")
                stop_buzzer()
                sleep(1)
            else:
                lcd.clear()
                lcd.putstr("FALL PREDICTED!\nEmail Sent")
                sound_buzzer(3000, 50000)
                if mqtt_client:
                    try:
                        mqtt_client.publish(TOPIC_FALL, f"warning ⚠️ fall predicted at {gps_loc}")
                    except:
                        pass
                sleep(3)
                stop_buzzer()
            continue

        lcd.clear()
        lcd.putstr(f"T:{temperature}C SpO2:{spo2_percentage}%\nActive & Safe")

        sleep(1.5)

    except Exception as e:
        lcd.clear()
        lcd.putstr("Sensor Error!")
        stop_buzzer()
        sleep(2)


class LcdApi:
    LCD_CLR = 0x01
    LCD_HOME = 0x02
    LCD_ENTRY_MODE = 0x04
    LCD_ENTRY_INC = 0x02
    LCD_ENTRY_SHIFT = 0x01
    LCD_ON_CTRL = 0x08
    LCD_ON_DISPLAY = 0x04
    LCD_ON_CURSOR = 0x02
    LCD_ON_BLINK = 0x01
    LCD_MOVE_DISP = 0x10
    LCD_MOVE_CURSOR = 0x00
    LCD_MOVE_RIGHT = 0x04
    LCD_MOVE_LEFT = 0x00
    LCD_FUNCTION = 0x20
    LCD_FUNCTION_8BIT = 0x10
    LCD_FUNCTION_2LINES = 0x08
    LCD_FUNCTION_5x10DOTS = 0x04

    def __init__(self, num_lines, num_columns):
        self.num_lines = num_lines
        if self.num_lines > 4:
            self.num_lines = 4
        self.num_columns = num_columns
        if self.num_columns > 40:
            self.num_columns = 40
        self.cursor_x = 0
        self.cursor_y = 0
        self.implied_newline = False

    def clear(self):
        self.hal_write_command(self.LCD_CLR)
        self.hal_write_command(self.LCD_HOME)
        self.cursor_x = 0
        self.cursor_y = 0

    def putstr(self, string):
        for char in string:
            if char == '\n':
                self.cursor_x = 0
                self.cursor_y = (self.cursor_y + 1) % self.num_lines
                self.move_to(self.cursor_x, self.cursor_y)
            else:
                self.hal_write_data(ord(char))
                self.cursor_x += 1
                if self.cursor_x >= self.num_columns:
                    self.cursor_x = 0
                    self.cursor_y = (self.cursor_y + 1) % self.num_lines
                    self.move_to(self.cursor_x, self.cursor_y)

    def move_to(self, cursor_x, cursor_y):
        self.cursor_x = cursor_x
        self.cursor_y = cursor_y
        addr = cursor_x & 0x3f
        if cursor_y & 1:
            addr += 0x40
        if cursor_y & 2:
            addr += self.num_columns
        self.hal_write_command(0x80 | addr)

    def hal_write_command(self, cmd):
        raise NotImplementedError

    def hal_write_data(self, data):
        raise NotImplementedError

MASK_RS = 0x01
MASK_RW = 0x02
MASK_E  = 0x04
SHIFT_BACKLIGHT = 3
SHIFT_DATA = 4

class I2cLcd(LcdApi):
    def __init__(self, i2c, i2c_addr, num_lines, num_columns):
        self.i2c = i2c
        self.i2c_addr = i2c_addr
        self.backlight = True
        self.i2c.writeto(self.i2c_addr, bytes([0]))
        sleep(0.02)
        self.hal_write_init_nibble(0x03)
        sleep(0.005)
        self.hal_write_init_nibble(0x03)
        sleep(0.001)
        self.hal_write_init_nibble(0x03)
        self.hal_write_init_nibble(0x02)
        super().__init__(num_lines, num_columns)
        cmd = self.LCD_FUNCTION | self.LCD_FUNCTION_2LINES
        self.hal_write_command(cmd)
        self.cmd_display = self.LCD_ON_CTRL | self.LCD_ON_DISPLAY
        self.hal_write_command(self.cmd_display)
        self.clear()
        self.hal_write_command(self.LCD_ENTRY_MODE | self.LCD_ENTRY_INC)

    def hal_write_init_nibble(self, nibble):
        byte = (nibble << SHIFT_DATA) | (int(self.backlight) << SHIFT_BACKLIGHT)
        self.i2c.writeto(self.i2c_addr, bytes([byte | MASK_E]))
        self.i2c.writeto(self.i2c_addr, bytes([byte]))

    def hal_write_command(self, cmd):
        byte = (self.backlight << SHIFT_BACKLIGHT)
        self.i2c.writeto(self.i2c_addr, bytes([byte | (cmd & 0xF0) | MASK_E]))
        self.i2c.writeto(self.i2c_addr, bytes([byte | (cmd & 0xF0)]))
        self.i2c.writeto(self.i2c_addr, bytes([byte | ((cmd << 4) & 0xF0) | MASK_E]))
        self.i2c.writeto(self.i2c_addr, bytes([byte | ((cmd << 4) & 0xF0)]))
        if cmd <= 3:
            sleep(0.005)

    def hal_write_data(self, data):
        byte = (self.backlight << SHIFT_BACKLIGHT) | MASK_RS
        self.i2c.writeto(self.i2c_addr, bytes([byte | (data & 0xF0) | MASK_E]))
        self.i2c.writeto(self.i2c_addr, bytes([byte | (data & 0xF0)]))
        self.i2c.writeto(self.i2c_addr, bytes([byte | ((data << 4) & 0xF0) | MASK_E]))
        self.i2c.writeto(self.i2c_addr, bytes([byte | ((data << 4) & 0xF0)]))