import pygame
import sys
import random
import os

pygame.init()
pygame.mixer.init()
WIDTH, HEIGHT = 1200, 800
CELL_SIZE = 10
FPS = 45
SPEED_BOOST_DURATION = 200
POWERUP_COLOR = (255, 0, 255)
NUM_POWERUPS = 200
NUM_ENEMIES = 5

WORLD_WIDTH = 3000
WORLD_HEIGHT = 3000
ENEMY_RESPAWN_DELAY = 120
TAIL_SHRINK_ON_HIT = 5

screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Snake IO Style Arena")
clock = pygame.time.Clock()

WHITE = (255, 255, 255)
BG_COLOR = (30, 30, 30)
PLAYER1_COLOR = (0, 200, 255)
PLAYER1_TRAIL = (0, 100, 200)
PLAYER2_COLOR = (255, 200, 0)
PLAYER2_TRAIL = (200, 100, 0)
ENEMY_COLOR = (255, 80, 80)
ENEMY_TRAIL = (180, 50, 50)
TERRITORY_COLOR = (0, 255, 100)
UI_COLOR = (200, 200, 200)

font = pygame.font.SysFont("arial", 25)

try:
    capture_sound = pygame.mixer.Sound("capture.wav")
    die_sound = pygame.mixer.Sound("die.wav")
except pygame.error:
    capture_sound = None
    die_sound = None

LEADERBOARD_FILE = "leaderboard.txt"

shop_items = [
    {"name": "Speed Boost", "price": 50, "effect": lambda p: setattr(p, 'speed_boost', SPEED_BOOST_DURATION)},
    {"name": "Shield", "price": 75, "effect": lambda p: setattr(p, 'shield', True)}
]

def draw_ui(player):
    draw_text(f"Score: {player.score}", (10, 10))
    draw_text("[S]hop", (10, 40))

def draw_shop(player):
    screen.fill(BG_COLOR)
    draw_text("SHOP - Press Number to Buy, ESC to Exit", (WIDTH // 2, 100), center=True)
    for i, item in enumerate(shop_items):
        draw_text(f"{i+1}. {item['name']} - {item['price']} pts", (WIDTH // 2, 150 + i * 30), center=True)
    pygame.display.flip()
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    return
                for i in range(len(shop_items)):
                    if event.key == getattr(pygame, f'K_{i+1}'):
                        item = shop_items[i]
                        if player.score >= item["price"]:
                            player.score -= item["price"]
                            item["effect"](player)
                        return

class Player:
    def __init__(self, x, y, color, trail_color, controls):
        self.pos = [x, y]
        self.dir = [CELL_SIZE, 0]
        self.trail = []
        self.territory = set()
        self.color = color
        self.trail_color = trail_color
        self.alive = True
        self.score = 0
        self.controls = controls
        self.speed_boost = 0
        self.shield = False
        self.territory.add(tuple(self.pos))

    def handle_input(self):
        keys = pygame.key.get_pressed()
        if keys[self.controls['up']] and self.dir != [0, CELL_SIZE]:
            self.dir = [0, -CELL_SIZE]
        elif keys[self.controls['down']] and self.dir != [0, -CELL_SIZE]:
            self.dir = [0, CELL_SIZE]
        elif keys[self.controls['left']] and self.dir != [CELL_SIZE, 0]:
            self.dir = [-CELL_SIZE, 0]
        elif keys[self.controls['right']] and self.dir != [-CELL_SIZE, 0]:
            self.dir = [CELL_SIZE, 0]

    def move(self):
        if not self.alive:
            return
        step_factor = 2 if self.speed_boost > 0 else 1
        self.pos[0] = (self.pos[0] + self.dir[0] * step_factor) % WORLD_WIDTH
        self.pos[1] = (self.pos[1] + self.dir[1] * step_factor) % WORLD_HEIGHT
        pos_tuple = tuple(self.pos)
        if self.speed_boost > 0:
            self.speed_boost -= 1
        if pos_tuple in self.territory:
            self.fill_area()
            if capture_sound:
                capture_sound.play()
        else:
            if tuple(self.pos) not in [tuple(t) for t in self.trail]:
                self.trail.append(list(self.pos))

    def fill_area(self):
        self.territory.update(map(tuple, self.trail))
        self.score += len(self.trail)
        self.trail.clear()

    def die(self):
        if not self.shield:
            self.alive = False
            if die_sound:
                die_sound.play()
        else:
            self.shield = False

    def draw(self, offset):
        for t in self.trail:
            pygame.draw.rect(screen, self.trail_color, (t[0]-offset[0], t[1]-offset[1], CELL_SIZE, CELL_SIZE))
        for z in self.territory:
            pygame.draw.rect(screen, TERRITORY_COLOR, (z[0]-offset[0], z[1]-offset[1], CELL_SIZE, CELL_SIZE))
        pygame.draw.rect(screen, self.color, (self.pos[0]-offset[0], self.pos[1]-offset[1], CELL_SIZE, CELL_SIZE))

    def check_collision(self, others):
        for other in others:
            if self != other and tuple(self.pos) in other.trail:
                if len(self.trail) > TAIL_SHRINK_ON_HIT:
                    self.trail = self.trail[:-TAIL_SHRINK_ON_HIT]
                else:
                    self.die()
            if isinstance(other, Enemy) and tuple(other.pos) == tuple(self.pos):
                other.reset()

class Enemy(Player):
    def __init__(self, x, y):
        super().__init__(x, y, ENEMY_COLOR, ENEMY_TRAIL, {'up': None, 'down': None, 'left': None, 'right': None})
        self.respawn_timer = 0

    def move_ai(self, players):
        if not self.alive:
            if self.respawn_timer > 0:
                self.respawn_timer -= 1
                return
            else:
                self.reset()
                return
        self.avoid_players(players)
        super().move()

    def avoid_players(self, players):
        safest_dir = self.dir
        max_dist = -1
        for d in [[CELL_SIZE, 0], [-CELL_SIZE, 0], [0, CELL_SIZE], [0, -CELL_SIZE]]:
            future = [(self.pos[0] + d[0]) % WORLD_WIDTH, (self.pos[1] + d[1]) % WORLD_HEIGHT]
            dist = min([abs(future[0]-p.pos[0]) + abs(future[1]-p.pos[1]) for p in players if p.alive and p != self], default=1000)
            if dist > max_dist:
                max_dist = dist
                safest_dir = d
        self.dir = safest_dir

    def reset(self):
        self.alive = False
        self.respawn_timer = ENEMY_RESPAWN_DELAY
        self.score = 0
        self.pos = [random.randint(0, WORLD_WIDTH // CELL_SIZE) * CELL_SIZE, random.randint(0, WORLD_HEIGHT // CELL_SIZE) * CELL_SIZE]
        self.trail.clear()
        self.territory.clear()
        self.territory.add(tuple(self.pos))
        self.alive = True

def draw_text(text, pos, color=UI_COLOR, center=False):
    img = font.render(text, True, color)
    rect = img.get_rect()
    rect.center = pos if center else pos
    screen.blit(img, rect)

def save_score(score):
    with open(LEADERBOARD_FILE, 'a') as f:
        f.write(f"{score}\n")

def show_leaderboard():
    if not os.path.exists(LEADERBOARD_FILE):
        return []
    with open(LEADERBOARD_FILE, 'r') as f:
        scores = [int(s.strip()) for s in f.readlines() if s.strip().isdigit()]
    return sorted(scores, reverse=True)[:5]

def game_loop():
    player1 = Player(100, 100, PLAYER1_COLOR, PLAYER1_TRAIL, {'up':pygame.K_UP,'down':pygame.K_DOWN,'left':pygame.K_LEFT,'right':pygame.K_RIGHT})
    player2 = Player(700, 500, PLAYER2_COLOR, PLAYER2_TRAIL, {'up':pygame.K_w,'down':pygame.K_s,'left':pygame.K_a,'right':pygame.K_d})
    enemies = [Enemy(random.randint(0, WORLD_WIDTH), random.randint(0, WORLD_HEIGHT)) for _ in range(NUM_ENEMIES)]
    players = [player1, player2] + enemies

    powerups = []
    for _ in range(NUM_POWERUPS):
        while True:
            px, py = random.randint(0, WORLD_WIDTH//CELL_SIZE) * CELL_SIZE, random.randint(0, WORLD_HEIGHT//CELL_SIZE) * CELL_SIZE
            if (px, py) not in [(x, y) for x, y, _ in powerups]:
                powerups.append((px, py, random.choice(['boost', 'shield', 'nuke'])))
                break

    paused = False

    while True:
        clock.tick(FPS)
        offset = [player1.pos[0] - WIDTH // 2, player1.pos[1] - HEIGHT // 2]
        screen.fill(BG_COLOR)

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_p:
                    paused = not paused
                elif event.key == pygame.K_s:
                    draw_shop(player1)

        if not paused:
            for player in [player1, player2]:
                if player.alive:
                    player.handle_input()
                    player.move()
                    player.check_collision(players)
            for enemy in enemies:
                enemy.move_ai(players)
                if enemy.alive:
                    enemy.check_collision(players)

        for p in players:
            if p.alive:
                p.draw(offset)

        for p in powerups[:]:
            rect = (p[0]-offset[0], p[1]-offset[1], CELL_SIZE, CELL_SIZE)
            pygame.draw.rect(screen, POWERUP_COLOR, rect)
            for player in players:
                if player.alive and tuple(player.pos) == (p[0], p[1]):
                    if p[2] == 'boost':
                        player.speed_boost = SPEED_BOOST_DURATION
                    elif p[2] == 'shield':
                        player.shield = True
                    elif p[2] == 'nuke':
                        for enemy in enemies:
                            if enemy.alive:
                                enemy.die()
                    powerups.remove(p)

        draw_ui(player1)

        if paused:
            draw_text("Paused - Press P to Resume", (WIDTH // 2, HEIGHT // 2), color=(255, 255, 0), center=True)

        if not player1.alive and not player2.alive:
            save_score(player1.score + player2.score)
            draw_text("Game Over! Press R to Restart or Q to Quit", (WIDTH // 2, HEIGHT // 2), color=(255, 0, 0), center=True)
            pygame.display.flip()
            while True:
                event = pygame.event.wait()
                if event.type == pygame.QUIT:
                    pygame.quit()
                    sys.exit()
                if event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_r:
                        game_loop()
                    elif event.key == pygame.K_q:
                        pygame.quit()
                        sys.exit()

        pygame.display.flip()

def main_menu():
    while True:
        screen.fill(BG_COLOR)
        draw_text("Snake IO Arena", (WIDTH // 2, 100), center=True)
        draw_text("1. Start Game", (WIDTH // 2, 200), center=True)
        draw_text("2. Leaderboard", (WIDTH // 2, 240), center=True)
        draw_text("3. Quit", (WIDTH // 2, 280), center=True)
        pygame.display.flip()

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_1:
                    game_loop()
                elif event.key == pygame.K_2:
                    screen.fill(BG_COLOR)
                    draw_text("Leaderboard", (WIDTH // 2, 100), center=True)
                    top_scores = show_leaderboard()
                    for i, score in enumerate(top_scores):
                        draw_text(f"{i+1}. {score}", (WIDTH // 2, 150 + i * 30), center=True)
                    draw_text("Press any key to return", (WIDTH // 2, 400), center=True)
                    pygame.display.flip()
                    waiting = True
                    while waiting:
                        for event in pygame.event.get():
                            if event.type == pygame.KEYDOWN:
                                waiting = False
                elif event.key == pygame.K_3:
                    pygame.quit()
                    sys.exit()

main_menu()