import os
import gradio as gr
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
import numpy as np

model = Llama(
    model_path=hf_hub_download(
        repo_id=os.environ.get("REPO_ID", "Lyte/QuadConnect2.5-0.5B-v0.0.4b"),
        filename=os.environ.get("MODEL_FILE", "unsloth.Q8_0.gguf"),
    ),
    n_ctx=16384
)

SYSTEM_PROMPT = """You are an expert Connect Four player. The game is played on a 6x7 grid where pieces fall to the lowest available position in each column due to gravity.

Board representation:
- The board is described as a list of occupied positions in the format: <column><row>(<piece>).
- Columns are labeled a-g (from left to right) and rows are numbered 1-6 (with 1 as the bottom row).
- For example: 'a1(O), a2(X), b1(O)' indicates that cell a1 has an O, a2 has an X, and b1 has an O.
- An empty board is simply represented by an empty list.
- To win, you must connect 4 of your pieces horizontally, vertically, or diagonally.

Respond in the following XML format:
<reasoning>
Explain your reasoning, including:
- Identifying winning opportunities for yourself.
- Blocking your opponent's potential wins.
- Strategic positioning, such as center control and setting up future moves.
</reasoning>
<move>
Indicate the column letter (a-g) where you want to drop your piece.
</move>
"""

class ConnectFour:
    def __init__(self):
        self.board = np.zeros((6, 7))
        self.current_player = 1  # 1 for player (X), 2 for AI (O)
        self.game_over = False
        
    def make_move(self, col):
        if self.game_over:
            return False, -1
            
        # Find the lowest empty row in the selected column
        for row in range(5, -1, -1):
            if self.board[row][col] == 0:
                self.board[row][col] = self.current_player
                return True, row
        return False, -1

    def check_winner(self):
        # Check horizontal
        for row in range(6):
            for col in range(4):
                if (self.board[row][col] != 0 and
                    self.board[row][col] == self.board[row][col+1] == 
                    self.board[row][col+2] == self.board[row][col+3]):
                    return self.board[row][col]
                    
        # Check vertical
        for row in range(3):
            for col in range(7):
                if (self.board[row][col] != 0 and
                    self.board[row][col] == self.board[row+1][col] ==
                    self.board[row+2][col] == self.board[row+3][col]):
                    return self.board[row][col]
                    
        # Check diagonal (positive slope)
        for row in range(3):
            for col in range(4):
                if (self.board[row][col] != 0 and
                    self.board[row][col] == self.board[row+1][col+1] ==
                    self.board[row+2][col+2] == self.board[row+3][col+3]):
                    return self.board[row][col]
                    
        # Check diagonal (negative slope)
        for row in range(3, 6):
            for col in range(4):
                if (self.board[row][col] != 0 and
                    self.board[row][col] == self.board[row-1][col+1] ==
                    self.board[row-2][col+2] == self.board[row-3][col+3]):
                    return self.board[row][col]
                    
        return 0

    def board_to_string(self):
        moves = []
        for row in range(6):
            for col in range(7):
                if self.board[row][col] != 0:
                    col_letter = chr(ord('a') + col)
                    row_num = str(6 - row)  # Convert to 1-based indexing
                    piece = "X" if self.board[row][col] == 1 else "O"
                    moves.append(f"{col_letter}{row_num}({piece})")
        return ", ".join(moves) if moves else ""

    def parse_ai_move(self, move_str):
        # Parse move like 'a', 'b', etc.
        try:
            col = ord(move_str.strip().lower()) - ord('a')
            if 0 <= col <= 6:
                return col
            return -1
        except:
            return -1

def create_interface():
    game = ConnectFour()
    
    css = """
    .connect4-board {
        display: grid;
        grid-template-columns: repeat(7, 1fr);
        gap: 8px;
        max-width: 600px;
        margin: 10px auto;
        background: #2196F3;
        padding: 15px;
        border-radius: 15px;
        box-shadow: 0 4px 8px rgba(0,0,0,0.2);
    }
    .connect4-cell {
        aspect-ratio: 1;
        background: white;
        border-radius: 50%;
        display: flex;
        align-items: center;
        justify-content: center;
        font-size: 2em;
    }
    .player1 { background: #f44336 !important; }
    .player2 { background: #ffc107 !important; }
    #ai-status {
        font-size: 1.2em;
        margin: 10px 0;
        color: #2196F3;
        font-weight: bold;
    }
    #ai-reasoning {
        background: #22004d;
        border-radius: 10px;
        padding: 15px;
        margin: 15px 0;
        font-family: monospace;
        min-height: 100px;
    }
    .reasoning-box {
        border-left: 4px solid #2196F3;
        padding-left: 15px;
        margin: 10px 0;
        background: #22004d;
        border-radius: 0 10px 10px 0;
    }
    #column-buttons {
        display: flex;
        justify-content: center;
        align-items: anchor-center;
        max-width: 600px;
        margin: 0 auto;
        padding: 0 15px;
    }
    #column-buttons button {
      margin: 0px 5px;
    }
    div.svelte-1nguped {
      display: block;
    }
    """

    with gr.Blocks(css=css) as interface:
        gr.Markdown("# 🎮 Connect Four vs AI")
        gr.Markdown("### Play against an AI trained to be an expert Connect Four player!")
        
        with gr.Row():
            with gr.Column(scale=2):
                # Status display
                status = gr.Markdown("Your turn! Click a button to drop your piece!", elem_id="ai-status")
                
                # Column buttons
                with gr.Group(elem_id="column-buttons"):
                    col_buttons = []
                    for i in range(7):
                        btn = gr.Button(f"⬇️ {chr(ord('A') + i)}", scale=1)
                        col_buttons.append(btn)
                
                # Game board
                board_display = gr.HTML(render_board(), elem_id="board-display")
                reset_btn = gr.Button("🔄 New Game", variant="primary")
            
            with gr.Column(scale=1):
                # AI reasoning display
                gr.Markdown("### 🤖 AI's Thoughts")
                reasoning_display = gr.HTML(
                    value='<div id="ai-reasoning">Waiting for your move...</div>',
                    elem_id="ai-reasoning-container"
                )

        def handle_move(col):
            if game.game_over:
                return [
                    render_board(game.board),
                    "Game is over! Click New Game to play again.",
                    '<div id="ai-reasoning">Game Over!</div>'
                ]
            
            # Player move
            success, row = game.make_move(col)
            if not success:
                return [
                    render_board(game.board),
                    "Column is full! Try another one.",
                    '<div id="ai-reasoning">Invalid move!</div>'
                ]
            
            # Check for winner
            winner = game.check_winner()
            if winner == 1:
                game.game_over = True
                return [
                    render_board(game.board),
                    "🎉 You win! 🎉",
                    '<div id="ai-reasoning">Congratulations! You won!</div>'
                ]
            
            # AI move
            game.current_player = 2
            board_state = game.board_to_string()
            
            prompt = f"Current board state (you are O, opponent is X):\n{board_state}\n\nMake your move."
            
            # Get AI response
            response = model.create_chat_completion(
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.8,
                top_p=0.95,
                max_tokens=2048
            )
            
            ai_response = response['choices'][0]['message']['content']
            print(ai_response)
            
            # Extract reasoning and move
            try:
                reasoning = ai_response.split("<reasoning>")[1].split("</reasoning>")[0].strip()
                move_str = ai_response.split("<move>")[1].split("</move>")[0].strip()
                ai_col = game.parse_ai_move(move_str)
                
                if ai_col == -1:
                    raise ValueError("Invalid move format from AI")
                
                # Format reasoning for display
                reasoning_html = f'''
                <div id="ai-reasoning">
                    <div class="reasoning-box">
                        <p><strong>🤔 Reasoning:</strong></p>
                        <p>{reasoning}</p>
                        <p><strong>📍 Move chosen:</strong> Column {move_str.upper()}</p>
                    </div>
                </div>
                '''
                
                success, _ = game.make_move(ai_col)
                if success:
                    # Check for AI winner
                    winner = game.check_winner()
                    if winner == 2:
                        game.game_over = True
                        return [
                            render_board(game.board),
                            "🤖 AI wins! Better luck next time!",
                            reasoning_html
                        ]
                else:
                    return [
                        render_board(game.board),
                        "AI made invalid move! You win by default!",
                        '<div id="ai-reasoning">AI made an invalid move!</div>'
                    ]
            except Exception as e:
                game.game_over = True
                return [
                    render_board(game.board),
                    "AI error occurred! You win by default!",
                    f'<div id="ai-reasoning">Error: {str(e)}</div>'
                ]
            
            game.current_player = 1
            return [render_board(game.board), "Your turn!", reasoning_html]

        def reset_game():
            game.board = np.zeros((6, 7))
            game.current_player = 1
            game.game_over = False
            return [
                render_board(),
                "Your turn! Click a button to drop your piece!",
                '<div id="ai-reasoning">New game started! Make your move...</div>'
            ]

        # Event handlers
        for i, btn in enumerate(col_buttons):
            btn.click(
                fn=handle_move,
                inputs=[gr.Number(value=i, visible=False)],
                outputs=[board_display, status, reasoning_display]
            )
        
        reset_btn.click(
            fn=reset_game,
            outputs=[board_display, status, reasoning_display]
        )
        
    return interface

def render_board(board=None):
    if board is None:
        board = np.zeros((6, 7))
        
    html = '<div class="connect4-board">'
    
    for row in range(6):
        for col in range(7):
            cell_class = "connect4-cell"
            content = "⚪"
            
            if board[row][col] == 1:
                cell_class += " player1"
                content = "🔴"
            elif board[row][col] == 2:
                cell_class += " player2"
                content = "🟡"
                
            html += f'<div class="{cell_class}">{content}</div>'
    
    html += "</div>"
    return html

interface = create_interface()
interface.launch()