"""Frogger: guide a frog across busy traffic lanes to reach the far side.

The frog starts at the bottom row and must cross to the top. Traffic lanes
alternate direction; safe lanes (no cars) appear every few rows as rest zones.
The frog scores 100 each time it reaches the top and respawns at the bottom.
It earns +1 reward for each upward step and +100 for crossing, but the game
ends immediately if a car hits it.
"""

from retro.game import Game

BOARD_WIDTH = 20
BOARD_HEIGHT = 12
FRAMERATE = 8
GAME_DURATION = 60 * FRAMERATE  # ~60 seconds

# One entry per traffic row (rows 1–10). None marks a safe lane with no cars.
# Each traffic entry is (direction, speed, num_cars, car_length).
# direction: 1 = right, -1 = left
# speed: steps between moves (higher = slower)
# num_cars: how many cars in the lane
# car_length: how many cells each car occupies
LANES = [
    ( 1, 2, 2, 2),   # row 1  fast
    (-1, 4, 3, 2),   # row 2  slow, crowded
    ( 1, 1, 1, 3),   # row 3  very fast
    None,             # row 4  safe zone
    (-1, 3, 2, 2),   # row 5  medium
    ( 1, 2, 2, 3),   # row 6  fast, long cars
    (-1, 5, 3, 2),   # row 7  slow, crowded
    None,             # row 8  safe zone
    ( 1, 1, 2, 2),   # row 9  very fast
    (-1, 3, 2, 3),   # row 10 medium, long cars
]


class Frog:
    name = "Frog"
    character = 'O'
    color = "green_on_black"
    position = (0, 0)

    UP = (0, -1)
    DOWN = (0, 1)
    LEFT = (-1, 0)
    RIGHT = (1, 0)

    def __init__(self):
        self._pending_direction = None
        self._frames_remaining = GAME_DURATION

    def handle_keystroke(self, keystroke, game):
        if keystroke.name == "KEY_UP":
            self._pending_direction = self.UP
        elif keystroke.name == "KEY_DOWN":
            self._pending_direction = self.DOWN
        elif keystroke.name == "KEY_LEFT":
            self._pending_direction = self.LEFT
        elif keystroke.name == "KEY_RIGHT":
            self._pending_direction = self.RIGHT

    def play_turn(self, game):
        self._frames_remaining -= 1
        game.state['time_left'] = self._frames_remaining // FRAMERATE
        if self._frames_remaining <= 0:
            game.end()
            return

        if self._pending_direction is None:
            return
        dx, dy = self._pending_direction
        self._pending_direction = None

        bw, bh = game.board_size
        x, y = self.position
        nx, ny = x + dx, y + dy

        if not (0 <= nx < bw) or ny >= bh:
            game.state['reward'] -= 5
            return

        if ny < 0:
            game.state['score'] += 100
            game.state['reward'] += 100
            self._respawn(game)
            return

        self.position = (nx, ny)

        if ny < y:
            game.state['reward'] += 1

        for agent in game.agents:
            if hasattr(agent, '_is_car') and agent.position == self.position:
                game.state['reward'] -= 10
                game.end()
                return

    def _respawn(self, game):
        bw, bh = game.board_size
        self.position = (bw // 2, bh - 1)
        self._pending_direction = None


class CarSegment:
    """One cell of a multi-cell car. Moved by the lead Car each turn."""
    _is_car = True
    character = 'X'
    color = "red_on_black"

    def __init__(self, name, position):
        self.name = name
        self.position = position


class Car:
    """Lead cell of a multi-cell car. Moves itself and all trailing segments."""
    _is_car = True
    character = 'X'
    color = "red_on_black"

    def __init__(self, name, lane, speed, direction, start_x, board_width, segments):
        self.name = name
        self._lane = lane
        self._speed = speed
        self._direction = direction
        self._board_width = board_width
        self._step = 0
        self.position = (start_x, lane)
        self.segments = segments

    def play_turn(self, game):
        self._step += 1
        if self._step < self._speed:
            return
        self._step = 0
        bw = self._board_width
        x, y = self.position
        self.position = ((x + self._direction) % bw, y)
        for seg in self.segments:
            sx, sy = seg.position
            seg.position = ((sx + self._direction) % bw, sy)

        frog = game.get_agent_by_name("Frog")
        car_cells = {self.position} | {s.position for s in self.segments}
        if frog.position in car_cells:
            game.state['reward'] -= 10
            game.end()


def create_game():
    bw, bh = BOARD_WIDTH, BOARD_HEIGHT
    frog = Frog()
    frog.position = (bw // 2, bh - 1)
    agents = [frog]

    for lane_idx, row in enumerate(range(1, bh - 1)):
        spec = LANES[lane_idx] if lane_idx < len(LANES) else None
        if spec is None:
            continue
        direction, speed, num_cars, car_length = spec
        spacing = bw // num_cars
        for i in range(num_cars):
            start_x = (i * spacing + lane_idx * 3) % bw
            segments = []
            for j in range(1, car_length):
                seg_x = (start_x - direction * j) % bw
                seg = CarSegment(f"Car_{row}_{i}_seg{j}", (seg_x, row))
                segments.append(seg)
                agents.append(seg)
            agents.append(Car(f"Car_{row}_{i}", row, speed, direction, start_x, bw, segments))

    return Game(
        agents,
        {
            'score': 0,
            'reward': 0.0,
            'time_left': GAME_DURATION // FRAMERATE,
        },
        board_size=(bw, bh),
        framerate=FRAMERATE,
        show_state=['score', 'time_left'],
    )


if __name__ == '__main__':
    create_game().play()