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

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

# 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 convert SVG to PNG
def svg_to_png(svg_string):
    png_bytes = svg2png(bytestring=svg_string.encode('utf-8'))
    return Image.open(io.BytesIO(png_bytes))

# Streamlit app
st.set_page_config(layout="wide")
st.title("Chess against LLaMA 405B")

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

# Display the chessboard
svg_board = chess.svg.board(board=st.session_state.board)
st.image(svg_to_png(svg_board), width=400)

# Input for 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)
            
            # Display updated board
            svg_board = chess.svg.board(board=st.session_state.board)
            st.image(svg_to_png(svg_board), width=400)
            
            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.write(f"AI's move: {ai_move}")
                
                # Display updated board after AI move
                svg_board = chess.svg.board(board=st.session_state.board)
                st.image(svg_to_png(svg_board), width=400)
            
            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"Is check? {'Yes' if st.session_state.board.is_check() else 'No'}")
st.write(f"Is game over? {'Yes' if st.session_state.board.is_game_over() else 'No'}")

# Hide the default Streamlit watermark
hide_streamlit_style = """
<style>
#MainMenu {visibility: hidden;}
footer {visibility: hidden;}
</style>
"""
st.markdown(hide_streamlit_style, unsafe_allow_html=True)