import streamlit as st

# Set page config to wide mode
st.set_page_config(layout="wide")

# Initialize session state if not already done
if 'board' not in st.session_state:
    st.session_state.board = [
        ['♜', '♞', '♝', '♛', '♚', '♝', '♞', '♜'],
        ['♟', '♟', '♟', '♟', '♟', '♟', '♟', '♟'],
        [None] * 8,
        [None] * 8,
        [None] * 8,
        [None] * 8,
        ['♙', '♙', '♙', '♙', '♙', '♙', '♙', '♙'],
        ['♖', '♘', '♗', '♕', '♔', '♗', '♘', '♖']
    ]
if 'current_player' not in st.session_state:
    st.session_state.current_player = 'white'
if 'selected_piece' not in st.session_state:
    st.session_state.selected_piece = None

def handle_square_click(row, col):
    if st.session_state.selected_piece is None:
        # Select a piece
        if st.session_state.board[row][col] is not None:
            st.session_state.selected_piece = {'row': row, 'col': col}
            st.rerun()
    else:
        # Move the selected piece
        from_row = st.session_state.selected_piece['row']
        from_col = st.session_state.selected_piece['col']
        
        # Make the move
        st.session_state.board[row][col] = st.session_state.board[from_row][from_col]
        st.session_state.board[from_row][from_col] = None
        st.session_state.selected_piece = None
        st.session_state.current_player = 'black' if st.session_state.current_player == 'white' else 'white'
        st.rerun()

# Create buttons for each square using Streamlit's native components
st.write(f"## Chess Game")
st.write(f"Current Player: {'White' if st.session_state.current_player == 'white' else 'Black'}")

# Create the chessboard using a grid of columns
board_container = st.container()
with board_container:
    # CSS for the chess board
    st.markdown("""
        <style>
        .stButton > button {
            width: 100%;
            height: 0;
            padding-bottom: 100%;
            font-size: 3em;
            display: flex;
            align-items: center;
            justify-content: center;
            background-color: transparent;
            border: none;
        }
        .chess-square-white {
            background-color: #e8e8e8 !important;
        }
        .chess-square-black {
            background-color: #b0b0b0 !important;
        }
        .chess-square-selected {
            background-color: #ffd700 !important;
        }
        </style>
    """, unsafe_allow_html=True)

    # Create the board grid
    for row in range(8):
        cols = st.columns(8)
        for col in range(8):
            with cols[col]:
                piece = st.session_state.board[row][col] or ' '
                is_selected = (st.session_state.selected_piece and 
                             st.session_state.selected_piece['row'] == row and 
                             st.session_state.selected_piece['col'] == col)
                
                square_color = 'white' if (row + col) % 2 == 0 else 'black'
                if is_selected:
                    class_name = "chess-square-selected"
                else:
                    class_name = f"chess-square-{square_color}"
                
                # Create a button for each square
                if st.button(
                    piece,
                    key=f"square_{row}_{col}",
                    help=f"Row {row}, Col {col}",
                    type="secondary",
                    use_container_width=True
                ):
                    handle_square_click(row, col)

# Add a reset button
if st.button('Reset Game', key='reset'):
    st.session_state.board = [
        ['♜', '♞', '♝', '♛', '♚', '♝', '♞', '♜'],
        ['♟', '♟', '♟', '♟', '♟', '♟', '♟', '♟'],
        [None] * 8,
        [None] * 8,
        [None] * 8,
        [None] * 8,
        ['♙', '♙', '♙', '♙', '♙', '♙', '♙', '♙'],
        ['♖', '♘', '♗', '♕', '♔', '♗', '♘', '♖']
    ]
    st.session_state.current_player = 'white'
    st.session_state.selected_piece = None
    st.rerun()