import curses
import time

import numpy as np

PIECES = {
    "O": np.array([[1, 1], [1, 1]]),  # 2x2 block
    "I": np.array([[1, 1, 1, 1]]),  # 4x1 horizontal bar
    "L": np.array([[1, 0], [1, 0], [1, 1]]),  # L-shaped
}


def check_placement(piece, board, i, j):
    """Check if a piece can be placed on the board. Return True if possible, False otherwise."""
    n, m = piece.shape
    if i + n > board.shape[0] or j + m > board.shape[1]:
        return False  # Out of bounds
    for x in range(n):
        for y in range(m):
            if piece[x, y] == 1 and board[i + x, j + y] == 1:
                return False
    return True


def place_piece(piece, board, i, j):
    """Place a piece on the board at the specified position."""
    board[i : i + piece.shape[0], j : j + piece.shape[1]] += piece
    return board


def attempt_rotate(piece, board, i, j):
    """Attempt to rotate a piece."""
    rotated_piece = np.rot90(piece)
    if check_placement(rotated_piece, board, i, j):
        return rotated_piece
    return piece


def clear_rows(board):
    """Clear completed rows from the board."""
    clear_indices = np.where(np.all(board == 1, axis=1))[0]
    if clear_indices.size > 0:
        board = np.delete(board, clear_indices, axis=0)
        new_rows = np.zeros((len(clear_indices), board.shape[1]))
        board = np.vstack([new_rows, board])
    return board


def calculate_shadow(piece, board, i, j):
    """Calculate the shadow position for the current piece."""
    while check_placement(piece, board, i + 1, j):
        i += 1
    return i


def display_board(stdscr, board, piece, i, j):
    """Display the board and the piece."""
    stdscr.clear()
    shadow_i = calculate_shadow(piece, board, i, j)
    n, m = piece.shape
    for y in range(board.shape[0]):
        for x in range(board.shape[1]):
            char = " "  # Default background
            if (
                y >= shadow_i
                and y < shadow_i + n
                and x >= j
                and x < j + m
                and piece[y - shadow_i, x - j]
            ):
                char = "*"  # Shadow
            elif y >= i and y < i + n and x >= j and x < j + m and piece[y - i, x - j]:
                char = "#"  # Piece
            elif board[y, x]:
                char = "#"  # Static pieces
            stdscr.addstr(y, x * 2, char)
    stdscr.refresh()


def main(stdscr):
    curses.curs_set(0)
    stdscr.nodelay(True)
    stdscr.keypad(True)

    board = np.zeros((20, 10), dtype=int)
    piece_types = list(PIECES.keys())
    current_piece = PIECES[np.random.choice(piece_types)]
    START_POS = (0, 4)
    i, j = START_POS
    auto_gravity = True
    fall_rate = 0.5  # Seconds between falls

    # Initial instruction
    stdscr.addstr(
        22,
        0,
        "Press SPACE to start. Use keys (a=left, d=right, s=down, w=rotate, S=drop, G=toggle gravity, Q=quit):",
    )
    while stdscr.getch() != ord(" "):
        pass

    last_fall_time = time.time()

    while True:
        display_board(stdscr, board, current_piece, i, j)

        if not check_placement(current_piece, board, *START_POS):
            stdscr.addstr(22, 0, "GAME OVER! Press any key to exit.")
            stdscr.getch()
            break

        # Handle key inputs
        key = stdscr.getch()
        if key == ord("q"):
            break
        elif (
            key == ord("a")
            and j > 0
            and check_placement(current_piece, board, i, j - 1)
        ):
            j -= 1
        elif (
            key == ord("d")
            and j + current_piece.shape[1] < board.shape[1]
            and check_placement(current_piece, board, i, j + 1)
        ):
            j += 1
        elif key == ord("w"):
            new_piece = attempt_rotate(current_piece, board, i, j)
            if check_placement(new_piece, board, i, j):
                current_piece = new_piece
        elif key == ord("s"):
            if check_placement(current_piece, board, i + 1, j):
                i += 1
            else:
                board = place_piece(current_piece, board, i, j)
                board = clear_rows(board)
                current_piece = PIECES[np.random.choice(piece_types)]
                i, j = START_POS
        elif key == ord("S"):
            i = calculate_shadow(current_piece, board, i, j)
            board = place_piece(current_piece, board, i, j)
            board = clear_rows(board)
            current_piece = PIECES[np.random.choice(piece_types)]
            i, j = START_POS
        elif key == ord("g"):
            auto_gravity = not auto_gravity

        # Gravity effect
        if auto_gravity and time.time() - last_fall_time >= fall_rate:
            if check_placement(current_piece, board, i + 1, j):
                i += 1
            else:
                board = place_piece(current_piece, board, i, j)
                board = clear_rows(board)
                current_piece = PIECES[np.random.choice(piece_types)]
                i, j = START_POS
            last_fall_time = time.time()


if __name__ == "__main__":
    curses.wrapper(main)