import pygame
import sys
import random
import math

# --- 初始化 ---
pygame.init()
CELL = 30
COLS, ROWS = 10, 20
WIDTH, HEIGHT = CELL * COLS, CELL * ROWS
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Pygame 俄罗斯方块 - SRS + Ghost")
clock = pygame.time.Clock()

# 颜色
BG_COLOR = (20, 20, 30)
GRID_COLOR = (40, 40, 50)
GHOST_COLOR = (100, 100, 120, 100)
TEXT_COLOR = (255, 255, 255)

# 7种标准方块颜色
COLORS = [
    (0, 240, 240),  # I - 青
    (240, 240, 0),  # O - 黄
    (160, 0, 240),  # T - 紫
    (0, 240, 0),    # S - 绿
    (240, 0, 0),    # Z - 红
    (0, 0, 240),    # J - 蓝
    (240, 160, 0),  # L - 橙
]

# 方块形状 (4种旋转状态)
SHAPES = [
    # I
    [[(0,1), (1,1), (2,1), (3,1)], [(2,0), (2,1), (2,2), (2,3)], [(0,2), (1,2), (2,2), (3,2)], [(1,0), (1,1), (1,2), (1,3)]],
    # O
    [[(1,0), (2,0), (1,1), (2,1)], [(1,0), (2,0), (1,1), (2,1)], [(1,0), (2,0), (1,1), (2,1)], [(1,0), (2,0), (1,1), (2,1)]],
    # T
    [[(1,0), (0,1), (1,1), (2,1)], [(1,0), (1,1), (2,1), (1,2)], [(0,1), (1,1), (2,1), (1,2)], [(1,0), (0,1), (1,1), (1,2)]],
    # S
    [[(1,0), (2,0), (0,1), (1,1)], [(1,0), (1,1), (2,1), (2,2)], [(1,1), (2,1), (0,2), (1,2)], [(0,0), (0,1), (1,1), (1,2)]],
    # Z
    [[(0,0), (1,0), (1,1), (2,1)], [(2,0), (1,1), (2,1), (1,2)], [(0,1), (1,1), (1,2), (2,2)], [(1,0), (0,1), (1,1), (0,2)]],
    # J
    [[(0,0), (0,1), (1,1), (2,1)], [(1,0), (2,0), (1,1), (1,2)], [(0,1), (1,1), (2,1), (2,2)], [(1,0), (1,1), (0,2), (1,2)]],
    # L
    [[(2,0), (0,1), (1,1), (2,1)], [(1,0), (1,1), (1,2), (2,2)], [(0,1), (1,1), (2,1), (0,2)], [(0,0), (1,0), (1,1), (1,2)]],
]

# 简化版踢墙数据 (仅基础偏移，足够日常游玩)
KICKS = {
    'normal': [
        [(0,0), (-1,0), (-1,1), (0,-2), (-1,-2)],
        [(0,0), (1,0), (1,-1), (0,2), (1,2)],
        [(0,0), (1,0), (1,1), (0,-2), (1,-2)],
        [(0,0), (-1,0), (-1,-1), (0,2), (-1,2)]
    ],
    'I': [
        [(0,0), (-2,0), (1,0), (-2,-1), (1,2)],
        [(0,0), (2,0), (-1,0), (2,1), (-1,-2)],
        [(0,0), (-1,0), (2,0), (-1,2), (2,-1)],
        [(0,0), (1,0), (-2,0), (1,-2), (-2,1)]
    ]
}

class Piece:
    def __init__(self, type_idx=None):
        self.type = type_idx if type_idx is not None else random.randint(0, 6)
        self.rot = 0
        self.x = 3
        self.y = 0
        self.color = COLORS[self.type]
        self.lock_delay = 0
        self.max_lock_delay = 30  # 0.5秒 @60fps

    def get_cells(self, rot=None, x=None, y=None):
        r = rot if rot is not None else self.rot
        px = x if x is not None else self.x
        py = y if y is not None else self.y
        return [(px + cx, py + cy) for cx, cy in SHAPES[self.type][r]]

    def is_valid(self, board, rot=None, x=None, y=None):
        cells = self.get_cells(rot, x, y)
        for cx, cy in cells:
            if cx < 0 or cx >= COLS or cy >= ROWS: return False
            if cy >= 0 and board[cy][cx] != 0: return False
        return True

    def move(self, board, dx, dy):
        if self.is_valid(board, x=self.x+dx, y=self.y+dy):
            self.x += dx
            self.y += dy
            if dy == 0 and self.is_valid(board, y=self.y+1):
                self.lock_delay = 0  # 水平移动重置锁定
            return True
        return False

    def rotate(self, board):
        new_rot = (self.rot + 1) % 4
        kicks = KICKS['I'] if self.type == 0 else KICKS['normal']
        kick_data = kicks[self.rot]
        
        for kx, ky in kick_data:
            if self.is_valid(board, rot=new_rot, x=self.x+kx, y=self.y+ky):
                self.rot = new_rot
                self.x += kx
                self.y += ky
                self.lock_delay = 0
                return True
        return False

    def hard_drop(self, board):
        dist = 0
        while self.is_valid(board, y=self.y+1):
            self.y += 1
            dist += 1
        self.lock_delay = self.max_lock_delay  # 硬降直接锁
        return dist

    def get_ghost_y(self, board):
        gy = self.y
        while self.is_valid(board, y=gy+1):
            gy += 1
        return gy

# --- 主程序 ---
def main():
    board = [[0]*COLS for _ in range(ROWS)]
    current = Piece()
    next_piece = Piece()
    
    score = 0
    level = 1
    lines = 0
    game_over = False
    fall_timer = 0
    flash_timer = 0
    clear_rows = []

    def get_fall_speed():
        return max(2, 30 - (level - 1) * 3)  # 等级越高下落越快

    def lock_piece():
        nonlocal score, lines, level, flash_timer, clear_rows
        cells = current.get_cells()
        for cx, cy in cells:
            if cy < 0:
                return True  # Game Over
            board[cy][cx] = current.color
        
        # 消行检测
        clear_rows = [i for i in range(ROWS) if all(board[i])]
        if clear_rows:
            flash_timer = 20
            lines += len(clear_rows)
            score += [0, 100, 300, 500, 800][len(clear_rows)] * level
            level = lines // 10 + 1
        else:
            spawn_next()
        return False

    def spawn_next():
        nonlocal current, next_piece
        current = next_piece
        next_piece = Piece()
        if not current.is_valid(board):
            return True
        return False

    running = True
    while running:
        clock.tick(60)
        for event in pygame.event.get():
            if event.type == pygame.QUIT: pygame.quit(); sys.exit()
            if event.type == pygame.KEYDOWN and not game_over:
                if event.key == pygame.K_LEFT: current.move(board, -1, 0)
                if event.key == pygame.K_RIGHT: current.move(board, 1, 0)
                if event.key == pygame.K_DOWN: 
                    if current.move(board, 0, 1): score += 1
                if event.key == pygame.K_UP: current.rotate(board)
                if event.key == pygame.K_SPACE:
                    score += current.hard_drop(board) * 2
                    if lock_piece(): game_over = True
                if event.key == pygame.K_r and game_over:
                    board = [[0]*COLS for _ in range(ROWS)]
                    current = Piece()
                    next_piece = Piece()
                    score = 0; lines = 0; level = 1
                    game_over = False

        if not game_over and flash_timer <= 0:
            fall_timer += 1
            if fall_timer >= get_fall_speed():
                fall_timer = 0
                if not current.move(board, 0, 1):
                    current.lock_delay += 1
                    if current.lock_delay >= current.max_lock_delay:
                        if lock_piece(): game_over = True
                else:
                    current.lock_delay = 0

        # 消行动画
        if flash_timer > 0:
            flash_timer -= 1
            if flash_timer == 0:
                for row in sorted(clear_rows, reverse=True):
                    board.pop(row)
                    board.insert(0, [0]*COLS)
                spawn_next()

        # --- 绘图 ---
        screen.fill(BG_COLOR)
        
        # 网格
        for x in range(COLS):
            for y in range(ROWS):
                if board[y][x]:
                    pygame.draw.rect(screen, board[y][x], (x*CELL, y*CELL, CELL-1, CELL-1))
                else:
                    pygame.draw.rect(screen, GRID_COLOR, (x*CELL, y*CELL, CELL-1, CELL-1), 1)

        if not game_over and flash_timer <= 0:
            # Ghost
            ghost_y = current.get_ghost_y(board)
            for gx, gy in current.get_cells(y=ghost_y):
                if gy >= 0:
                    s = pygame.Surface((CELL-1, CELL-1), pygame.SRCALPHA)
                    s.fill((*current.color, 80))
                    screen.blit(s, (gx*CELL, gy*CELL))

            # Current
            for cx, cy in current.get_cells():
                if cy >= 0:
                    pygame.draw.rect(screen, current.color, (cx*CELL, cy*CELL, CELL-1, CELL-1))
                    pygame.draw.rect(screen, (255,255,255), (cx*CELL, cy*CELL, CELL-1, CELL-1), 1)

        # UI
        font = pygame.font.SysFont("consolas", 24)
        screen.blit(font.render(f"Score: {score}", True, TEXT_COLOR), (10, HEIGHT-30))
        screen.blit(font.render(f"Level: {level}", True, TEXT_COLOR), (WIDTH-100, HEIGHT-30))

        if game_over:
            overlay = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
            overlay.fill((0,0,0,180))
            screen.blit(overlay, (0,0))
            big = pygame.font.SysFont("consolas", 50)
            txt = big.render("GAME OVER", True, (255,80,80))
            sub = font.render("Press R to Restart", True, TEXT_COLOR)
            screen.blit(txt, txt.get_rect(center=(WIDTH//2, HEIGHT//2-30)))
            screen.blit(sub, sub.get_rect(center=(WIDTH//2, HEIGHT//2+30)))

        pygame.display.flip()

if __name__ == "__main__":
    main()