import streamlit as st
import chess
import chess.svg
import os
import git
import sys
import random

# Configuration de la page
st.set_page_config(page_title="Échecs contre IA", layout="wide")

def render_svg(svg_string):
    """Render un SVG dans Streamlit."""
    from streamlit.components.v1 import html
    html(f"""
        <div style="display: flex; justify-content: center;">
            {svg_string}
        </div>
    """, height=400)

def calculate_max_length(move_count):
    """
    Calcule le max_length optimal basé sur le nombre de coups joués
    - Début de partie: besoin minimum (15 tokens)
    - Milieu de partie: augmentation progressive
    - Fin de partie: plafonnement à 50 tokens
    """
    base_length = 15  # Longueur minimum pour un coup simple
    increment = 1     # Augmentation par coup
    max_length = 100  # Plafond maximum
    
    dynamic_length = base_length + (move_count * increment)
    return min(dynamic_length, max_length)

# Setup du modèle
@st.cache_resource
def setup_inference():
    if not os.path.exists('chess_char'):
        git.Repo.clone_from('https://github.com/l-pommeret/chess_char.git', 'chess_char')
    if 'chess_char' not in sys.path:
        sys.path.append('chess_char')
    from inference import InferenceConfig, ChessGenerator
    
    # Configuration initiale avec max_length minimal
    config = InferenceConfig(
        model_name="Zual/chess_char",
        temperature=0.1,
        max_length=6  # Sera ajusté dynamiquement pendant la partie
    )
    return ChessGenerator(config)

# Initialisation
try:
    generator = setup_inference()
except Exception as e:
    st.error("Erreur lors de l'initialisation.")
    st.stop()

# Initialisation de l'état
if 'board' not in st.session_state:
    st.session_state.board = chess.Board()
    st.session_state.moves = []
    st.session_state.last_move = None

def get_ai_move(prompt):
    """Obtient le coup de l'IA avec max_length dynamique"""
    print(f"\n=== Tour de l'IA ===")
    moves_count = len(st.session_state.moves)
    
    # Mise à jour du max_length en fonction du nombre de coups
    dynamic_max_length = calculate_max_length(moves_count)
    generator.config.max_length = dynamic_max_length
    print(f"max_length actuel: {dynamic_max_length} pour {moves_count} coups joués")
    
    print(f"État actuel de la partie: {prompt}")
    print(f"FEN actuel: {st.session_state.board.fen()}")
    print(f"Coups légaux: {[st.session_state.board.san(move) for move in st.session_state.board.legal_moves]}")
    
    try:
        # On ne génère que jusqu'au prochain coup
        if not prompt:  # Premier coup
            response = generator.generate("1.")
        else:
            # On prend les derniers coups pour ne pas surcharger le contexte
            moves = prompt.split()
            last_moves = " ".join(moves[-4:])  # Garder seulement les 2 derniers coups complets
            if len(moves) % 2 == 0:  # Si on vient de finir un coup noir
                next_move_num = f"{(len(moves)//2 + 1)}."
                response = generator.generate(f"{last_moves} {next_move_num}")
            else:  # Si on vient de finir un coup blanc
                response = generator.generate(f"{last_moves}")
                
        print(f"Réponse brute de l'IA: {response}")
        
        # Gestion de la réponse quelle que soit sa forme
        moves = response[0].split() if isinstance(response, list) else response.split()
        print(f"Coups extraits: {moves}")
        
        # On prend toujours le dernier coup généré
        next_move = moves[-1]
        if '.' in next_move and len(moves) > 1:
            next_move = moves[-2]
            
        print(f"Coup candidat de l'IA: {next_move}")
        
        # Vérification de la validité
        try:
            move = st.session_state.board.parse_san(next_move)
            print(f"Coup parsé en UCI: {move}")
            if move in st.session_state.board.legal_moves:
                print(f"Coup valide trouvé: {next_move}")
                return next_move
            else:
                print(f"Coup non légal: {next_move}")
        except ValueError as e:
            print(f"Erreur de parsing: {e}")
        
        # En cas d'échec, jouer un coup légal aléatoire
        legal_moves = list(st.session_state.board.legal_moves)
        if legal_moves:
            random_move = random.choice(legal_moves)
            random_san = st.session_state.board.san(random_move)
            print(f"Utilisation d'un coup aléatoire: {random_san}")
            return random_san
            
    except Exception as e:
        print(f"Erreur dans get_ai_move: {e}")
        return None

def try_move(move_str):
    """Applique un coup au plateau"""
    print(f"\n=== Tentative de coup: {move_str} ===")
    print(f"État avant le coup: {get_game_string()}")
    print(f"FEN avant: {st.session_state.board.fen()}")
    
    try:
        # Nettoyer le coup (enlever le numéro si présent)
        clean_move = move_str.split('.')[-1].strip()
        print(f"Coup nettoyé: {clean_move}")
        
        move = st.session_state.board.parse_san(clean_move)
        print(f"Coup parsé en UCI: {move}")
        
        if move in st.session_state.board.legal_moves:
            st.session_state.board.push(move)
            st.session_state.moves.append(clean_move)
            st.session_state.last_move = clean_move
            print(f"Coup appliqué avec succès")
            print(f"Nouvel état: {get_game_string()}")
            print(f"Nouveau FEN: {st.session_state.board.fen()}")
            return True
            
        print(f"Coup illégal")
        return False
    except ValueError as e:
        print(f"Erreur de parsing: {e}")
        return False

def get_game_string():
    """Renvoie la notation de la partie"""
    result = []
    for i in range(0, len(st.session_state.moves), 2):
        move_num = i//2 + 1
        result.append(f"{move_num}.{st.session_state.moves[i]}")
        if i+1 < len(st.session_state.moves):
            result.append(st.session_state.moves[i+1])
    return " ".join(result)

# Interface utilisateur
st.title("♟️ Échecs contre IA")

col1, col2 = st.columns([2, 1])

with col1:
    # Plateau avec la nouvelle méthode d'affichage
    board_svg = chess.svg.board(
        board=st.session_state.board,
        lastmove=st.session_state.board.peek() if st.session_state.board.move_stack else None,
        size=400  # Taille fixe pour l'échiquier
    )
    render_svg(board_svg)
    
    # Input du joueur avec key dynamique
    move = st.text_input(
        "Votre coup",
        key=f"move_input_{len(st.session_state.moves)}",
        placeholder="ex: e4, Nf3, O-O"
    )

with col2:
    # État actuel
    print(f"\n=== État actuel ===")
    print(f"Partie en cours: {get_game_string()}")
    print(f"FEN: {st.session_state.board.fen()}")
    
    # Affichage du max_length actuel
    current_max_length = calculate_max_length(len(st.session_state.moves))
    st.info(f"Longueur de génération actuelle: {current_max_length} tokens")
    
    # Historique
    st.subheader("Partie en cours")
    game_str = get_game_string()
    if game_str:
        st.text_area("Historique", value=game_str, height=100, label_visibility="collapsed")
    
    # Instructions
    st.markdown("""
    **Comment jouer:**
    - Pion: e4, d5
    - Cavalier: Nf3, Nc6
    - Fou: Bc4, Be7
    - Tour: Ra1, Rd8
    - Dame: Qd1, Qh4
    - Roi: Ke2, Kg8
    - Roque: O-O ou O-O-O
    """)
    
    # Nouvelle partie
    if st.button("Nouvelle partie"):
        st.session_state.board = chess.Board()
        st.session_state.moves = []
        st.session_state.last_move = None
        st.rerun()

# Logique du jeu
if move:
    if try_move(move):
        game_str = get_game_string()
        
        # Tour de l'IA
        with st.spinner("L'IA réfléchit..."):
            ai_move = get_ai_move(game_str)
            if ai_move and try_move(ai_move):
                if st.session_state.board.is_checkmate():
                    st.success("Échec et mat!")
                elif st.session_state.board.is_game_over():
                    st.info("Partie terminée!")
            else:
                st.error("Problème avec le coup de l'IA")
        st.rerun()
    else:
        st.error("Coup invalide")

# État du jeu
if st.session_state.board.is_check():
    st.warning("⚠️ Échec et mat !")