from machine import Pin, I2C
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)

# Controller Buttons
btn_up = Pin(10, Pin.IN, Pin.PULL_UP)
btn_down = Pin(11, Pin.IN, Pin.PULL_UP)
btn_left = Pin(12, Pin.IN, Pin.PULL_UP)
btn_right = Pin(13, Pin.IN, Pin.PULL_UP)
btn_a = Pin(14, Pin.IN, Pin.PULL_UP)
btn_b = Pin(15, Pin.IN, Pin.PULL_UP)

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

# Sound Effects
def beep(duration=0.05, pattern="single"):
    if pattern == "single":
        buzzer.on()
        time.sleep(duration)
        buzzer.off()
    elif pattern == "double":
        for _ in range(2):
            buzzer.on()
            time.sleep(0.05)
            buzzer.off()
            time.sleep(0.05)
    elif pattern == "game_over":
        for freq in [0.1, 0.1, 0.15]:
            buzzer.on()
            time.sleep(freq)
            buzzer.off()
            time.sleep(0.05)

# Button Debouncing
last_press = {"up": 0, "down": 0, "left": 0, "right": 0, "a": 0, "b": 0}
debounce_time = 150  # ms

def button_pressed(button_name):
    current = time.ticks_ms()
    if time.ticks_diff(current, last_press[button_name]) > debounce_time:
        last_press[button_name] = current
        return True
    return False

def read_buttons():
    buttons = {
        "up": btn_up.value() == 0,
        "down": btn_down.value() == 0,
        "left": btn_left.value() == 0,
        "right": btn_right.value() == 0,
        "a": btn_a.value() == 0,
        "b": btn_b.value() == 0
    }
    return buttons

# ===== SNAKE GAME =====
class SnakeGame:
    def __init__(self):
        self.reset()
    
    def reset(self):
        self.snake = [[64, 32], [60, 32], [56, 32]]
        self.direction = [4, 0]
        self.food = [random.randint(0, 30) * 4, random.randint(0, 14) * 4]
        self.score = 0
        self.game_over = False
    
    def update(self, buttons):
        if self.game_over:
            return
        
        # Change direction
        if buttons["up"] and self.direction[1] == 0:
            self.direction = [0, -4]
        elif buttons["down"] and self.direction[1] == 0:
            self.direction = [0, 4]
        elif buttons["left"] and self.direction[0] == 0:
            self.direction = [-4, 0]
        elif buttons["right"] and self.direction[0] == 0:
            self.direction = [4, 0]
        
        # Move snake
        head = [self.snake[0][0] + self.direction[0], self.snake[0][1] + self.direction[1]]
        
        # Check wall collision
        if head[0] < 0 or head[0] >= 128 or head[1] < 0 or head[1] >= 64:
            self.game_over = True
            beep(0.1, "game_over")
            return
        
        # Check self collision
        if head in self.snake:
            self.game_over = True
            beep(0.1, "game_over")
            return
        
        self.snake.insert(0, head)
        
        # Check food collision
        if head == self.food:
            self.score += 10
            beep(0.05, "double")
            self.food = [random.randint(0, 30) * 4, random.randint(0, 14) * 4]
        else:
            self.snake.pop()
    
    def draw(self):
        oled.fill(0)
        
        # Draw snake
        for segment in self.snake:
            oled.fill_rect(segment[0], segment[1], 4, 4, 1)
        
        # Draw food
        oled.fill_rect(self.food[0], self.food[1], 4, 4, 1)
        
        # Draw score
        oled.text(f"Score:{self.score}", 0, 0)
        
        if self.game_over:
            oled.fill_rect(20, 25, 88, 20, 0)
            oled.rect(20, 25, 88, 20, 1)
            oled.text("GAME OVER", 30, 30)
        
        oled.show()

# ===== PONG GAME =====
class PongGame:
    def __init__(self):
        self.reset()
    
    def reset(self):
        self.paddle_y = 24
        self.ball_x = 64
        self.ball_y = 32
        self.ball_dx = 2
        self.ball_dy = 1
        self.score = 0
        self.game_over = False
    
    def update(self, buttons):
        if self.game_over:
            return
        
        # Move paddle
        if buttons["up"] and self.paddle_y > 0:
            self.paddle_y -= 3
        if buttons["down"] and self.paddle_y < 48:
            self.paddle_y += 3
        
        # Move ball
        self.ball_x += self.ball_dx
        self.ball_y += self.ball_dy
        
        # Ball collision with top/bottom
        if self.ball_y <= 0 or self.ball_y >= 62:
            self.ball_dy *= -1
            beep(0.03)
        
        # Ball collision with paddle
        if self.ball_x <= 8 and self.paddle_y <= self.ball_y <= self.paddle_y + 16:
            self.ball_dx *= -1
            self.score += 1
            beep(0.05)
        
        # Ball out of bounds
        if self.ball_x < 0:
            self.game_over = True
            beep(0.1, "game_over")
        
        # Ball collision with right wall
        if self.ball_x >= 126:
            self.ball_dx *= -1
    
    def draw(self):
        oled.fill(0)
        
        # Draw paddle
        oled.fill_rect(2, self.paddle_y, 4, 16, 1)
        
        # Draw ball
        oled.fill_rect(self.ball_x, self.ball_y, 2, 2, 1)
        
        # Draw score
        oled.text(f"Score:{self.score}", 40, 0)
        
        if self.game_over:
            oled.fill_rect(20, 25, 88, 20, 0)
            oled.rect(20, 25, 88, 20, 1)
            oled.text("GAME OVER", 30, 30)
        
        oled.show()

# ===== SPACE SHOOTER GAME =====
class SpaceShooter:
    def __init__(self):
        self.reset()
    
    def reset(self):
        self.player_x = 60
        self.bullets = []
        self.enemies = [[random.randint(0, 120), random.randint(-30, -10)] for _ in range(3)]
        self.score = 0
        self.lives = 3
        self.game_over = False
    
    def update(self, buttons):
        if self.game_over:
            return
        
        # Move player
        if buttons["left"] and self.player_x > 0:
            self.player_x -= 3
        if buttons["right"] and self.player_x < 120:
            self.player_x += 3
        
        # Shoot
        if buttons["a"] and button_pressed("a"):
            self.bullets.append([self.player_x + 3, 50])
            beep(0.03)
        
        # Move bullets
        for bullet in self.bullets[:]:
            bullet[1] -= 4
            if bullet[1] < 0:
                self.bullets.remove(bullet)
        
        # Move enemies
        for enemy in self.enemies:
            enemy[1] += 1
            if enemy[1] > 64:
                enemy[0] = random.randint(0, 120)
                enemy[1] = random.randint(-30, -10)
                self.lives -= 1
                beep(0.1)
                if self.lives <= 0:
                    self.game_over = True
                    beep(0.1, "game_over")
        
        # Check collisions
        for bullet in self.bullets[:]:
            for enemy in self.enemies[:]:
                if abs(bullet[0] - enemy[0]) < 6 and abs(bullet[1] - enemy[1]) < 6:
                    self.bullets.remove(bullet)
                    enemy[0] = random.randint(0, 120)
                    enemy[1] = random.randint(-30, -10)
                    self.score += 10
                    beep(0.05, "double")
                    break
    
    def draw(self):
        oled.fill(0)
        
        # Draw player
        oled.fill_rect(self.player_x, 54, 8, 6, 1)
        oled.pixel(self.player_x + 4, 52, 1)
        
        # Draw bullets
        for bullet in self.bullets:
            oled.fill_rect(bullet[0], bullet[1], 2, 4, 1)
        
        # Draw enemies
        for enemy in self.enemies:
            oled.fill_rect(enemy[0], enemy[1], 6, 6, 1)
        
        # Draw HUD
        oled.text(f"S:{self.score}", 0, 0)
        oled.text(f"L:{self.lives}", 100, 0)
        
        if self.game_over:
            oled.fill_rect(20, 25, 88, 20, 0)
            oled.rect(20, 25, 88, 20, 1)
            oled.text("GAME OVER", 30, 30)
        
        oled.show()

# ===== BREAKOUT GAME =====
class BreakoutGame:
    def __init__(self):
        self.reset()
    
    def reset(self):
        self.paddle_x = 52
        self.ball_x = 64
        self.ball_y = 50
        self.ball_dx = 2
        self.ball_dy = -2
        self.bricks = [[x * 16, y * 6 + 10] for x in range(8) for y in range(4)]
        self.score = 0
        self.game_over = False
        self.won = False
    
    def update(self, buttons):
        if self.game_over or self.won:
            return
        
        # Move paddle
        if buttons["left"] and self.paddle_x > 0:
            self.paddle_x -= 4
        if buttons["right"] and self.paddle_x < 104:
            self.paddle_x += 4
        
        # Move ball
        self.ball_x += self.ball_dx
        self.ball_y += self.ball_dy
        
        # Ball collision with walls
        if self.ball_x <= 0 or self.ball_x >= 126:
            self.ball_dx *= -1
            beep(0.03)
        if self.ball_y <= 0:
            self.ball_dy *= -1
            beep(0.03)
        
        # Ball collision with paddle
        if (self.paddle_x <= self.ball_x <= self.paddle_x + 24 and 
            58 <= self.ball_y <= 60):
            self.ball_dy *= -1
            beep(0.05)
        
        # Ball out of bounds
        if self.ball_y > 64:
            self.game_over = True
            beep(0.1, "game_over")
        
        # Ball collision with bricks
        for brick in self.bricks[:]:
            if (brick[0] <= self.ball_x <= brick[0] + 14 and 
                brick[1] <= self.ball_y <= brick[1] + 4):
                self.bricks.remove(brick)
                self.ball_dy *= -1
                self.score += 5
                beep(0.04)
                if not self.bricks:
                    self.won = True
                    beep(0.1, "double")
                break
    
    def draw(self):
        oled.fill(0)
        
        # Draw paddle
        oled.fill_rect(self.paddle_x, 58, 24, 4, 1)
        
        # Draw ball
        oled.fill_rect(self.ball_x, self.ball_y, 2, 2, 1)
        
        # Draw bricks
        for brick in self.bricks:
            oled.rect(brick[0], brick[1], 14, 4, 1)
        
        # Draw score
        oled.text(f"Score:{self.score}", 0, 0)
        
        if self.game_over:
            oled.fill_rect(20, 25, 88, 20, 0)
            oled.rect(20, 25, 88, 20, 1)
            oled.text("GAME OVER", 30, 30)
        elif self.won:
            oled.fill_rect(20, 25, 88, 20, 0)
            oled.rect(20, 25, 88, 20, 1)
            oled.text("YOU WIN!", 35, 30)
        
        oled.show()

# ===== GAME MENU =====
class GameMenu:
    def __init__(self):
        self.games = ["Snake", "Pong", "Shooter", "Breakout"]
        self.selected = 0
    
    def update(self, buttons):
        if buttons["up"] and button_pressed("up"):
            self.selected = (self.selected - 1) % len(self.games)
            beep(0.03)
        if buttons["down"] and button_pressed("down"):
            self.selected = (self.selected + 1) % len(self.games)
            beep(0.03)
        if buttons["a"] and button_pressed("a"):
            beep(0.05, "double")
            return self.selected
        return None
    
    def draw(self):
        oled.fill(0)
        oled.rect(0, 0, 128, 64, 1)
        oled.text("RETRO ARCADE", 20, 5)
        oled.line(0, 15, 128, 15, 1)
        
        for i, game in enumerate(self.games):
            y = 22 + i * 10
            if i == self.selected:
                oled.text(f">{game}", 15, y)
            else:
                oled.text(f" {game}", 15, y)
        
        oled.show()

# ===== MAIN GAME LOOP =====
print("=" * 40)
print("🎮 RETRO GAMING CONSOLE")
print("=" * 40)
print("\nControls:")
print("  D-Pad: UP/DOWN/LEFT/RIGHT")
print("  A Button: Action/Select")
print("  B Button: Back/Menu")
print("\nStarting...\n")

# Boot animation
oled.fill(0)
oled.text("RETRO ARCADE", 25, 20)
oled.text("Loading...", 30, 35)
oled.show()
time.sleep(1)

menu = GameMenu()
current_game = None
game_mode = "menu"

try:
    while True:
        buttons = read_buttons()
        
        if game_mode == "menu":
            selected = menu.update(buttons)
            menu.draw()
            
            if selected is not None:
                if selected == 0:
                    current_game = SnakeGame()
                elif selected == 1:
                    current_game = PongGame()
                elif selected == 2:
                    current_game = SpaceShooter()
                elif selected == 3:
                    current_game = BreakoutGame()
                game_mode = "playing"
                print(f"🎮 Starting {menu.games[selected]}...")
        
        elif game_mode == "playing":
            current_game.update(buttons)
            current_game.draw()
            
            # Back to menu
            if buttons["b"] and button_pressed("b"):
                game_mode = "menu"
                beep(0.05)
                print("📋 Returning to menu...")
        
        time.sleep(0.05)  # 20 FPS

except KeyboardInterrupt:
    print("\n\n🛑 Shutting down console...")
    oled.fill(0)
    oled.text("GAME OVER", 35, 28)
    oled.show()
    buzzer.off()