import streamlit as st
import chess
import chess.svg
from gradio_client import Client
import cairosvg
from PIL import Image
import io

# Initialize the Gradio client
client = Client("xianbao/SambaNova-fast")

# Initialize the chess board
board = chess.Board()

# Function to get AI move
def get_ai_move(board):
    fen = board.fen()
    prompt = f"You are playing as Black in a chess game. The current board position in FEN notation is: {fen}. What is your next move? Please respond with the move in UCI notation (e.g., 'e2e4')."
    
    result = client.predict(
        message=prompt,
        system_message="You are a chess engine based on LLaMA 405B. Provide only the next move in UCI notation.",
        max_tokens=1024,
        temperature=0.6,
        top_p=0.9,
        top_k=50,
        api_name="/chat"
    )
    
    return result.strip()

# Function to render the chess board
def render_board(board):
    svg = chess.svg.board(board=board)
    png = cairosvg.svg2png(bytestring=svg.encode('utf-8'))
    return Image.open(io.BytesIO(png))

# Streamlit app
st.title("Chess against LLaMA 405B")

# Initialize session state
if 'board' not in st.session_state:
    st.session_state.board = chess.Board()

# Display the current board state
st.image(render_board(st.session_state.board))

# Get user's move
user_move = st.text_input("Enter your move (e.g., 'e2e4'):")

if st.button("Make Move"):
    try:
        move = chess.Move.from_uci(user_move)
        if move in st.session_state.board.legal_moves:
            st.session_state.board.push(move)
            st.image(render_board(st.session_state.board))
            
            if not st.session_state.board.is_game_over():
                with st.spinner("AI is thinking..."):
                    ai_move = get_ai_move(st.session_state.board)
                ai_move_obj = chess.Move.from_uci(ai_move)
                st.session_state.board.push(ai_move_obj)
                st.image(render_board(st.session_state.board))
                st.write(f"AI's move: {ai_move}")
            
            if st.session_state.board.is_game_over():
                st.write("Game Over!")
                st.write(f"Result: {st.session_state.board.result()}")
        else:
            st.write("Invalid move. Please try again.")
    except ValueError:
        st.write("Invalid input. Please enter a move in UCI notation (e.g., 'e2e4').")

# Reset button
if st.button("Reset Game"):
    st.session_state.board = chess.Board()
    st.experimental_rerun()

# Display game status
st.write(f"Current turn: {'White' if st.session_state.board.turn else 'Black'}")
st.write(f"Fullmove number: {st.session_state.board.fullmove_number}")
st.write(f"Halfmove clock: {st.session_state.board.halfmove_clock}")
st.write(f"Is check? {st.session_state.board.is_check()}")
st.write(f"Is checkmate? {st.session_state.board.is_checkmate()}")
st.write(f"Is stalemate? {st.session_state.board.is_stalemate()}")