import pygame
import chess
import chess.engine

# Initialize Pygame
pygame.init()

# Constants
WIDTH, HEIGHT = 900, 600
SQ_SIZE = 60
EVAL_BAR_WIDTH = 60
BUTTON_HEIGHT = 40
BUTTON_WIDTH = 80
IMAGES = {}
MOVE_DELAY_TIME = 1000  # Delay in milliseconds

# Colors
LIGHT_SQ_COLOR = pygame.Color("#F0D9B5")
DARK_SQ_COLOR = pygame.Color("#B58863")
HIGHLIGHT_COLOR = pygame.Color("yellow")
PLAYER_MOVE_COLOR = pygame.Color("blue")
ENGINE_MOVE_COLOR = pygame.Color("red")
TOP_LINE_COLOR = pygame.Color((255, 255, 0, 100))

def load_images():
    pieces = ['wp', 'wN', 'wB', 'wR', 'wQ', 'wK', 'bp', 'bN', 'bB', 'bR', 'bQ', 'bK']
    for piece in pieces:
        try:
            image_path = f"F:/Chess Piece Images/{piece}.png"
            IMAGES[piece] = pygame.transform.scale(pygame.image.load(image_path), (SQ_SIZE, SQ_SIZE))
        except pygame.error:
            print(f"Error loading image: {piece}.png from {image_path}")

def draw_board(screen):
    colors = [LIGHT_SQ_COLOR, DARK_SQ_COLOR]
    for r in range(8):
        for c in range(8):
            color = colors[(r + c) % 2]
            pygame.draw.rect(screen, color, pygame.Rect(EVAL_BAR_WIDTH + c * SQ_SIZE, r * SQ_SIZE, SQ_SIZE, SQ_SIZE))

def highlight_square(screen, square, color, width=5):
    if square is not None:
        row, col = divmod(square, 8)
        pygame.draw.rect(screen, color, pygame.Rect(EVAL_BAR_WIDTH + col * SQ_SIZE, row * SQ_SIZE, SQ_SIZE, SQ_SIZE), width)

def draw_pieces(screen, board, moves=[], pv_moves=[]):
    piece_mapping = {
        'P': 'wp', 'N': 'wN', 'B': 'wB', 'R': 'wR', 'Q': 'wQ', 'K': 'wK',
        'p': 'bp', 'n': 'bN', 'b': 'bB', 'r': 'bR', 'q': 'bQ', 'k': 'bK'
    }
    for sq in chess.SQUARES:
        piece = board.piece_at(sq)
        if piece:
            piece_image = piece_mapping[piece.symbol()]
            screen.blit(IMAGES[piece_image], (EVAL_BAR_WIDTH + chess.square_file(sq) * SQ_SIZE, chess.square_rank(sq) * SQ_SIZE))

    for move in moves:
        highlight_square(screen, move.from_square, PLAYER_MOVE_COLOR if move in moves else ENGINE_MOVE_COLOR)
        highlight_square(screen, move.to_square, PLAYER_MOVE_COLOR if move in moves else ENGINE_MOVE_COLOR)

    for move in pv_moves:
        highlight_square(screen, move.from_square, TOP_LINE_COLOR, width=2)
        highlight_square(screen, move.to_square, TOP_LINE_COLOR, width=2)

def draw_button(screen):
    pygame.draw.rect(screen, pygame.Color("gray"), pygame.Rect(720, 50, BUTTON_WIDTH, BUTTON_HEIGHT))
    font = pygame.font.Font(None, 36)
    label = font.render("Show PV", True, pygame.Color("white"))
    screen.blit(label, (720 + (BUTTON_WIDTH - label.get_width()) // 2, 50 + (BUTTON_HEIGHT - label.get_height()) // 2))

def handle_button_click(mouse_pos, top_lines):
    if pygame.Rect(720, 50, BUTTON_WIDTH, BUTTON_HEIGHT).collidepoint(mouse_pos):
        return top_lines
    return []

def show_pv_line(screen, board, pv_line):
    for move in pv_line:
        board.push(move)
        draw_board(screen)
        draw_pieces(screen, board)
        pygame.display.flip()
        pygame.time.delay(MOVE_DELAY_TIME // 2)
    for _ in pv_line:
        board.pop()

def main():
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    pygame.display.set_caption("Chess GUI")
    clock = pygame.time.Clock()
    load_images()

    board = chess.Board()
    try:
        engine = chess.engine.SimpleEngine.popen_uci("G:/LcZero/lc0-v0.31.0-windows-cpu-dnnl/lc0")
    except FileNotFoundError:
        print("Error: Chess engine not found.")
        pygame.quit()
        return

    running = True
    player_color = chess.BLACK
    engine_color = chess.WHITE

    selected_square = None
    player_move = None
    top_lines = []
    current_move = None
    is_player_turn = player_color == chess.WHITE
    last_move_time = 0
    pv_to_show = []

    while running:
        screen.fill(pygame.Color("white"))
        current_time = pygame.time.get_ticks()

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

            elif event.type == pygame.MOUSEBUTTONDOWN:
                x, y = event.pos

                if x > EVAL_BAR_WIDTH and x < EVAL_BAR_WIDTH + 8 * SQ_SIZE and y < 8 * SQ_SIZE:
                    col = (x - EVAL_BAR_WIDTH) // SQ_SIZE
                    row = y // SQ_SIZE
                    square = chess.square(col, row)
                    if selected_square is None:
                        selected_square = square
                    else:
                        player_move = chess.Move(selected_square, square)
                        if player_move in board.legal_moves:
                            board.push(player_move)
                            selected_square = None
                            last_move_time = current_time
                            is_player_turn = False
                        else:
                            selected_square = square
                else:
                    pv_to_show = handle_button_click(event.pos, top_lines)

            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_f:
                    player_color = not player_color
                    engine_color = not engine_color

        if not board.is_game_over() and not is_player_turn:
            if current_time - last_move_time >= MOVE_DELAY_TIME:
                info = engine.analyse(board, chess.engine.Limit(time=2.0), multipv=1)
                top_lines = info[0]['pv']
                show_pv_line(screen, board, top_lines)
                current_move = engine.play(board, chess.engine.Limit(time=2.0)).move
                board.push(current_move)
                is_player_turn = True

        draw_board(screen)
        highlight_square(screen, selected_square, HIGHLIGHT_COLOR)
        draw_pieces(screen, board, moves=[player_move] if player_move else [], pv_moves=pv_to_show)
        draw_button(screen)

        pygame.display.flip()
        clock.tick(30)

    engine.quit()
    pygame.quit()

if __name__ == "__main__":
    main()
code working fine except Show PV button but i dont need it just remove that button and  a section where i can  go back and forth of engine PV line of cureent move just like you see buttons in chess.com review section . this section should be updated with current move its PV line. the section should look professional and  simple and should do its job .