from machine import Pin, I2C, ADC
import time
import framebuf
import dht

# SSD1306 OLED Driver (128x64)
class SSD1306_I2C:
    def __init__(self, width, height, i2c, addr=0x3C):
        self.i2c = i2c
        self.addr = addr
        self.width = width
        self.height = height
        self.pages = height // 8
        self.buffer = bytearray(self.pages * width)
        self.framebuf = framebuf.FrameBuffer(self.buffer, width, height, framebuf.MONO_VLSB)
        self.init_display()
    
    def init_display(self):
        for cmd in (
            0xAE, 0xD5, 0x80, 0xA8, 0x3F, 0xD3, 0x00, 0x40, 0x8D, 0x14,
            0x20, 0x00, 0xA1, 0xC8, 0xDA, 0x12, 0x81, 0xCF, 0xD9, 0xF1,
            0xDB, 0x40, 0xA4, 0xA6, 0xAF
        ):
            self.write_cmd(cmd)
        self.fill(0)
        self.show()
    
    def write_cmd(self, cmd):
        self.i2c.writeto(self.addr, bytearray([0x00, cmd]))
    
    def write_data(self, buf):
        self.i2c.writeto(self.addr, b'\x40' + buf)
    
    def show(self):
        for page in range(self.pages):
            self.write_cmd(0xB0 + page)
            self.write_cmd(0x00)
            self.write_cmd(0x10)
            self.write_data(self.buffer[page * self.width:(page + 1) * self.width])
    
    def fill(self, color):
        self.framebuf.fill(color)
    
    def text(self, string, x, y, color=1):
        self.framebuf.text(string, x, y, color)
    
    def rect(self, x, y, w, h, color):
        self.framebuf.rect(x, y, w, h, color)
    
    def fill_rect(self, x, y, w, h, color):
        self.framebuf.fill_rect(x, y, w, h, color)

# Initialize Components
i2c = I2C(0, scl=Pin(1), sda=Pin(0), freq=400000)
oled = SSD1306_I2C(128, 64, i2c)

# DHT22 Temperature Sensor
dht_sensor = dht.DHT22(Pin(15))

# LED Indicators (Smart Devices)
led_living = Pin(10, Pin.OUT)
led_bedroom = Pin(11, Pin.OUT)
led_kitchen = Pin(12, Pin.OUT)

# Buzzer for Audio Feedback
buzzer = Pin(13, Pin.OUT)

# Wake Button
wake_button = Pin(14, Pin.IN, Pin.PULL_UP)

# System State
devices = {
    "living_room": False,
    "bedroom": False,
    "kitchen": False
}
temperature = 0.0
humidity = 0.0
listening = False

# Text-to-Speech Simulator (Buzzer Tones)
def speak(message, tone_pattern="short"):
    """Simulate speech with buzzer tones and display text"""
    print(f"🔊 Assistant: {message}")
    
    # Display on OLED
    oled.fill(0)
    oled.rect(0, 0, 128, 64, 1)
    oled.text("ASSISTANT", 30, 5)
    oled.fill_rect(0, 16, 128, 1, 1)
    
    # Split message into lines (max 16 chars per line)
    words = message.split()
    lines = []
    current_line = ""
    
    for word in words:
        if len(current_line + word) < 16:
            current_line += word + " "
        else:
            lines.append(current_line.strip())
            current_line = word + " "
    if current_line:
        lines.append(current_line.strip())
    
    # Display up to 4 lines
    y_pos = 25
    for line in lines[:4]:
        oled.text(line, 5, y_pos)
        y_pos += 10
    
    oled.show()
    
    # Audio feedback
    if tone_pattern == "short":
        buzzer.on()
        time.sleep(0.1)
        buzzer.off()
    elif tone_pattern == "long":
        for _ in range(2):
            buzzer.on()
            time.sleep(0.15)
            buzzer.off()
            time.sleep(0.05)
    elif tone_pattern == "alert":
        for _ in range(3):
            buzzer.on()
            time.sleep(0.05)
            buzzer.off()
            time.sleep(0.05)
    
    time.sleep(2)

# Update Temperature
def update_temperature():
    global temperature, humidity
    try:
        dht_sensor.measure()
        temperature = dht_sensor.temperature()
        humidity = dht_sensor.humidity()
    except:
        pass

# Display Idle Screen
def show_idle_screen():
    oled.fill(0)
    oled.rect(0, 0, 128, 64, 1)
    oled.text("AI ASSISTANT", 20, 5)
    oled.fill_rect(0, 16, 128, 1, 1)
    oled.text("Say 'Hey'", 30, 25)
    oled.text("or", 52, 35)
    oled.text("Press Button", 20, 45)
    oled.show()

# Display Listening Screen
def show_listening_screen():
    oled.fill(0)
    oled.rect(0, 0, 128, 64, 1)
    oled.text("LISTENING...", 25, 20)
    oled.fill_rect(20, 35, 88, 8, 1)
    oled.show()

# Command Processing
def process_command(command):
    """Parse and execute voice commands"""
    cmd = command.lower().strip()
    
    # Device Control Commands
    if "living" in cmd or "living room" in cmd:
        if "on" in cmd:
            led_living.on()
            devices["living_room"] = True
            speak("Living room on", "short")
        elif "off" in cmd:
            led_living.off()
            devices["living_room"] = False
            speak("Living room off", "short")
    
    elif "bedroom" in cmd:
        if "on" in cmd:
            led_bedroom.on()
            devices["bedroom"] = True
            speak("Bedroom on", "short")
        elif "off" in cmd:
            led_bedroom.off()
            devices["bedroom"] = False
            speak("Bedroom off", "short")
    
    elif "kitchen" in cmd:
        if "on" in cmd:
            led_kitchen.on()
            devices["kitchen"] = True
            speak("Kitchen on", "short")
        elif "off" in cmd:
            led_kitchen.off()
            devices["kitchen"] = False
            speak("Kitchen off", "short")
    
    # All Lights Control
    elif "all lights" in cmd or "all light" in cmd:
        if "on" in cmd:
            led_living.on()
            led_bedroom.on()
            led_kitchen.on()
            devices["living_room"] = True
            devices["bedroom"] = True
            devices["kitchen"] = True
            speak("All lights on", "long")
        elif "off" in cmd:
            led_living.off()
            led_bedroom.off()
            led_kitchen.off()
            devices["living_room"] = False
            devices["bedroom"] = False
            devices["kitchen"] = False
            speak("All lights off", "long")
    
    # Temperature Query
    elif "temperature" in cmd or "hot" in cmd or "cold" in cmd:
        update_temperature()
        speak(f"Temp {temperature}C", "short")
    
    # Humidity Query
    elif "humidity" in cmd:
        update_temperature()
        speak(f"Humidity {humidity}%", "short")
    
    # Weather Query
    elif "weather" in cmd:
        update_temperature()
        speak(f"{temperature}C {humidity}%", "short")
    
    # Time Query
    elif "time" in cmd:
        t = time.localtime()
        speak(f"{t[3]}:{t[4]:02d}", "short")
    
    # Good Morning Routine
    elif "good morning" in cmd or "morning" in cmd:
        led_living.on()
        led_bedroom.on()
        led_kitchen.on()
        devices["living_room"] = True
        devices["bedroom"] = True
        devices["kitchen"] = True
        update_temperature()
        speak("Good morning!", "long")
        time.sleep(1)
        speak(f"{temperature}C", "short")
    
    # Goodnight Routine
    elif "goodnight" in cmd or "good night" in cmd:
        led_living.off()
        led_bedroom.off()
        led_kitchen.off()
        devices["living_room"] = False
        devices["bedroom"] = False
        devices["kitchen"] = False
        speak("Goodnight!", "long")
    
    # Status Check
    elif "status" in cmd:
        count = sum(devices.values())
        speak(f"{count} devices on", "short")
    
    # Help
    elif "help" in cmd:
        speak("Say turn on or off device", "long")
    
    else:
        speak("Not understood", "alert")

# Boot Screen
def show_boot_screen():
    oled.fill(0)
    oled.rect(0, 0, 128, 64, 1)
    oled.text("AI ASSISTANT", 20, 10)
    oled.text("============", 15, 20)
    oled.text("Voice Control", 18, 32)
    oled.text("Initializing..", 15, 45)
    oled.show()
    time.sleep(2)

# Main Program
print("=" * 40)
print("🎤 AI VOICE ASSISTANT STARTING...")
print("=" * 40)

show_boot_screen()
update_temperature()

print("\n📋 AVAILABLE COMMANDS:")
print("  • 'Turn on/off [living room/bedroom/kitchen]'")
print("  • 'Turn on/off all lights'")
print("  • 'What's the temperature?'")
print("  • 'What's the humidity?'")
print("  • 'What time is it?'")
print("  • 'Good morning' / 'Goodnight'")
print("  • 'Status' / 'Help'")
print("\n💡 Type commands in Serial Monitor to simulate voice!")
print("   Or press the wake button.\n")

show_idle_screen()

try:
    while True:
        # Check wake button
        if wake_button.value() == 0:  # Button pressed (active low)
            listening = True
            show_listening_screen()
            speak("Yes?", "short")
            time.sleep(0.5)
            
            # Wait for button release
            while wake_button.value() == 0:
                time.sleep(0.1)
        
        # Simulate voice input via serial (for Wokwi testing)
        # In real implementation, this would be actual speech recognition
        user_input = input("🎤 You: ").strip() if listening else ""
        
        if user_input:
            # Check for wake word
            if "hey" in user_input.lower() or "assistant" in user_input.lower():
                listening = True
                show_listening_screen()
                speak("Yes?", "short")
            elif listening:
                # Process command
                process_command(user_input)
                listening = False
                show_idle_screen()
        else:
            # Update display periodically
            time.sleep(0.1)

except KeyboardInterrupt:
    print("\n\n🛑 Shutting down AI Assistant...")
    led_living.off()
    led_bedroom.off()
    led_kitchen.off()
    buzzer.off()
    oled.fill(0)
    oled.text("GOODBYE!", 35, 28)
    oled.show()