import random
import streamlit as st

# Candy types
CANDY_TYPES = ["🍬", "🍫", "🍭", "🍪", "🍓"]

# Initialize the game grid
def initialize_grid(rows, cols):
    return [[random.choice(CANDY_TYPES) for _ in range(cols)] for _ in range(rows)]

# Display the grid as buttons
def display_grid(grid, selected):
    for r, row in enumerate(grid):
        cols = st.columns(len(row))
        for c, candy in enumerate(row):
            # Use row and column indices to generate a unique key
            button_key = f"btn_{r}_{c}"
            if (r, c) == selected:
                button_label = f"**{candy}**"  # Highlight selected candy
            else:
                button_label = candy
            if cols[c].button(button_label, key=button_key):
                return (r, c)  # Return the coordinates of the clicked candy
    return None

# Check for matches (horizontal and vertical)
def check_matches(grid):
    rows, cols = len(grid), len(grid[0])
    matches = set()

    # Check horizontal matches
    for r in range(rows):
        for c in range(cols - 2):
            if grid[r][c] == grid[r][c + 1] == grid[r][c + 2]:
                matches.update([(r, c), (r, c + 1), (r, c + 2)])

    # Check vertical matches
    for r in range(rows - 2):
        for c in range(cols):
            if grid[r][c] == grid[r + 1][c] == grid[r + 2][c]:
                matches.update([(r, c), (r + 1, c), (r + 2, c)])

    return matches

# Remove matched candies and drop new candies
def remove_matches(grid, matches):
    rows, cols = len(grid), len(grid[0])

    # Set matched candies to None
    for r, c in matches:
        grid[r][c] = None

    # Drop candies down and refill the grid
    for c in range(cols):
        column = [grid[r][c] for r in range(rows) if grid[r][c] is not None]
        while len(column) < rows:
            column.insert(0, random.choice(CANDY_TYPES))
        for r in range(rows):
            grid[r][c] = column[r]

    return grid

# Swap candies in the grid
def swap_candies(grid, first, second):
    if first and second:
        r1, c1 = first
        r2, c2 = second
        if (
            0 <= r1 < len(grid)
            and 0 <= c1 < len(grid[0])
            and 0 <= r2 < len(grid)
            and 0 <= c2 < len(grid[0])
        ):
            grid[r1][c1], grid[r2][c2] = grid[r2][c2], grid[r1][c1]
            return True
    return False

# Main Streamlit app
def main():
    st.title("Candy Crush Game")
    st.write("Match 3 or more candies to crush them! 🍬🍫🍭")

    rows, cols = 6, 6  # Grid size

    # Initialize or reset the grid
    if "grid" not in st.session_state:
        st.session_state.grid = initialize_grid(rows, cols)
        st.session_state.selected_candy = None

    # Display the current grid as buttons
    st.write("### Current Grid")
    selected_candy = display_grid(st.session_state.grid, st.session_state.selected_candy)

    # Handle candy selection
    if selected_candy:
        if st.session_state.selected_candy is None:
            st.session_state.selected_candy = selected_candy
            st.info(f"Selected candy at {selected_candy}")
        else:
            # Attempt to swap the two candies
            if swap_candies(st.session_state.grid, st.session_state.selected_candy, selected_candy):
                matches = check_matches(st.session_state.grid)
                if matches:
                    st.success("Valid move! Candies matched!")
                    st.session_state.grid = remove_matches(st.session_state.grid, matches)
                else:
                    st.warning("No match found. Swap reverted!")
                    swap_candies(st.session_state.grid, st.session_state.selected_candy, selected_candy)  # Undo the swap
            else:
                st.error("Invalid swap!")
            st.session_state.selected_candy = None  # Reset selection

    # Reset the game
    if st.button("Reset Game"):
        st.session_state.grid = initialize_grid(rows, cols)
        st.session_state.selected_candy = None
        st.write("Game reset!")

if __name__ == "__main__":
    main()