from machine import Pin, I2C, PWM, ADC
import time
import framebuf
import random

# 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 pixel(self, x, y, color):
        self.framebuf.pixel(x, y, 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)
    
    def line(self, x1, y1, x2, y2, color):
        self.framebuf.line(x1, y1, x2, y2, color)

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

# RGB LED (PWM for smooth colors)
led_r = PWM(Pin(6))
led_g = PWM(Pin(7))
led_b = PWM(Pin(8))
led_r.freq(1000)
led_g.freq(1000)
led_b.freq(1000)

# Rotary Encoder
encoder_clk = Pin(10, Pin.IN, Pin.PULL_UP)
encoder_dt = Pin(11, Pin.IN, Pin.PULL_UP)
encoder_sw = Pin(12, Pin.IN, Pin.PULL_UP)

# Control Buttons (Simulating IR Remote)
btn_up = Pin(15, Pin.IN, Pin.PULL_UP)
btn_down = Pin(16, Pin.IN, Pin.PULL_UP)
btn_select = Pin(17, Pin.IN, Pin.PULL_UP)
btn_back = Pin(18, Pin.IN, Pin.PULL_UP)
btn_play = Pin(19, Pin.IN, Pin.PULL_UP)

# Buzzer
buzzer = Pin(13, Pin.OUT)

# System Variables
volume = 50
last_encoder_clk = 1
button_debounce = {"up": 0, "down": 0, "select": 0, "back": 0, "play": 0}
DEBOUNCE_TIME = 200

# Media Library
MEDIA_LIBRARY = {
    "Movies": {
        "Action": ["The Matrix", "Die Hard", "Mad Max Fury Road", "John Wick"],
        "Comedy": ["Superbad", "The Hangover", "Anchorman", "Step Brothers"],
        "Drama": ["Shawshank Redemption", "Forrest Gump", "The Godfather"],
        "Sci-Fi": ["Inception", "Interstellar", "Blade Runner 2049"]
    },
    "TV Shows": {
        "Sitcoms": ["Friends S01E01", "The Office S01E01", "Parks & Rec S01E01"],
        "Drama": ["Breaking Bad S01E01", "Game of Thrones S01E01", "The Crown S01E01"],
        "Sci-Fi": ["Stranger Things S01E01", "Black Mirror S01E01", "Westworld S01E01"]
    },
    "Music": {
        "Rock": ["Queen - Bohemian Rhapsody", "AC/DC - Thunderstruck", "Led Zeppelin - Stairway"],
        "Pop": ["Michael Jackson - Thriller", "Madonna - Like a Prayer", "Prince - Purple Rain"],
        "Classical": ["Beethoven - Symphony No.9", "Mozart - Requiem", "Bach - Air on G String"],
        "Jazz": ["Miles Davis - So What", "John Coltrane - Giant Steps", "Ella Fitzgerald - Summertime"]
    },
    "Photos": {
        "Albums": ["Vacation 2024", "Family Events", "Nature Photography", "Cityscapes"]
    }
}

# Playback State
current_state = "main_menu"  # main_menu, category_menu, item_list, playing
current_category = None
current_subcategory = None
current_item = None
playback_status = "stopped"  # stopped, playing, paused
playback_progress = 0
selected_index = 0

# RGB Ambient Lighting Presets
LIGHTING_MODES = {
    "Movies": (255, 0, 100),      # Purple/Pink
    "TV Shows": (0, 150, 255),    # Blue
    "Music": (255, 100, 0),       # Orange
    "Photos": (0, 255, 100),      # Green
    "Playing": (100, 0, 255),     # Deep Purple
}

# Helper Functions
def set_rgb(r, g, b):
    """Set RGB LED color (0-255 values)"""
    led_r.duty_u16(int((255 - r) / 255 * 65535))
    led_g.duty_u16(int((255 - g) / 255 * 65535))
    led_b.duty_u16(int((255 - b) / 255 * 65535))

def beep(pattern="single"):
    """Play notification sound"""
    if pattern == "single":
        buzzer.on()
        time.sleep(0.05)
        buzzer.off()
    elif pattern == "double":
        for _ in range(2):
            buzzer.on()
            time.sleep(0.04)
            buzzer.off()
            time.sleep(0.04)
    elif pattern == "select":
        buzzer.on()
        time.sleep(0.08)
        buzzer.off()

def button_pressed(btn_name, btn_pin):
    """Check if button pressed with debouncing"""
    current = time.ticks_ms()
    if btn_pin.value() == 0 and time.ticks_diff(current, button_debounce[btn_name]) > DEBOUNCE_TIME:
        button_debounce[btn_name] = current
        return True
    return False

def read_encoder():
    """Read rotary encoder for volume control"""
    global volume, last_encoder_clk
    
    clk_state = encoder_clk.value()
    if clk_state != last_encoder_clk:
        if encoder_dt.value() != clk_state:
            volume = min(100, volume + 2)
        else:
            volume = max(0, volume - 2)
        last_encoder_clk = clk_state
        beep("single")

# Display Functions
def draw_header(title):
    """Draw screen header"""
    oled.fill_rect(0, 0, 128, 12, 1)
    oled.text(title[:16], 2, 2, 0)
    oled.line(0, 12, 128, 12, 1)

def draw_volume_bar():
    """Draw volume indicator"""
    bar_width = int(volume * 1.2)
    oled.rect(2, 54, 124, 8, 1)
    oled.fill_rect(3, 55, bar_width, 6, 1)
    oled.text(f"{volume}%", 96, 55, 0 if volume > 80 else 1)

def draw_progress_bar(progress):
    """Draw playback progress bar"""
    bar_width = int(progress * 1.2)
    oled.rect(2, 45, 124, 6, 1)
    oled.fill_rect(3, 46, bar_width, 4, 1)

def draw_menu(items, selected, start_y=16):
    """Draw scrollable menu list"""
    visible_items = 4
    scroll_offset = max(0, selected - visible_items + 1)
    
    for i in range(visible_items):
        idx = scroll_offset + i
        if idx >= len(items):
            break
        
        y = start_y + i * 11
        item_text = items[idx][:15]
        
        if idx == selected:
            oled.fill_rect(0, y, 128, 10, 1)
            oled.text(f">{item_text}", 4, y + 1, 0)
        else:
            oled.text(f" {item_text}", 4, y + 1, 1)

def show_main_menu():
    """Display main menu"""
    oled.fill(0)
    draw_header("MEDIA CENTER")
    
    categories = list(MEDIA_LIBRARY.keys())
    draw_menu(categories, selected_index)
    
    # Show volume at bottom
    oled.text("VOL:", 2, 54)
    oled.text(f"{volume}%", 30, 54)
    
    # Show time
    oled.text("12:45", 92, 54)
    
    oled.show()

def show_category_menu():
    """Display category/genre menu"""
    oled.fill(0)
    draw_header(current_category[:14])
    
    subcategories = list(MEDIA_LIBRARY[current_category].keys())
    draw_menu(subcategories, selected_index)
    
    oled.text("B:Back", 2, 54)
    oled.show()

def show_item_list():
    """Display media items"""
    oled.fill(0)
    draw_header(current_subcategory[:14])
    
    items = MEDIA_LIBRARY[current_category][current_subcategory]
    draw_menu(items, selected_index)
    
    oled.text("B:Back SEL:Play", 2, 54)
    oled.show()

def show_now_playing():
    """Display now playing screen"""
    global playback_progress
    
    oled.fill(0)
    draw_header("NOW PLAYING")
    
    # Show media info
    if current_item:
        # Title (wrap if needed)
        title = current_item[:21]
        oled.text(title, 4, 18)
        
        # Category
        oled.text(current_subcategory[:16], 4, 28, 1)
    
    # Playback status
    status_icon = ">" if playback_status == "playing" else "||" if playback_status == "paused" else "[]"
    oled.text(status_icon, 4, 38)
    
    # Progress bar
    draw_progress_bar(playback_progress)
    
    # Time
    current_time = int(playback_progress / 100 * 180)  # Max 3 min
    total_time = 180
    oled.text(f"{current_time//60}:{current_time%60:02d}/{total_time//60}:{total_time%60:02d}", 30, 38)
    
    # Controls
    oled.text("PLAY BACK STOP", 2, 54)
    
    oled.show()

def show_splash_screen():
    """Boot animation"""
    oled.fill(0)
    oled.rect(10, 10, 108, 44, 1)
    oled.text("MEDIA CENTER", 22, 20)
    oled.text("===========", 22, 30)
    oled.text("Loading...", 30, 40)
    oled.show()
    
    # RGB rainbow effect
    for i in range(0, 255, 15):
        set_rgb(i, 255-i, 128)
        time.sleep(0.05)
    
    set_rgb(0, 0, 0)

# State Management
def handle_main_menu():
    """Main menu navigation"""
    global current_state, current_category, selected_index
    
    categories = list(MEDIA_LIBRARY.keys())
    
    if button_pressed("up", btn_up):
        selected_index = (selected_index - 1) % len(categories)
        beep("single")
    
    if button_pressed("down", btn_down):
        selected_index = (selected_index + 1) % len(categories)
        beep("single")
    
    if button_pressed("select", btn_select):
        current_category = categories[selected_index]
        current_state = "category_menu"
        selected_index = 0
        beep("select")
        
        # Set ambient lighting
        if current_category in LIGHTING_MODES:
            r, g, b = LIGHTING_MODES[current_category]
            set_rgb(r, g, b)
    
    show_main_menu()

def handle_category_menu():
    """Category/genre menu navigation"""
    global current_state, current_subcategory, selected_index
    
    subcategories = list(MEDIA_LIBRARY[current_category].keys())
    
    if button_pressed("up", btn_up):
        selected_index = (selected_index - 1) % len(subcategories)
        beep("single")
    
    if button_pressed("down", btn_down):
        selected_index = (selected_index + 1) % len(subcategories)
        beep("single")
    
    if button_pressed("select", btn_select):
        current_subcategory = subcategories[selected_index]
        current_state = "item_list"
        selected_index = 0
        beep("select")
    
    if button_pressed("back", btn_back):
        current_state = "main_menu"
        selected_index = 0
        beep("single")
        set_rgb(0, 0, 0)
    
    show_category_menu()

def handle_item_list():
    """Media item list navigation"""
    global current_state, current_item, selected_index, playback_status, playback_progress
    
    items = MEDIA_LIBRARY[current_category][current_subcategory]
    
    if button_pressed("up", btn_up):
        selected_index = (selected_index - 1) % len(items)
        beep("single")
    
    if button_pressed("down", btn_down):
        selected_index = (selected_index + 1) % len(items)
        beep("single")
    
    if button_pressed("select", btn_select):
        current_item = items[selected_index]
        current_state = "playing"
        playback_status = "playing"
        playback_progress = 0
        beep("double")
        
        # Set playing ambient light
        r, g, b = LIGHTING_MODES["Playing"]
        set_rgb(r, g, b)
    
    if button_pressed("back", btn_back):
        current_state = "category_menu"
        selected_index = 0
        beep("single")
    
    show_item_list()

def handle_playing():
    """Playback screen"""
    global current_state, playback_status, playback_progress
    
    if button_pressed("play", btn_play):
        if playback_status == "playing":
            playback_status = "paused"
            beep("single")
        elif playback_status == "paused":
            playback_status = "playing"
            beep("single")
    
    if button_pressed("back", btn_back):
        playback_status = "stopped"
        current_state = "item_list"
        playback_progress = 0
        beep("single")
        
        # Restore category lighting
        if current_category in LIGHTING_MODES:
            r, g, b = LIGHTING_MODES[current_category]
            set_rgb(r, g, b)
    
    # Update progress
    if playback_status == "playing":
        playback_progress = min(100, playback_progress + 0.5)
        
        # Auto-return when finished
        if playback_progress >= 100:
            playback_status = "stopped"
            current_state = "item_list"
            playback_progress = 0
            beep("double")
    
    show_now_playing()

# Main Program
print("=" * 50)
print("📺 SMART HOME ENTERTAINMENT SYSTEM")
print("=" * 50)
print("\n🎬 Media Library Loaded:")
for category, subcats in MEDIA_LIBRARY.items():
    total_items = sum(len(items) for items in subcats.values())
    print(f"  • {category}: {total_items} items")

print("\n🎮 Controls:")
print("  UP/DOWN: Navigate")
print("  SELECT: Choose/Play")
print("  BACK: Return")
print("  PLAY: Play/Pause")
print("  Encoder: Volume Control")

print("\n🚀 Starting Media Center...\n")

show_splash_screen()
time.sleep(2)

try:
    while True:
        # Read encoder for volume
        read_encoder()
        
        # State machine
        if current_state == "main_menu":
            handle_main_menu()
        elif current_state == "category_menu":
            handle_category_menu()
        elif current_state == "item_list":
            handle_item_list()
        elif current_state == "playing":
            handle_playing()
        
        time.sleep(0.05)  # 20Hz update rate

except KeyboardInterrupt:
    print("\n\n🛑 Shutting down Media Center...")
    set_rgb(0, 0, 0)
    buzzer.off()
    oled.fill(0)
    oled.text("GOODBYE!", 40, 28)
    oled.show()
    time.sleep(1)
    oled.fill(0)
    oled.show()