from machine import Pin
import neopixel
import network
import time
from umqtt.simple import MQTTClient
import json

# WiFi configuration
WIFI_SSID = "Wokwi-GUEST"
WIFI_PASSWORD = ""

# MQTT configuration
MQTT_BROKER = "broker.hivemq.com"
MQTT_PORT = 1883
MQTT_CLIENT_ID = "esp32_led_ring"
MQTT_TOPIC = b"home/ledring"

# NeoPixel setup
PIN = 13
NUM_PIXELS = 16  # Number of LEDs in the ring
np = neopixel.NeoPixel(Pin(PIN, Pin.OUT), NUM_PIXELS)

# Color presets
COLORS = {
    'morning': (255, 200, 100),
    'day': (255, 255, 255),
    'evening': (255, 147, 41),
    'night': (50, 0, 255)
}

# Global state
auto_mode = True  # Toggle between auto and manual mode
current_brightness = 100  # Brightness percentage (0-100)
current_color = (0, 0, 0)  # Store current color for brightness adjustments

def connect_wifi():
    """Connect to WiFi network"""
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    if not wlan.isconnected():
        print('Connecting to WiFi...')
        wlan.connect(WIFI_SSID, WIFI_PASSWORD)
        while not wlan.isconnected():
            time.sleep(1)
    print('WiFi connected!')
    print('Network config:', wlan.ifconfig())

def apply_brightness(r, g, b, brightness=100):
    """Apply brightness to a color"""
    brightness_factor = brightness / 100.0
    return (
        int(r * brightness_factor),
        int(g * brightness_factor),
        int(b * brightness_factor)
    )

def set_color(r, g, b, brightness=None):
    """Set all pixels to the same color with optional brightness"""
    global current_color, current_brightness
    
    # If no specific brightness provided, use global brightness
    if brightness is None:
        brightness = current_brightness
    
    # Store the original color and current brightness
    current_color = (r, g, b)
    current_brightness = brightness
    
    # Apply brightness
    adjusted_r, adjusted_g, adjusted_b = apply_brightness(r, g, b, brightness)
    
    for i in range(NUM_PIXELS):
        np[i] = (adjusted_r, adjusted_g, adjusted_b)
    np.write()

def fade_to_color(target_r, target_g, target_b, steps=10, brightness=None):
    """Smooth transition to target color with optional brightness"""
    global current_brightness
    
    # If no specific brightness provided, use global brightness
    if brightness is None:
        brightness = current_brightness
    
    current = np[0]
    r_step = (target_r - current[0]) / steps
    g_step = (target_g - current[1]) / steps
    b_step = (target_b - current[2]) / steps

    for i in range(steps):
        r = int(current[0] + (r_step * i))
        g = int(current[1] + (g_step * i))
        b = int(current[2] + (b_step * i))
        
        # Apply brightness during fade
        adjusted_r, adjusted_g, adjusted_b = apply_brightness(r, g, b, brightness)
        
        for j in range(NUM_PIXELS):
            np[j] = (adjusted_r, adjusted_g, adjusted_b)
        np.write()
        time.sleep_ms(20)  # Reduced delay
    
    # Ensure final color is set
    set_color(target_r, target_g, target_b, brightness)

def rainbow_effect(wait=5, brightness=None):
    """Create a faster rainbow effect with optional brightness"""
    # If no specific brightness provided, use global brightness
    if brightness is None:
        brightness = current_brightness
    
    def wheel(pos):
        if pos < 85:
            return (pos * 3, 255 - pos * 3, 0)
        elif pos < 170:
            pos -= 85
            return (255 - pos * 3, 0, pos * 3)
        else:
            pos -= 170
            return (0, pos * 3, 255 - pos * 3)

    for j in range(255):
        for i in range(NUM_PIXELS):
            rc_index = (i * 256 // NUM_PIXELS) + j
            base_color = wheel(rc_index & 255)
            
            # Apply brightness
            adjusted_r, adjusted_g, adjusted_b = apply_brightness(
                base_color[0], base_color[1], base_color[2], brightness
            )
            np[i] = (adjusted_r, adjusted_g, adjusted_b)
        np.write()
        time.sleep_ms(wait)  # Reduced wait time

def breathing_effect(r, g, b, cycles=2, step_delay=10, brightness=None):
    """Create a faster breathing effect with specified color and optional brightness"""
    # If no specific brightness provided, use global brightness
    if brightness is None:
        brightness = current_brightness
    
    for _ in range(cycles):
        # Fade in
        for i in range(50):  # Reduced steps from 100 to 50
            breath_brightness = (i / 50.0) * brightness
            adjusted_r, adjusted_g, adjusted_b = apply_brightness(r, g, b, breath_brightness)
            set_color(adjusted_r, adjusted_g, adjusted_b)
            time.sleep_ms(step_delay)  # Reduced delay
        
        # Fade out
        for i in range(50, -1, -1):  
            breath_brightness = (i / 50.0) * brightness
            adjusted_r, adjusted_g, adjusted_b = apply_brightness(r, g, b, breath_brightness)
            set_color(adjusted_r, adjusted_g, adjusted_b)
            time.sleep_ms(step_delay)

def get_demo_hour():
    """Simulated time function for demo"""
    return (time.time() // 10) % 24

def get_time_mode(hour):
    """Determine color mode based on hour"""
    if 5 <= hour < 9:
        return 'morning'
    elif 9 <= hour < 17:
        return 'day'
    elif 17 <= hour < 21:
        return 'evening'
    else:
        return 'night'

def mqtt_callback(topic, msg):
    """Handle incoming MQTT messages"""
    global auto_mode, current_brightness
    try:
        data = json.loads(msg)
        print("Received MQTT message:", data)
        
        # Mode control
        if 'mode' in data:
            auto_mode = data['mode'] == 'auto'
            print(f"Mode set to: {'auto' if auto_mode else 'manual'}")
            return
        
        # Brightness control
        if 'brightness' in data:
            brightness = max(0, min(100, data['brightness']))  # Clamp between 0-100
            current_brightness = brightness
            print(f"Brightness set to: {brightness}%")
            
            # Reapply current color with new brightness
            if current_color != (0, 0, 0):
                set_color(current_color[0], current_color[1], current_color[2], brightness)
            return
        
        if not auto_mode:  
            # Color control
            if 'color' in data:
                color = data['color']
                r = color.get('r', 0)
                g = color.get('g', 0)
                b = color.get('b', 0)
                # Use optional brightness if provided
                brightness = color.get('brightness', current_brightness)
                fade_to_color(r, g, b, brightness=brightness)
                print(f'Set color to R:{r} G:{g} B:{b} Brightness:{brightness}%')
            
            # Preset color control
            elif 'preset' in data:
                if data['preset'] in COLORS:
                    r, g, b = COLORS[data['preset']]
                    # Use optional brightness if provided
                    brightness = data.get('brightness', current_brightness)
                    fade_to_color(r, g, b, brightness=brightness)
                    print(f'Applied preset: {data["preset"]} with {brightness}% brightness')
            
            # Effect control
            elif 'effect' in data:
                effect = data['effect']
                # Use optional brightness if provided
                brightness = data.get('brightness', current_brightness)
                
                if effect == 'rainbow':
                    rainbow_effect(brightness=brightness)
                elif effect == 'breathing':
                    color = data.get('color', {'r': 255, 'g': 0, 'b': 0})
                    breathing_effect(
                        color['r'], 
                        color['g'], 
                        color['b'], 
                        brightness=brightness
                    )
                    
    except Exception as e:
        print('Error processing message:', e)

def connect_mqtt():
    """Connect to MQTT broker"""
    client = MQTTClient(MQTT_CLIENT_ID, MQTT_BROKER, port=MQTT_PORT)
    client.set_callback(mqtt_callback)
    client.connect()
    print('Connected to MQTT broker')
    client.subscribe(MQTT_TOPIC)
    return client

def main():
    try:
        # Connect to WiFi
        connect_wifi()
        
        # Connect to MQTT broker
        client = connect_mqtt()
        print(f'Subscribed to {MQTT_TOPIC}')
        
        # Initialize with rainbow effect
        rainbow_effect(wait=5)
        
        last_check = time.time()
        
        # Main loop
        while True:
            client.check_msg()  
            
            # Handle auto mode color changes efficiently
            current_time = time.time()
            if auto_mode and current_time - last_check >= 2:  
                hour = get_demo_hour()
                mode = get_time_mode(hour)
                print(f"Auto mode - Hour: {hour}, Mode: {mode}")
                r, g, b = COLORS[mode]
                fade_to_color(r, g, b)  
                last_check = current_time
            
            time.sleep(0.05)  
            
    except Exception as e:
        print('Error:', e)

# Start the program
if __name__ == '__main__':
    main()