import streamlit as st
import math

# --- Initialize session state ---
if "board" not in st.session_state:
    st.session_state.board = [' ']*9
if "game_over" not in st.session_state:
    st.session_state.game_over = False
if "winner" not in st.session_state:
    st.session_state.winner = None

# --- Game Logic ---
def check_winner(player):
    win_combos = [
        [0,1,2], [3,4,5], [6,7,8],
        [0,3,6], [1,4,7], [2,5,8],
        [0,4,8], [2,4,6]
    ]
    for combo in win_combos:
        if all(st.session_state.board[i] == player for i in combo):
            return True
    return False

def is_full():
    return ' ' not in st.session_state.board

def minimax(is_maximizing, alpha, beta):
    if check_winner('O'):
        return 1
    if check_winner('X'):
        return -1
    if is_full():
        return 0

    if is_maximizing:
        max_eval = -math.inf
        for i in range(9):
            if st.session_state.board[i] == ' ':
                st.session_state.board[i] = 'O'
                eval_score = minimax(False, alpha, beta)
                st.session_state.board[i] = ' '
                max_eval = max(max_eval, eval_score)
                alpha = max(alpha, eval_score)
                if beta <= alpha:
                    break
        return max_eval
    else:
        min_eval = math.inf
        for i in range(9):
            if st.session_state.board[i] == ' ':
                st.session_state.board[i] = 'X'
                eval_score = minimax(True, alpha, beta)
                st.session_state.board[i] = ' '
                min_eval = min(min_eval, eval_score)
                beta = min(beta, eval_score)
                if beta <= alpha:
                    break
        return min_eval

def ai_move():
    best_score = -math.inf
    move = None
    for i in range(9):
        if st.session_state.board[i] == ' ':
            st.session_state.board[i] = 'O'
            score = minimax(False, -math.inf, math.inf)
            st.session_state.board[i] = ' '
            if score > best_score:
                best_score = score
                move = i
    st.session_state.board[move] = 'O'

# --- Reset Game ---
def reset_game():
    st.session_state.board = [' ']*9
    st.session_state.game_over = False
    st.session_state.winner = None

# --- Handle Player Move ---
def handle_click(i):
    if st.session_state.board[i] == ' ' and not st.session_state.game_over:
        st.session_state.board[i] = 'X'
        if check_winner('X'):
            st.session_state.game_over = True
            st.session_state.winner = "You"
            return
        if is_full():
            st.session_state.game_over = True
            st.session_state.winner = "Draw"
            return

        ai_move()
        if check_winner('O'):
            st.session_state.game_over = True
            st.session_state.winner = "AI"
        elif is_full():
            st.session_state.game_over = True
            st.session_state.winner = "Draw"

# --- UI Layout ---
st.title("🤖 Tic-Tac-Toe with AI")
st.write("You are X. The AI is O. Click on a box to make your move.")

# Display Board as 3x3 buttons
cols = st.columns(3)
for row in range(3):
    for col in range(3):
        idx = row*3 + col
        with cols[col]:
            if st.button(st.session_state.board[idx] if st.session_state.board[idx] != ' ' else ' ', key=idx):
                handle_click(idx)

# Game Status
if st.session_state.game_over:
    if st.session_state.winner == "Draw":
        st.subheader("🤝 It's a Draw!")
    else:
        st.subheader(f"🏆 {st.session_state.winner} Wins!")

# Reset Button
if st.button("🔄 Reset Game"):
    reset_game()