//! Neural network bot using ONNX Runtime for inference.
//!
//! Loads trained neural network models exported as ONNX files and uses them
//! to select actions. The ONNX model is expected to have observation
//! normalization baked in, making it fully self-contained.
//!
//! # ONNX Model Interface
//!
//! **Inputs:**
//! - `obs`: `[batch, 328]` float32 - game state tensor
//! - `action_mask`: `[batch, 422]` bool - legal action mask (true = legal)
//!
//! **Outputs:**
//! - `policy_logits`: `[batch, 422]` float32 - raw action logits
//! - `value`: `[batch]` float32 - position evaluation

use std::path::{Path, PathBuf};
use std::sync::Mutex;

use ort::session::Session;
use ort::value::Tensor;

use crate::actions::Action;
use crate::tensor::STATE_TENSOR_SIZE;

use super::Bot;

/// A bot that uses a trained ONNX neural network for action selection.
pub struct NeuralBot {
    /// Session wrapped in Mutex because run() takes &mut self.
    session: Mutex<Session>,
    model_path: PathBuf,
    name: String,
    /// Whether to use deterministic (argmax) or stochastic (sample) action selection.
    deterministic: bool,
    /// Random seed for stochastic action selection.
    seed: u64,
    rng_state: u64,
}

// Safety: Session is Send but not Sync. We wrap in Mutex for interior mutability.
unsafe impl Send for NeuralBot {}

impl NeuralBot {
    /// Create a new NeuralBot from an ONNX model file.
    ///
    /// The model must have been exported with the expected input/output interface
    /// (see module docs). Models exported via `essence_wars.agents.export_onnx`
    /// automatically satisfy this.
    ///
    /// # Arguments
    /// * `model_path` - Path to the .onnx model file
    /// * `deterministic` - If true, always picks the highest-scoring action
    /// * `seed` - Random seed (used when `deterministic` is false)
    pub fn new(model_path: &Path, deterministic: bool, seed: u64) -> Result<Self, NeuralBotError> {
        let session = Session::builder()
            .map_err(|e| NeuralBotError::SessionCreation(e.to_string()))?
            .with_intra_threads(1)
            .map_err(|e| NeuralBotError::SessionCreation(e.to_string()))?
            .commit_from_file(model_path)
            .map_err(|e| NeuralBotError::ModelLoad(model_path.to_path_buf(), e.to_string()))?;

        // Derive a display name from the filename
        let name = model_path
            .file_stem()
            .and_then(|s| s.to_str())
            .unwrap_or("neural")
            .to_string();

        Ok(Self {
            session: Mutex::new(session),
            model_path: model_path.to_path_buf(),
            name: format!("Neural({})", name),
            deterministic,
            seed,
            rng_state: seed,
        })
    }

    /// Run inference and return (policy_logits, value).
    fn infer(
        &self,
        state_tensor: &[f32; STATE_TENSOR_SIZE],
        legal_mask: &[f32; Action::ACTION_SPACE_SIZE],
    ) -> Result<(Vec<f32>, f32), NeuralBotError> {
        // Prepare observation tensor (batch size 1)
        let obs_data: Vec<f32> = state_tensor.to_vec();
        let obs = Tensor::from_array(([1, STATE_TENSOR_SIZE], obs_data))
            .map_err(|e| NeuralBotError::Inference(format!("obs tensor: {}", e)))?;

        // Convert float mask to bool mask tensor
        let mask_bools: Vec<bool> = legal_mask.iter().map(|&v| v > 0.5).collect();
        let mask = Tensor::from_array(([1, Action::ACTION_SPACE_SIZE], mask_bools))
            .map_err(|e| NeuralBotError::Inference(format!("mask tensor: {}", e)))?;

        // Run inference (session.run requires &mut self)
        let mut session = self.session.lock().unwrap();
        let outputs = session
            .run(ort::inputs![obs, mask])
            .map_err(|e| NeuralBotError::Inference(e.to_string()))?;

        // Extract policy logits: try_extract_tensor returns (Shape, &[T])
        let (_shape, policy_data) = outputs[0]
            .try_extract_tensor::<f32>()
            .map_err(|e| NeuralBotError::Inference(format!("policy extraction: {}", e)))?;
        let policy_logits: Vec<f32> = policy_data.to_vec();

        // Extract value
        let (_shape, value_data) = outputs[1]
            .try_extract_tensor::<f32>()
            .map_err(|e| NeuralBotError::Inference(format!("value extraction: {}", e)))?;
        let value = value_data.first().copied().unwrap_or(0.0);

        Ok((policy_logits, value))
    }

    /// Simple xorshift64 PRNG for stochastic sampling.
    fn next_random(&mut self) -> f64 {
        self.rng_state ^= self.rng_state << 13;
        self.rng_state ^= self.rng_state >> 7;
        self.rng_state ^= self.rng_state << 17;
        (self.rng_state as f64) / (u64::MAX as f64)
    }

    /// Get the model file path.
    pub fn model_path(&self) -> &Path {
        &self.model_path
    }

    /// Run inference and return structured results for introspection.
    pub fn infer_with_details(
        &self,
        state_tensor: &[f32; STATE_TENSOR_SIZE],
        legal_mask: &[f32; Action::ACTION_SPACE_SIZE],
    ) -> Result<NeuralInferenceResult, NeuralBotError> {
        let (policy_logits, value) = self.infer(state_tensor, legal_mask)?;
        Ok(NeuralInferenceResult {
            policy_logits,
            value,
            model_name: self.name.clone(),
        })
    }
}

/// Result of neural network inference for introspection.
pub struct NeuralInferenceResult {
    /// Raw policy logits (422 values).
    pub policy_logits: Vec<f32>,
    /// Value estimate.
    pub value: f32,
    /// Model name.
    pub model_name: String,
}

impl Bot for NeuralBot {
    fn name(&self) -> &str {
        &self.name
    }

    fn select_action(
        &mut self,
        state_tensor: &[f32; STATE_TENSOR_SIZE],
        legal_mask: &[f32; Action::ACTION_SPACE_SIZE],
        legal_actions: &[Action],
    ) -> Action {
        if legal_actions.len() == 1 {
            return legal_actions[0];
        }

        let (policy_logits, _value) = match self.infer(state_tensor, legal_mask) {
            Ok(result) => result,
            Err(e) => {
                log::error!("Neural inference failed: {}, falling back to first legal action", e);
                return legal_actions[0];
            }
        };

        if self.deterministic {
            // Argmax over legal actions
            let mut best_action = legal_actions[0];
            let mut best_logit = f32::NEG_INFINITY;

            for &action in legal_actions {
                let idx = action.to_index() as usize;
                if idx < policy_logits.len() && policy_logits[idx] > best_logit {
                    best_logit = policy_logits[idx];
                    best_action = action;
                }
            }

            best_action
        } else {
            // Softmax sampling over legal actions
            let mut max_logit = f32::NEG_INFINITY;
            for &action in legal_actions {
                let idx = action.to_index() as usize;
                if idx < policy_logits.len() && policy_logits[idx] > max_logit {
                    max_logit = policy_logits[idx];
                }
            }

            let mut exp_sum = 0.0f64;
            let mut probs: Vec<(Action, f64)> = Vec::with_capacity(legal_actions.len());
            for &action in legal_actions {
                let idx = action.to_index() as usize;
                let logit = if idx < policy_logits.len() {
                    policy_logits[idx]
                } else {
                    f32::NEG_INFINITY
                };
                let exp = ((logit - max_logit) as f64).exp();
                exp_sum += exp;
                probs.push((action, exp));
            }

            // Sample from the distribution
            let r = self.next_random() * exp_sum;
            let mut cumsum = 0.0;
            for (action, prob) in &probs {
                cumsum += prob;
                if cumsum >= r {
                    return *action;
                }
            }

            // Fallback (shouldn't happen)
            legal_actions[legal_actions.len() - 1]
        }
    }

    fn reset(&mut self) {
        self.rng_state = self.seed;
    }

    fn clone_box(&self) -> Box<dyn Bot> {
        // Re-load the model for the clone
        match NeuralBot::new(&self.model_path, self.deterministic, self.seed) {
            Ok(bot) => Box::new(bot),
            Err(e) => {
                log::error!("Failed to clone NeuralBot: {}", e);
                // Fallback to random bot
                Box::new(super::RandomBot::new(self.seed))
            }
        }
    }
}

/// Errors that can occur when working with neural bots.
#[derive(Debug, Clone)]
pub enum NeuralBotError {
    /// Failed to create ONNX Runtime session.
    SessionCreation(String),
    /// Failed to load the ONNX model file.
    ModelLoad(PathBuf, String),
    /// Failed during model inference.
    Inference(String),
}

impl std::fmt::Display for NeuralBotError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            NeuralBotError::SessionCreation(e) => write!(f, "ONNX session creation failed: {}", e),
            NeuralBotError::ModelLoad(path, e) => {
                write!(f, "Failed to load model {:?}: {}", path, e)
            }
            NeuralBotError::Inference(e) => write!(f, "Neural inference failed: {}", e),
        }
    }
}

impl std::error::Error for NeuralBotError {}
