level = [
    ['#', '#', '#', '#', '#'],
    ['#', '@', ' ', '.', '#'],
    ['#', ' ', '$', ' ', '#'],
    ['#', '.', '$', ' ', '#'],
    ['#', ' ', '$', '.', '#'],
    ['#', '.', '$', '.', '#'],
    ['#', '.', '*', ' ', '#'],
    ['#', ' ', '*', '.', '#'],
    ['#', ' ', '*', ' ', '#'],
    ['#', '.', '*', '.', '#'],
    ['#', '#', '#', '#', '#'],
]

player = '@'
player_on_storage = '+'
box = '$'
box_on_storage = '*'
storage = '.'
wall = '#'
empty = ' '

def on_key_down(key):
    if key in (keys.UP, keys.DOWN, keys.LEFT, keys.RIGHT):
        for test_y, row in enumerate(level):
            for test_x, cell in enumerate(row):
                if cell == player or cell == player_on_storage:
                    player_x = test_x
                    player_y = test_y

        dx = 0
        dy = 0
        if key == keys.LEFT:
            dx = -1
        elif key == keys.RIGHT:
            dx = 1
        elif key == keys.UP:
            dy = -1
        elif key == keys.DOWN:
            dy = 1

        current = level[player_y][player_x]
        adjacent = level[player_y + dy][player_x + dx]
        beyond = ''
        if (
            0 <= player_y + dy + dy < len(level)
            and 0 <= player_x + dx + dx < len(level[player_y + dy + dy])
        ):
            beyond = level[player_y + dy + dy][player_x + dx + dx]

        next_adjacent = {
            empty: player,
            storage: player_on_storage,
        }
        next_current = {
            player: empty,
            player_on_storage: storage,
        }

        if adjacent in next_adjacent:
            level[player_y][player_x] = next_current[current]
            level[player_y + dy][player_x + dx] = next_adjacent[adjacent]

        elif adjacent == box:
            if beyond == empty:
                level[player_y][player_x] = next_current[current]
                level[player_y + dy][player_x + dx] = player
                level[player_y + dy + dy][player_x + dx + dx] = box
            elif beyond == storage:
                level[player_y][player_x] = next_current[current]
                level[player_y + dy][player_x + dx] = player
                level[player_y + dy + dy][player_x + dx + dx] = box_on_storage

def draw():
    screen.fill((255, 255, 190))

    for y, row in enumerate(level):
        for x, cell in enumerate(row):
            if cell != empty:
                cell_size = 23

                colors = {
                    player: (167, 135, 255),
                    player_on_storage: (158, 119, 255),
                    box: (255, 201, 126),
                    box_on_storage: (150, 255, 127),
                    storage: (156, 229, 255),
                    wall: (255, 147, 209),
                }

                screen.draw.filled_rect(
                    Rect(
                        (x * cell_size, y * cell_size),
                        (cell_size, cell_size)
                    ),
                    color=colors[cell]
                )

                screen.draw.text(
                    cell,
                    (x * cell_size, y * cell_size),
                    color=(255, 255, 255)
                )