import gradio as gr
from RockPaperScissor.models import RandomAI # Import the AI model

# --- Initialize AI ---
ai_model = RandomAI()

# --- Helper Function ---
def _determine_winner(player_move, ai_move):
    """Determines the winner and returns a result code and message."""
    if player_move == ai_move:
        return "draw", f"It's a draw! (Both chose {player_move})"
    elif (player_move == "rock" and ai_move == "scissors") or \
         (player_move == "paper" and ai_move == "rock") or \
         (player_move == "scissors" and ai_move == "paper"):
        return "player_win", f"You win! ({player_move} beats {ai_move})"
    else:
        return "ai_win", f"AI wins! ({ai_move} beats {player_move})"

# --- Main Gradio Function ---
def play_round(player_choice, current_stats):
    """
    Handles one round of the game.
    Updates stats stored in the Gradio state.
    """
    print(f"Player chose: {player_choice}") # Optional: for debugging in Space logs
    print(f"Current stats before round: {current_stats}")

    # Get AI's move (state is ignored by RandomAI)
    ai_choice, _ = ai_model.make_move()
    print(f"AI chose: {ai_choice}")

    # Determine winner
    result_code, result_message = _determine_winner(player_choice, ai_choice)

    # Update stats (modify the dictionary directly)
    if result_code == "player_win":
        current_stats["wins"] += 1
    elif result_code == "ai_win":
        current_stats["losses"] += 1
    else:
        current_stats["draws"] += 1

    print(f"Current stats after round: {current_stats}")

    # Format output strings
    full_result_message = f"You chose: {player_choice}\nAI chose: {ai_choice}\n\n{result_message}"
    stats_string = f"Wins: {current_stats['wins']} | Losses: {current_stats['losses']} | Draws: {current_stats['draws']}"

    # Return updated components AND the modified state dictionary
    return full_result_message, stats_string, current_stats

# --- Build Gradio Interface ---
with gr.Blocks(title="Rock Paper Scissors Demo") as demo:
    gr.Markdown("# Rock Paper Scissors - Simple Demo")
    gr.Markdown("Play against a Random AI. Choose your move:")

    # State to store session scores
    session_state = gr.State(value={"wins": 0, "losses": 0, "draws": 0})

    with gr.Row():
        rock_btn = gr.Button("🪨 Rock")
        paper_btn = gr.Button("📄 Paper")
        scissors_btn = gr.Button("✂️ Scissors")

    result_output = gr.Textbox(label="Round Result", lines=3) # Use Textbox for multi-line
    stats_output = gr.Label(label="Session Score")

    # --- Connect Buttons to Function ---
    # Use lambda functions to pass the specific choice string
    rock_btn.click(
        fn=play_round,
        inputs=[gr.Textbox("rock", visible=False), session_state], # Pass "rock" implicitly
        outputs=[result_output, stats_output, session_state] # Update outputs AND state
    )
    paper_btn.click(
        fn=play_round,
        inputs=[gr.Textbox("paper", visible=False), session_state], # Pass "paper" implicitly
        outputs=[result_output, stats_output, session_state]
    )
    scissors_btn.click(
        fn=play_round,
        inputs=[gr.Textbox("scissors", visible=False), session_state], # Pass "scissors" implicitly
        outputs=[result_output, stats_output, session_state]
    )

# --- Launch the App ---
demo.launch() # share=True creates a public link when running locally