//! Python bindings for the Essence Wars game engine.
//!
//! This module provides PyO3 bindings to expose the Rust game engine to Python.
//! It enables high-performance training with frameworks like PyTorch while
//! keeping all game logic in Rust.
//!
//! # Usage from Python
//!
//! ```python
//! from essence_wars._core import PyGame, PyParallelGames
//!
//! # Single game
//! game = PyGame()
//! game.reset(seed=42)
//! obs = game.observe()  # numpy array (328,)
//! mask = game.action_mask()  # numpy array (422,)
//! reward, done = game.step(action)
//!
//! # Batched games for vectorized training
//! games = PyParallelGames(num_envs=64)
//! games.reset([seed1, seed2, ...])
//! obs_batch = games.observe_batch()  # numpy array (64, 328)
//! ```

// Allow clippy::useless_conversion for PyO3 method returns
// PyO3 macros require PyResult even though clippy 1.92.0 sees it as redundant
#![allow(clippy::useless_conversion)]

use pyo3::prelude::*;
use pyo3::exceptions::{PyValueError, PyRuntimeError};
use numpy::{PyArray1, PyArray2, PyReadonlyArray1, PyArrayMethods};

/// Type alias for the complex return type used in batched step operations
type BatchStepResult<'py> = (Bound<'py, PyArray1<f32>>, Bound<'py, PyArray1<bool>>);

use crate::core::actions::Action;
use crate::core::cards::CardDatabase;
use crate::core::config::tensor::STATE_TENSOR_SIZE;
use crate::core::engine::GameEngine;
use crate::core::state::GameMode;
use crate::core::types::PlayerId;
use crate::decks::{DeckDefinition, DeckRegistry};
use crate::bots::{AlphaBetaBot, AlphaBetaConfig, Bot, GreedyBot, MctsBot, MctsConfig, RandomBot};

/// Opponent type for vectorized training environments.
#[derive(Clone, Debug)]
enum OpponentType {
    /// No built-in opponent — agent plays both sides (self-play)
    SelfPlay,
    /// Built-in GreedyBot opponent for Player 2
    Greedy,
    /// Built-in RandomBot opponent for Player 2
    Random,
    /// MCTS opponent with configurable simulation count
    Mcts { simulations: u32 },
    /// Alpha-Beta opponent with configurable search depth
    AlphaBeta { depth: u32 },
}

/// Single game wrapper for Python.
///
/// Provides a Pythonic interface to the Rust game engine with zero-copy
/// tensor transfer via NumPy arrays.
#[pyclass]
pub struct PyGame {
    engine: GameEngine<'static>,
    #[allow(dead_code)]
    card_db: &'static CardDatabase,
    #[allow(dead_code)]
    deck_registry: &'static DeckRegistry,
    deck1: DeckDefinition,
    deck2: DeckDefinition,
    game_mode: GameMode,
}

// Static storage for card database and deck registry
// These are loaded once from embedded data and shared across all PyGame instances
lazy_static::lazy_static! {
    static ref CARD_DB: CardDatabase = {
        crate::embedded_data::load_embedded_cards_with_commanders()
            .expect("Failed to load embedded card database with commanders")
    };
    static ref DECK_REGISTRY: DeckRegistry = {
        crate::embedded_data::load_embedded_decks()
            .expect("Failed to load embedded deck registry")
    };
}

#[pymethods]
impl PyGame {
    /// Create a new game instance.
    ///
    /// Args:
    ///     deck1: Name of deck for player 1 (default: "architect_fortify")
    ///     deck2: Name of deck for player 2 (default: "broodmother_pack")
    ///     game_mode: "attrition" (default) or "essence_duel"
    #[new]
    #[pyo3(signature = (deck1=None, deck2=None, game_mode=None))]
    fn new(
        deck1: Option<&str>,
        deck2: Option<&str>,
        game_mode: Option<&str>,
    ) -> PyResult<Self> {
        let card_db: &'static CardDatabase = &CARD_DB;
        let deck_registry: &'static DeckRegistry = &DECK_REGISTRY;

        // Load decks
        let deck1_name = deck1.unwrap_or("artificer_tokens");
        let deck2_name = deck2.unwrap_or("broodmother_pack");

        let deck1_def = deck_registry.get(deck1_name)
            .ok_or_else(|| PyValueError::new_err(format!("Unknown deck: {}", deck1_name)))?;
        let deck2_def = deck_registry.get(deck2_name)
            .ok_or_else(|| PyValueError::new_err(format!("Unknown deck: {}", deck2_name)))?;

        // Parse game mode (essence_war is the new default, but accept legacy names for compatibility)
        let mode = match game_mode.unwrap_or("essence_war") {
            "attrition" => GameMode::Attrition,
            "essence_war" | "essence-war" | "essence_duel" | "essence-duel" => GameMode::EssenceWar,
            other => return Err(PyValueError::new_err(format!("Unknown game mode: {}", other))),
        };

        let engine = GameEngine::new(card_db);

        Ok(Self {
            engine,
            card_db,
            deck_registry,
            deck1: deck1_def.clone(),
            deck2: deck2_def.clone(),
            game_mode: mode,
        })
    }

    /// Reset the game to initial state.
    ///
    /// Args:
    ///     seed: Random seed for deck shuffling and game randomness
    fn reset(&mut self, seed: u64) {
        self.engine
            .start_game_with_mode(&self.deck1, &self.deck2, seed, self.game_mode)
            .expect("Failed to start game - commander not found in card database");
    }

    /// Get the current state tensor as a NumPy array.
    ///
    /// Returns:
    ///     numpy.ndarray: Float32 array of shape (328,) representing game state
    fn observe<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<f32>> {
        let tensor = self.engine.get_state_tensor();
        PyArray1::from_slice_bound(py, &tensor)
    }

    /// Get the legal action mask as a NumPy array.
    ///
    /// Returns:
    ///     numpy.ndarray: Float32 array of shape (422,) where 1.0 = legal, 0.0 = illegal
    fn action_mask<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<f32>> {
        let mask = self.engine.get_legal_action_mask();
        PyArray1::from_slice_bound(py, &mask)
    }

    /// Apply an action and return (reward, done).
    ///
    /// Args:
    ///     action: Action index (0-421)
    ///
    /// Returns:
    ///     tuple: (reward: float, done: bool)
    ///         - reward is from player 0's perspective: +1 win, -1 loss, 0 otherwise
    ///         - done is True if the game has ended
    fn step(&mut self, action: u16) -> PyResult<(f32, bool)> {
        self.engine.apply_action_by_index(action)
            .map_err(PyValueError::new_err)?;

        let done = self.engine.is_game_over();
        let reward = if done {
            self.engine.get_reward(PlayerId::PLAYER_ONE)
        } else {
            0.0
        };

        Ok((reward, done))
    }

    /// Get the current player (0 or 1).
    fn current_player(&self) -> u8 {
        self.engine.current_player().0
    }

    /// Check if the game is over.
    fn is_done(&self) -> bool {
        self.engine.is_game_over()
    }

    /// Get the reward for a specific player.
    ///
    /// Args:
    ///     player: Player index (0 or 1)
    ///
    /// Returns:
    ///     float: +1.0 for win, -1.0 for loss, 0.0 for ongoing/draw
    fn get_reward(&self, player: u8) -> f32 {
        self.engine.get_reward(PlayerId(player))
    }

    /// Clone the game state for tree search (e.g., MCTS).
    ///
    /// Returns:
    ///     PyGame: A new PyGame instance with copied state
    fn fork(&self) -> Self {
        Self {
            engine: self.engine.fork(),
            card_db: self.card_db,
            deck_registry: self.deck_registry,
            deck1: self.deck1.clone(),
            deck2: self.deck2.clone(),
            game_mode: self.game_mode,
        }
    }

    /// Get action from the built-in GreedyBot.
    ///
    /// Useful as a baseline opponent during training.
    ///
    /// Returns:
    ///     int: Action index chosen by GreedyBot
    fn greedy_action(&self) -> PyResult<u16> {
        let tensor = self.engine.get_state_tensor();
        let mask = self.engine.get_legal_action_mask();
        let legal_actions = self.engine.get_legal_actions();

        if legal_actions.is_empty() {
            return Err(PyRuntimeError::new_err("No legal actions available"));
        }

        let mut bot = GreedyBot::new(self.card_db, 42);
        let action = bot.select_action(&tensor, &mask, &legal_actions);

        Ok(action.to_index())
    }

    /// Get action from the built-in RandomBot.
    ///
    /// Returns:
    ///     int: Action index chosen uniformly at random from legal actions
    fn random_action(&self) -> PyResult<u16> {
        let tensor = self.engine.get_state_tensor();
        let mask = self.engine.get_legal_action_mask();
        let legal_actions = self.engine.get_legal_actions();

        if legal_actions.is_empty() {
            return Err(PyRuntimeError::new_err("No legal actions available"));
        }

        let mut bot = RandomBot::new(42);
        let action = bot.select_action(&tensor, &mask, &legal_actions);

        Ok(action.to_index())
    }

    /// Get action from the built-in MCTS bot.
    ///
    /// MCTS (Monte Carlo Tree Search) provides a strong baseline that
    /// uses tree search with UCB1 selection and GreedyBot rollouts.
    ///
    /// Args:
    ///     simulations: Number of MCTS simulations (default: 100)
    ///
    /// Returns:
    ///     int: Action index chosen by MCTS
    #[pyo3(signature = (simulations=100))]
    fn mcts_action(&self, simulations: u32) -> PyResult<u16> {
        let legal_actions = self.engine.get_legal_actions();

        if legal_actions.is_empty() {
            return Err(PyRuntimeError::new_err("No legal actions available"));
        }

        let config = MctsConfig {
            simulations,
            ..Default::default()
        };
        let mut bot = MctsBot::with_config(self.card_db, config, 42);

        // MCTS needs engine access for simulation
        let action = bot.select_action_with_engine(&self.engine);

        Ok(action.to_index())
    }

    /// Get action from the built-in Alpha-Beta bot.
    ///
    /// Alpha-Beta is a highly optimized minimax search with:
    /// - Transposition table for position caching
    /// - Killer moves and history heuristic for move ordering
    /// - Aspiration windows for faster iterative deepening
    /// - Late Move Reduction (LMR) for deeper effective search
    ///
    /// This bot is 18-30x faster than MCTS at comparable strength.
    ///
    /// Args:
    ///     depth: Search depth in plies (default: 6, recommended: 6-8)
    ///
    /// Returns:
    ///     int: Action index chosen by Alpha-Beta search
    #[pyo3(signature = (depth=6))]
    fn alphabeta_action(&self, depth: u32) -> PyResult<u16> {
        let legal_actions = self.engine.get_legal_actions();

        if legal_actions.is_empty() {
            return Err(PyRuntimeError::new_err("No legal actions available"));
        }

        let config = AlphaBetaConfig::with_depth(depth);
        let mut bot = AlphaBetaBot::with_config(self.card_db, config, 42);

        // Alpha-Beta needs engine access for tree search
        let action = bot.select_action_with_engine(&self.engine);

        Ok(action.to_index())
    }

    /// Get all legal actions ranked by Alpha-Beta search score.
    ///
    /// Returns every legal action with its minimax evaluation, sorted
    /// best-first. Useful for generating reasoning traces / chain-of-thought
    /// training data where the model learns to evaluate candidate moves.
    ///
    /// Args:
    ///     depth: Search depth in plies (default: 6, recommended: 3-6)
    ///
    /// Returns:
    ///     list[tuple[int, float]]: List of (action_index, score) pairs,
    ///         sorted by score descending (best move first).
    #[pyo3(signature = (depth=6))]
    fn alphabeta_ranked_actions(&self, depth: u32) -> PyResult<Vec<(u16, f32)>> {
        let legal_actions = self.engine.get_legal_actions();

        if legal_actions.is_empty() {
            return Ok(vec![]);
        }

        let config = AlphaBetaConfig::with_depth(depth);
        let mut bot = AlphaBetaBot::with_config(self.card_db, config, 42);

        let ranked = bot.search_ranked(&self.engine);

        Ok(ranked.iter().map(|(action, score)| (action.to_index(), *score)).collect())
    }

    /// Get the current turn number.
    fn turn_number(&self) -> u16 {
        self.engine.turn_number()
    }

    /// List all available decks.
    #[staticmethod]
    fn list_decks() -> Vec<String> {
        DECK_REGISTRY.deck_ids().into_iter().map(|s| s.to_string()).collect()
    }
}

/// Batched game wrapper for vectorized training.
///
/// Runs multiple games in parallel using Rayon for high throughput.
/// Achieves 60k+ steps per second on modern CPUs.
#[pyclass]
pub struct PyParallelGames {
    engines: Vec<GameEngine<'static>>,
    #[allow(dead_code)]
    card_db: &'static CardDatabase,
    #[allow(dead_code)]
    deck_registry: &'static DeckRegistry,
    deck1: DeckDefinition,
    deck2: DeckDefinition,
    game_mode: GameMode,
    num_envs: usize,
    opponent_type: OpponentType,
    /// Per-env greedy bots (reused across steps for performance)
    greedy_bots: Vec<GreedyBot<'static>>,
    /// Per-env random bots
    random_bots: Vec<RandomBot>,
    /// Per-env MCTS bots (only populated when opponent is Mcts)
    mcts_bots: Vec<MctsBot<'static>>,
    /// Per-env Alpha-Beta bots (only populated when opponent is AlphaBeta)
    alphabeta_bots: Vec<AlphaBetaBot<'static>>,
}

#[pymethods]
impl PyParallelGames {
    /// Create a batch of game environments.
    ///
    /// Args:
    ///     num_envs: Number of parallel environments
    ///     deck1: Name of deck for player 1 (default: "architect_fortify")
    ///     deck2: Name of deck for player 2 (default: "broodmother_pack")
    ///     game_mode: "attrition" (default) or "essence_duel"
    ///     opponent: Built-in opponent for player 2: "greedy" (default), "random", or "self" (agent plays both sides)
    #[new]
    #[pyo3(signature = (num_envs, deck1=None, deck2=None, game_mode=None, opponent=None))]
    fn new(
        num_envs: usize,
        deck1: Option<&str>,
        deck2: Option<&str>,
        game_mode: Option<&str>,
        opponent: Option<&str>,
    ) -> PyResult<Self> {
        let card_db: &'static CardDatabase = &CARD_DB;
        let deck_registry: &'static DeckRegistry = &DECK_REGISTRY;

        // Load decks
        let deck1_name = deck1.unwrap_or("artificer_tokens");
        let deck2_name = deck2.unwrap_or("broodmother_pack");

        let deck1_def = deck_registry.get(deck1_name)
            .ok_or_else(|| PyValueError::new_err(format!("Unknown deck: {}", deck1_name)))?;
        let deck2_def = deck_registry.get(deck2_name)
            .ok_or_else(|| PyValueError::new_err(format!("Unknown deck: {}", deck2_name)))?;

        // Parse game mode (essence_war is the new default, but accept legacy names for compatibility)
        let mode = match game_mode.unwrap_or("essence_war") {
            "attrition" => GameMode::Attrition,
            "essence_war" | "essence-war" | "essence_duel" | "essence-duel" => GameMode::EssenceWar,
            other => return Err(PyValueError::new_err(format!("Unknown game mode: {}", other))),
        };

        // Parse opponent type (default: greedy for training against a real opponent)
        let opponent_str = opponent.unwrap_or("greedy");
        let opponent_type = match opponent_str {
            "greedy" => OpponentType::Greedy,
            "random" => OpponentType::Random,
            "self" | "self-play" | "selfplay" => OpponentType::SelfPlay,
            s if s.starts_with("mcts") => {
                let sims: u32 = s[4..].parse().map_err(|_| PyValueError::new_err(format!(
                    "Invalid MCTS simulation count in '{}'. Expected format: 'mcts25', 'mcts100', etc.", s
                )))?;
                OpponentType::Mcts { simulations: sims }
            }
            s if s.starts_with("alphabeta") => {
                let depth: u32 = s[9..].parse().map_err(|_| PyValueError::new_err(format!(
                    "Invalid Alpha-Beta depth in '{}'. Expected format: 'alphabeta3', 'alphabeta6', etc.", s
                )))?;
                OpponentType::AlphaBeta { depth }
            }
            other => return Err(PyValueError::new_err(format!(
                "Unknown opponent type: '{}'. Valid options: 'greedy', 'random', 'self', 'mcts<N>' (e.g. 'mcts25'), 'alphabeta<N>' (e.g. 'alphabeta3')", other
            ))),
        };

        // Create engines
        let engines: Vec<_> = (0..num_envs)
            .map(|_| GameEngine::new(card_db))
            .collect();

        // Create per-env bots (only for the selected opponent type)
        let greedy_bots: Vec<_> = if matches!(opponent_type, OpponentType::Greedy) {
            (0..num_envs).map(|i| GreedyBot::new(card_db, 1000 + i as u64)).collect()
        } else {
            Vec::new()
        };
        let random_bots: Vec<_> = if matches!(opponent_type, OpponentType::Random) {
            (0..num_envs).map(|i| RandomBot::new(2000 + i as u64)).collect()
        } else {
            Vec::new()
        };
        let mcts_bots: Vec<_> = if let OpponentType::Mcts { simulations } = opponent_type {
            (0..num_envs).map(|i| {
                let config = MctsConfig {
                    simulations,
                    // Disable TT for low sim counts to save memory (8MB per bot)
                    use_transposition_table: simulations > 50,
                    ..MctsConfig::batch(simulations)
                };
                MctsBot::with_config(card_db, config, 3000 + i as u64)
            }).collect()
        } else {
            Vec::new()
        };
        let alphabeta_bots: Vec<_> = if let OpponentType::AlphaBeta { depth } = opponent_type {
            (0..num_envs).map(|i| {
                let config = AlphaBetaConfig {
                    max_depth: depth,
                    // 2MB TT per bot (vs 16MB default) to limit memory with 64 bots
                    tt_size_mb: 2,
                    ..Default::default()
                };
                AlphaBetaBot::with_config(card_db, config, 4000 + i as u64)
            }).collect()
        } else {
            Vec::new()
        };

        Ok(Self {
            engines,
            card_db,
            deck_registry,
            deck1: deck1_def.clone(),
            deck2: deck2_def.clone(),
            game_mode: mode,
            num_envs,
            opponent_type,
            greedy_bots,
            random_bots,
            mcts_bots,
            alphabeta_bots,
        })
    }

    /// Reset all games with given seeds.
    ///
    /// Args:
    ///     seeds: List of random seeds, one per environment
    fn reset(&mut self, seeds: Vec<u64>) -> PyResult<()> {
        if seeds.len() != self.num_envs {
            return Err(PyValueError::new_err(format!(
                "Expected {} seeds, got {}",
                self.num_envs,
                seeds.len()
            )));
        }

        for (engine, seed) in self.engines.iter_mut().zip(seeds.iter()) {
            engine
                .start_game_with_mode(&self.deck1, &self.deck2, *seed, self.game_mode)
                .expect("Failed to start game - commander not found in card database");
        }

        Ok(())
    }

    /// Get batch observations as a 2D NumPy array.
    ///
    /// Returns:
    ///     numpy.ndarray: Float32 array of shape (num_envs, 328)
    fn observe_batch<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray2<f32>> {
        let mut data = vec![0.0f32; self.num_envs * STATE_TENSOR_SIZE];

        for (i, engine) in self.engines.iter().enumerate() {
            let tensor = engine.get_state_tensor();
            let start = i * STATE_TENSOR_SIZE;
            data[start..start + STATE_TENSOR_SIZE].copy_from_slice(&tensor);
        }

        // Create 1D array and reshape to 2D
        let arr = PyArray1::from_slice_bound(py, &data);
        arr.reshape([self.num_envs, STATE_TENSOR_SIZE])
            .expect("Failed to reshape observation batch")
    }

    /// Get batch action masks as a 2D NumPy array.
    ///
    /// Returns:
    ///     numpy.ndarray: Float32 array of shape (num_envs, 422)
    fn action_mask_batch<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray2<f32>> {
        let mut data = vec![0.0f32; self.num_envs * Action::ACTION_SPACE_SIZE];

        for (i, engine) in self.engines.iter().enumerate() {
            let mask = engine.get_legal_action_mask();
            let start = i * Action::ACTION_SPACE_SIZE;
            data[start..start + Action::ACTION_SPACE_SIZE].copy_from_slice(&mask);
        }

        let arr = PyArray1::from_slice_bound(py, &data);
        arr.reshape([self.num_envs, Action::ACTION_SPACE_SIZE])
            .expect("Failed to reshape action mask batch")
    }

    /// Apply actions to all environments.
    ///
    /// The agent's action is applied for Player 1 (index 0). If a built-in
    /// opponent is configured (greedy/random), the opponent then plays all
    /// of Player 2's turns automatically. The returned observations and
    /// masks are always from Player 1's perspective, ready for the next
    /// agent action.
    ///
    /// Args:
    ///     actions: NumPy array of action indices, shape (num_envs,)
    ///
    /// Returns:
    ///     tuple: (rewards, dones) as NumPy arrays of shape (num_envs,)
    fn step_batch<'py>(
        &mut self,
        py: Python<'py>,
        actions: PyReadonlyArray1<'py, u16>,
    ) -> PyResult<BatchStepResult<'py>> {
        let actions = actions.as_slice()?;

        if actions.len() != self.num_envs {
            return Err(PyValueError::new_err(format!(
                "Expected {} actions, got {}",
                self.num_envs,
                actions.len()
            )));
        }

        let mut rewards = vec![0.0f32; self.num_envs];
        let mut dones = vec![false; self.num_envs];

        for i in 0..self.num_envs {
            // Apply the agent's action (Player 1)
            if let Err(e) = self.engines[i].apply_action_by_index(actions[i]) {
                return Err(PyValueError::new_err(format!(
                    "Env {}: {}",
                    i, e
                )));
            }

            // Play opponent turns (Player 2) if a built-in opponent is configured
            if !self.engines[i].is_game_over() {
                match self.opponent_type {
                    OpponentType::SelfPlay => {
                        // No opponent — agent plays both sides (legacy behavior)
                    }
                    OpponentType::Greedy => {
                        play_opponent_turns(&mut self.engines[i], &mut self.greedy_bots[i]);
                    }
                    OpponentType::Random => {
                        play_opponent_turns(&mut self.engines[i], &mut self.random_bots[i]);
                    }
                    OpponentType::Mcts { .. } => {
                        play_opponent_turns(&mut self.engines[i], &mut self.mcts_bots[i]);
                    }
                    OpponentType::AlphaBeta { .. } => {
                        play_opponent_turns(&mut self.engines[i], &mut self.alphabeta_bots[i]);
                    }
                }
            }

            dones[i] = self.engines[i].is_game_over();
            if dones[i] {
                rewards[i] = self.engines[i].get_reward(PlayerId::PLAYER_ONE);
            }
        }

        Ok((
            PyArray1::from_slice_bound(py, &rewards),
            PyArray1::from_slice_bound(py, &dones),
        ))
    }

    /// Get the number of environments.
    #[getter]
    fn num_envs(&self) -> usize {
        self.num_envs
    }

    /// Get current players for all environments.
    ///
    /// Returns:
    ///     list: List of current player indices (0 or 1)
    fn current_players(&self) -> Vec<u8> {
        self.engines.iter().map(|e| e.current_player().0).collect()
    }

    /// Check which environments are done.
    ///
    /// Returns:
    ///     list: List of booleans indicating game over status
    fn dones(&self) -> Vec<bool> {
        self.engines.iter().map(|e| e.is_game_over()).collect()
    }

    /// Reset a single environment by index.
    ///
    /// Args:
    ///     idx: Environment index (0 to num_envs - 1)
    ///     seed: Random seed for this environment
    fn reset_single(&mut self, idx: usize, seed: u64) -> PyResult<()> {
        if idx >= self.num_envs {
            return Err(PyValueError::new_err(format!(
                "Index {} out of bounds for {} environments",
                idx, self.num_envs
            )));
        }

        self.engines[idx]
            .start_game_with_mode(&self.deck1, &self.deck2, seed, self.game_mode)
            .expect("Failed to start game - commander not found in card database");

        // Play opponent's opening turns if Player 2 goes first
        match self.opponent_type {
            OpponentType::SelfPlay => {}
            OpponentType::Greedy => play_opponent_turns(&mut self.engines[idx], &mut self.greedy_bots[idx]),
            OpponentType::Random => play_opponent_turns(&mut self.engines[idx], &mut self.random_bots[idx]),
            OpponentType::Mcts { .. } => play_opponent_turns(&mut self.engines[idx], &mut self.mcts_bots[idx]),
            OpponentType::AlphaBeta { .. } => play_opponent_turns(&mut self.engines[idx], &mut self.alphabeta_bots[idx]),
        }

        Ok(())
    }

    /// Reset all environments with a base seed.
    ///
    /// Each environment gets base_seed + env_index as its seed.
    /// If a built-in opponent is configured, the opponent plays first
    /// if Player 2 has the opening turn.
    ///
    /// Args:
    ///     base_seed: Base random seed
    fn reset_all(&mut self, base_seed: u64) {
        for (i, engine) in self.engines.iter_mut().enumerate() {
            engine
                .start_game_with_mode(&self.deck1, &self.deck2, base_seed + i as u64, self.game_mode)
                .expect("Failed to start game - commander not found in card database");
        }

        // Play opponent's opening turns if Player 2 goes first
        match self.opponent_type {
            OpponentType::SelfPlay => {}
            OpponentType::Greedy => {
                for i in 0..self.num_envs {
                    play_opponent_turns(&mut self.engines[i], &mut self.greedy_bots[i]);
                }
            }
            OpponentType::Random => {
                for i in 0..self.num_envs {
                    play_opponent_turns(&mut self.engines[i], &mut self.random_bots[i]);
                }
            }
            OpponentType::Mcts { .. } => {
                for i in 0..self.num_envs {
                    play_opponent_turns(&mut self.engines[i], &mut self.mcts_bots[i]);
                }
            }
            OpponentType::AlphaBeta { .. } => {
                for i in 0..self.num_envs {
                    play_opponent_turns(&mut self.engines[i], &mut self.alphabeta_bots[i]);
                }
            }
        }
    }
}

/// Internal helper: play opponent turns on an engine using a greedy bot.
/// Separate function to avoid borrow conflicts with PyParallelGames fields.
fn play_opponent_turns(engine: &mut GameEngine<'static>, bot: &mut dyn Bot) {
    while engine.current_player() == PlayerId::PLAYER_TWO && !engine.is_game_over() {
        let action = bot.select_action_with_engine(engine);
        engine
            .apply_action(action)
            .expect("Opponent bot produced invalid action");
    }
}

/// Python module definition.
#[pymodule]
fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyGame>()?;
    m.add_class::<PyParallelGames>()?;

    // Add constants
    m.add("STATE_TENSOR_SIZE", STATE_TENSOR_SIZE)?;
    m.add("ACTION_SPACE_SIZE", Action::ACTION_SPACE_SIZE)?;

    Ok(())
}
