import curses
import random
import time

# Tetris shapes: (rotations)
TETROMINOES = {
    'I': [[1, 1, 1, 1]],
    'O': [[1, 1], [1, 1]],
    'T': [[0, 1, 0], [1, 1, 1]],
    'S': [[0, 1, 1], [1, 1, 0]],
    'Z': [[1, 1, 0], [0, 1, 1]],
    'J': [[1, 0, 0], [1, 1, 1]],
    'L': [[0, 0, 1], [1, 1, 1]],
}

COLORS = {
    'I': 1,
    'O': 2,
    'T': 3,
    'S': 4,
    'Z': 5,
    'J': 6,
    'L': 7,
}

WIDTH, HEIGHT = 10, 20

class Tetromino:
    def __init__(self, shape):
        self.shape = shape
        self.rot = 0
        self.x = WIDTH // 2 - len(self.current()[0]) // 2
        self.y = 0

    def current(self):
        s = TETROMINOES[self.shape]
        for _ in range(self.rot % 4):
            s = [list(row) for row in zip(*s[::-1])]
        return s

    def rotate(self):
        self.rot = (self.rot + 1) % 4

    def unrotate(self):
        self.rot = (self.rot - 1) % 4


def create_board():
    return [[0 for _ in range(WIDTH)] for _ in range(HEIGHT)]


def check_collision(board, tetro, dx=0, dy=0):
    shape = tetro.current()
    for y, row in enumerate(shape):
        for x, cell in enumerate(row):
            if cell:
                nx, ny = tetro.x + x + dx, tetro.y + y + dy
                if nx < 0 or nx >= WIDTH or ny < 0 or ny >= HEIGHT:
                    return True
                if ny >= 0 and board[ny][nx]:
                    return True
    return False


def merge_tetromino(board, tetro):
    shape = tetro.current()
    for y, row in enumerate(shape):
        for x, cell in enumerate(row):
            if cell:
                nx, ny = tetro.x + x, tetro.y + y
                if 0 <= ny < HEIGHT and 0 <= nx < WIDTH:
                    board[ny][nx] = COLORS[tetro.shape]


def clear_lines(board):
    new_board = [row for row in board if any(cell == 0 for cell in row)]
    lines_cleared = HEIGHT - len(new_board)
    for _ in range(lines_cleared):
        new_board.insert(0, [0 for _ in range(WIDTH)])
    return new_board, lines_cleared


def draw_board(stdscr, board, tetro, score):
    stdscr.clear()
    for y, row in enumerate(board):
        for x, cell in enumerate(row):
            if cell:
                stdscr.addstr(y, x * 2, '[]', curses.color_pair(cell))
            else:
                stdscr.addstr(y, x * 2, '  ')
    # Draw current tetromino
    shape = tetro.current()
    for y2, row in enumerate(shape):
        for x2, cell in enumerate(row):
            if cell:
                nx, ny = tetro.x + x2, tetro.y + y2
                if 0 <= ny < HEIGHT and 0 <= nx < WIDTH:
                    stdscr.addstr(ny, nx * 2, '[]', curses.color_pair(COLORS[tetro.shape]))
    stdscr.addstr(0, WIDTH * 2 + 2, f'Score: {score}')
    stdscr.refresh()


def main(stdscr):
    curses.curs_set(0)
    stdscr.nodelay(True)
    stdscr.timeout(100)
    for i in range(1, 8):
        curses.init_pair(i, i, 0)

    board = create_board()
    current = Tetromino(random.choice(list(TETROMINOES.keys())))
    next_drop = time.time() + 0.5
    score = 0

    while True:
        draw_board(stdscr, board, current, score)
        key = stdscr.getch()
        if key == curses.KEY_LEFT:
            if not check_collision(board, current, dx=-1):
                current.x -= 1
        elif key == curses.KEY_RIGHT:
            if not check_collision(board, current, dx=1):
                current.x += 1
        elif key == curses.KEY_DOWN:
            if not check_collision(board, current, dy=1):
                current.y += 1
        elif key == curses.KEY_UP:
            current.rotate()
            if check_collision(board, current):
                current.unrotate()
        elif key == ord('q'):
            break

        if time.time() > next_drop:
            if not check_collision(board, current, dy=1):
                current.y += 1
            else:
                merge_tetromino(board, current)
                board, lines = clear_lines(board)
                score += lines * 100
                current = Tetromino(random.choice(list(TETROMINOES.keys())))
                if check_collision(board, current):
                    stdscr.addstr(HEIGHT // 2, WIDTH, 'GAME OVER')
                    stdscr.refresh()
                    time.sleep(2)
                    break
            next_drop = time.time() + 0.5
        time.sleep(0.01)

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