game.py

import random
import pygame

from assets import load_deck_from_folder_unique, load_images_scaled
from animations import GlideAnim, ShakeAnim, add_stack_draw_anims
from models import Card
from rules import can_play, color_name, top_card_label
from settings import (
    WIDTH, HEIGHT, FPS, UI_BG, UI_TEXT, UI_MUTED, UI_BORDER,
    CENTER_Y, DISCARD_CENTER, DRAW_CENTER, DRAW_EDGE,
    STATUS_DURATION
)
from ui import (
    draw_text, draw_rotated_card, draw_turn_overlay, draw_game_over_menu,
    ensure_discard_transforms, draw_discard_pile, slot_position_for_card,
    draw_player_hand, draw_status_message, get_hand_positions
)


DrawEvent = tuple[int, Card, int]


class UnoGame:
    def __init__(self, deck: list[Card]):
        self.player_count = 2
        self.deck = deck[:]
        random.shuffle(self.deck)

        self.discard = []
        self.hands = [[] for _ in range(self.player_count)]

        self.turn = 0
        self.direction = 1
        self.current_color = None

        self.awaiting_color_choice = False
        self.pending_wild_card = None
        self.winner = None

        self.deal(7)
        self.start_discard()

    @property
    def top_card(self):
        return self.discard[-1] if self.discard else None

    def draw_card(self):
        if not self.deck:
            if len(self.discard) > 1:
                top = self.discard[-1]
                rest = self.discard[:-1]
                random.shuffle(rest)
                self.deck = rest
                self.discard = [top]
            else:
                return None
        return self.deck.pop()

    def deal(self, n: int):
        for _ in range(n):
            for p in range(self.player_count):
                c = self.draw_card()
                if c:
                    self.hands[p].append(c)

    def start_discard(self):
        while True:
            c = self.draw_card()
            if not c:
                break
            self.discard.append(c)
            if not c.is_wild():
                self.current_color = c.color
                break

        if self.current_color is None:
            self.current_color = random.choice(["r", "g", "b", "y"])

    def next_player(self, skip=0):
        steps = 1 + skip
        self.turn = (self.turn + steps * self.direction) % self.player_count

    def _apply_effect(self, card: Card):
        draw_events = []

        if card.kind == "skip":
            return 1, draw_events

        if card.kind == "reverse":
            self.direction *= -1
            return 1, draw_events

        if card.kind in ("+2", "+4"):
            target = (self.turn + self.direction) % self.player_count
            count = 2 if card.kind == "+2" else 4

            start_len = len(self.hands[target])
            for j in range(count):
                d = self.draw_card()
                if d:
                    self.hands[target].append(d)
                    draw_events.append((target, d, start_len + j))
            return 1, draw_events

        return 0, draw_events

    def play_card(self, player_idx: int, hand_index: int):
        if self.winner is not None:
            return None, [], "Game is over"
        if self.awaiting_color_choice:
            return None, [], "Choose a color first"
        if player_idx != self.turn:
            return None, [], "Not your turn"
        if hand_index < 0 or hand_index >= len(self.hands[player_idx]):
            return None, [], "Invalid card"

        card = self.hands[player_idx][hand_index]
        if not can_play(card, self.top_card, self.current_color):
            return None, [], "Invalid move"

        self.hands[player_idx].pop(hand_index)
        self.discard.append(card)

        if len(self.hands[player_idx]) == 0:
            self.winner = player_idx
            return card, [], None

        if card.is_wild():
            self.awaiting_color_choice = True
            self.pending_wild_card = card
            return card, [], "Choose a color"

        self.current_color = card.color
        skip, draw_events = self._apply_effect(card)
        self.next_player(skip)
        return card, draw_events, None

    def choose_color(self, chosen_color: str):
        if not self.awaiting_color_choice:
            return [], "No color choice needed"

        self.current_color = chosen_color
        wild_card = self.pending_wild_card

        self.awaiting_color_choice = False
        self.pending_wild_card = None

        skip, draw_events = self._apply_effect(wild_card)
        self.next_player(skip)
        return draw_events, f"Color set to {color_name(chosen_color)}"

    def draw_for_current(self):
        if self.winner is not None or self.awaiting_color_choice:
            return None, None, None, "Cannot draw now"

        drawing_player = self.turn
        before = len(self.hands[drawing_player])

        c = self.draw_card()
        if c:
            self.hands[drawing_player].append(c)

        self.next_player(skip=0)
        return drawing_player, c, before, f"Player {drawing_player + 1} drew a card"


def run_game():
    pygame.init()
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    pygame.display.set_caption("UNO")
    clock = pygame.time.Clock()

    small = pygame.font.SysFont("arial", 16)
    font = pygame.font.SysFont("arial", 20)
    big = pygame.font.SysFont("arial", 32, bold=True)
    mid = pygame.font.SysFont("arial", 22, bold=True)

    images, back_img = load_images_scaled("cards")

    def new_game():
        deck = load_deck_from_folder_unique("cards")
        return UnoGame(deck)

    game = new_game()

    pile_rng = random.Random()
    pile_rng.seed()
    discard_transforms = []
    ensure_discard_transforms(discard_transforms, len(game.discard), pile_rng)

    turn_gate = True
    animations = []

    hand_shakes = {0: None, 1: None}
    status_text = ""
    status_timer = 0.0
    hover_index = None

    color_buttons = {
        "r": pygame.Rect(16, CENTER_Y - 70, 130, 40),
        "g": pygame.Rect(16, CENTER_Y - 25, 130, 40),
        "b": pygame.Rect(16, CENTER_Y + 20, 130, 40),
        "y": pygame.Rect(16, CENTER_Y + 65, 130, 40),
    }

    running = True
    while running:
        dt = clock.tick(FPS) / 1000.0

        for a in animations:
            a.update(dt)
        animations = [a for a in animations if not a.done]

        for player, shake in hand_shakes.items():
            if shake:
                shake.update(dt)
                if shake.done:
                    hand_shakes[player] = None

        if status_timer > 0:
            status_timer -= dt
            if status_timer <= 0:
                status_text = ""

        ensure_discard_transforms(discard_transforms, len(game.discard), pile_rng)

        screen.fill(UI_BG)

        if game.winner is None:
            draw_text(screen, font, f"P{game.turn + 1} | color: {color_name(game.current_color)}", 12, 10, UI_TEXT)
        else:
            draw_text(screen, font, "Game Over", 12, 10, UI_TEXT)

        draw_text(screen, small, top_card_label(game.top_card, game.current_color), 12, 36, UI_TEXT)
        draw_discard_pile(screen, images, game.discard, discard_transforms)

        draw_pile_rect = draw_rotated_card(screen, back_img, DRAW_CENTER, angle_deg=12)
        draw_text(screen, font, "DRAW", DRAW_CENTER[0] - 22, DRAW_CENTER[1] + 72, UI_MUTED)

        if game.winner is None and game.awaiting_color_choice and not turn_gate:
            draw_text(screen, mid, f"P{game.turn + 1}: choose color", 12, CENTER_Y - 120, UI_TEXT)
            fill_map = {
                "r": (220, 70, 70),
                "g": (70, 200, 110),
                "b": (80, 130, 240),
                "y": (240, 210, 80),
            }
            for c, rect in color_buttons.items():
                pygame.draw.rect(screen, fill_map[c], rect, border_radius=12)
                pygame.draw.rect(screen, UI_BORDER, rect, 2, border_radius=12)
                draw_text(screen, font, color_name(c), rect.x + 18, rect.y + 10, UI_TEXT)

        p1_offset = hand_shakes[0].offset() if hand_shakes[0] else 0
        p2_offset = hand_shakes[1].offset() if hand_shakes[1] else 0

        click_targets = []
        hover_index = None

        if turn_gate or game.winner is not None:
            draw_player_hand(screen, game.hands[1], False, 1, images, back_img, None, p2_offset)
            draw_player_hand(screen, game.hands[0], False, 0, images, back_img, None, p1_offset)
        else:
            mx, my = pygame.mouse.get_pos()

            if game.turn == 0:
                draw_player_hand(screen, game.hands[1], False, 1, images, back_img, None, p2_offset)

                positions = get_hand_positions(0, len(game.hands[0]))
                for i in reversed(range(len(positions))):
                    x, y, ang = positions[i]
                    x += p1_offset
                    test_rect = images[game.hands[0][i].filename].get_rect(center=(x, y))
                    if test_rect.collidepoint(mx, my):
                        hover_index = i
                        break

                click_targets = draw_player_hand(
                    screen, game.hands[0], True, 0, images, back_img, hover_index, p1_offset
                )

            else:
                draw_player_hand(screen, game.hands[0], False, 0, images, back_img, None, p1_offset)

                positions = get_hand_positions(1, len(game.hands[1]))
                for i in reversed(range(len(positions))):
                    x, y, ang = positions[i]
                    x += p2_offset
                    test_rect = images[game.hands[1][i].filename].get_rect(center=(x, y))
                    if test_rect.collidepoint(mx, my):
                        hover_index = i
                        break

                click_targets = draw_player_hand(
                    screen, game.hands[1], True, 1, images, back_img, hover_index, p2_offset
                )

        draw_text(screen, font, f"P1: {len(game.hands[0])}", 12, HEIGHT - 28, UI_MUTED)
        draw_text(screen, font, f"P2: {len(game.hands[1])}", 12, 56, UI_MUTED)

        for a in animations:
            a.draw(screen)

        draw_status_message(screen, small, status_text)

        if game.winner is None and turn_gate:
            draw_turn_overlay(screen, big, font, game.turn)

        play_again_rect = exit_rect = None
        if game.winner is not None:
            play_again_rect, exit_rect = draw_game_over_menu(screen, big, font, game.winner)

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

            if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
                mx, my = event.pos

                if game.winner is not None and play_again_rect and exit_rect:
                    if play_again_rect.collidepoint(mx, my):
                        game = new_game()
                        animations.clear()
                        turn_gate = True
                        discard_transforms.clear()
                        ensure_discard_transforms(discard_transforms, len(game.discard), pile_rng)
                        hand_shakes = {0: None, 1: None}
                        status_text = "New game started"
                        status_timer = STATUS_DURATION
                    elif exit_rect.collidepoint(mx, my):
                        running = False
                    continue

                if game.winner is None and turn_gate:
                    turn_gate = False
                    continue

                if game.winner is not None:
                    continue

                if game.awaiting_color_choice:
                    for c, rect in color_buttons.items():
                        if rect.collidepoint(mx, my):
                            prev_turn = game.turn
                            draw_events, msg = game.choose_color(c)

                            status_text = msg
                            status_timer = STATUS_DURATION

                            for j, (target_player, drawn_card, hand_idx) in enumerate(draw_events):
                                total = len(game.hands[target_player])
                                end_pos, end_ang = slot_position_for_card(target_player, hand_idx, total)
                                delay = 0.12 * j
                                add_stack_draw_anims(
                                    animations, images[drawn_card.filename], end_pos, end_ang, delay,
                                    DRAW_CENTER, DRAW_EDGE
                                )

                            if game.turn != prev_turn:
                                turn_gate = True
                            break
                    continue

                if draw_pile_rect.collidepoint(mx, my):
                    prev_turn = game.turn
                    drawing_player, drawn, hand_idx, msg = game.draw_for_current()

                    status_text = msg
                    status_timer = STATUS_DURATION

                    if drawn and drawing_player is not None and hand_idx is not None:
                        total = len(game.hands[drawing_player])
                        end_pos, end_ang = slot_position_for_card(drawing_player, hand_idx, total)
                        add_stack_draw_anims(
                            animations, images[drawn.filename], end_pos, end_ang, 0.0,
                            DRAW_CENTER, DRAW_EDGE
                        )

                    if game.turn != prev_turn:
                        turn_gate = True
                    continue

                for rect, idx, card_center, card_angle in click_targets:
                    if rect.collidepoint(mx, my):
                        prev_turn = game.turn
                        current_player = game.turn
                        card_obj = game.hands[current_player][idx]

                        played, draw_events, msg = game.play_card(current_player, idx)

                        if played is None:
                            hand_shakes[current_player] = ShakeAnim()
                            status_text = msg
                            status_timer = STATUS_DURATION
                            break

                        animations.append(
                            GlideAnim(
                                image=images[card_obj.filename],
                                start=card_center,
                                end=DISCARD_CENTER,
                                angle_start=card_angle,
                                angle_end=-10,
                                duration=0.20,
                                delay=0.0
                            )
                        )

                        ensure_discard_transforms(discard_transforms, len(game.discard), pile_rng)

                        for j, (target_player, drawn_card, hand_idx) in enumerate(draw_events):
                            total = len(game.hands[target_player])
                            end_pos, end_ang = slot_position_for_card(target_player, hand_idx, total)
                            delay = 0.12 * j
                            add_stack_draw_anims(
                                animations, images[drawn_card.filename], end_pos, end_ang, delay,
                                DRAW_CENTER, DRAW_EDGE
                            )

                        if msg:
                            status_text = msg
                            status_timer = STATUS_DURATION

                        if game.turn != prev_turn and not game.awaiting_color_choice and game.winner is None:
                            turn_gate = True
                        break

        pygame.display.flip()

    pygame.quit()