//! xLSTM inference adapter for ensemble prediction
//!
//! Wraps an `XLSTMNetwork` with a sequence buffer. Feature vectors accumulate
//! until `sequence_length` frames are collected, then the sLSTM+mLSTM blocks
//! run a forward pass and the scalar output is mapped via sigmoid to
//! direction + confidence.

use std::collections::VecDeque;
use std::sync::{Arc, Mutex};

use cudarc::driver::CudaStream;
use ml_core::device::MlDevice;

use crate::ensemble::inference_adapter::{
    EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, RawPrediction,
};
use crate::gpu::DeviceConfig;
use crate::xlstm::config::XLSTMConfig;
use crate::xlstm::network::XLSTMNetwork;
use crate::{MLError, MLResult};

use ml_supervised::gpu_tensor::GpuTensor;

/// Extract a CudaStream from an MlDevice.
fn extract_cuda_stream(device: &MlDevice) -> MLResult<Arc<CudaStream>> {
    device.cuda_stream().cloned().map_err(|e| MLError::ConfigError(format!("CUDA required for xLSTM: {e}")))
}

/// Inference adapter for xLSTM (sLSTM + mLSTM blocks) with sequence buffering.
///
/// Maintains a ring buffer of feature vectors. Returns a neutral prediction
/// (direction=0, confidence=0) until `sequence_length` frames have been
/// collected. Once the buffer is full, each call builds a
/// `[1, seq_len, input_dim]` tensor and runs the xLSTM forward pass.
/// The scalar output is mapped via sigmoid to direction [-1, 1] and
/// confidence [0, 1].
#[allow(missing_debug_implementations)]
pub struct XlstmInferenceAdapter {
    model: Mutex<XLSTMNetwork>,
    buffer: Mutex<VecDeque<Vec<f32>>>,
    #[allow(dead_code)]
    _placeholder: (),  // VarMap removed: GpuLinear weights are not Candle Vars
    sequence_length: usize,
    input_dim: usize,
    device: MlDevice,
}

// SAFETY: XLSTMNetwork uses GpuTensor/CudaSlice which are Send+Sync.
// The Mutex provides exclusive access for inference calls and buffer mutation.
#[allow(unsafe_code)]
unsafe impl Send for XlstmInferenceAdapter {}
#[allow(unsafe_code)]
unsafe impl Sync for XlstmInferenceAdapter {}

impl XlstmInferenceAdapter {
    /// Create a new xLSTM inference adapter.
    ///
    /// # Arguments
    /// * `config` - xLSTM model configuration
    /// * `sequence_length` - Number of feature vectors to buffer before inference
    pub fn new(config: XLSTMConfig, sequence_length: usize) -> MLResult<Self> {
        let device = DeviceConfig::Auto.resolve().map_err(|e| MLError::DeviceError(format!("CUDA required: {e}")))?;
        let input_dim = config.input_dim;
        let stream = extract_cuda_stream(&device)?;
        let network = XLSTMNetwork::new(&config, &stream)?;

        Ok(Self {
            model: Mutex::new(network),
            buffer: Mutex::new(VecDeque::with_capacity(sequence_length)),
            _placeholder: (),
            sequence_length,
            input_dim,
            device,
        })
    }

    /// Create an xLSTM inference adapter and load weights from a safetensors checkpoint.
    pub fn from_checkpoint(
        config: XLSTMConfig,
        sequence_length: usize,
        path: &str,
    ) -> MLResult<Self> {
        let device = DeviceConfig::Auto.resolve().map_err(|e| MLError::DeviceError(format!("CUDA required: {e}")))?;
        let input_dim = config.input_dim;
        let stream = extract_cuda_stream(&device)?;
        let network = XLSTMNetwork::new(&config, &stream)?;

        // TODO: Load GpuLinear weights from safetensors checkpoint
        tracing::warn!("xLSTM checkpoint loading not yet supported for GpuLinear weights: {}", path);

        Ok(Self {
            model: Mutex::new(network),
            buffer: Mutex::new(VecDeque::with_capacity(sequence_length)),
            _placeholder: (),
            sequence_length,
            input_dim,
            device,
        })
    }

    /// Pad or truncate feature values to match input_dim, converting f64 to f32.
    fn pad_features(&self, values: &[f64]) -> Vec<f32> {
        let mut padded = vec![0.0_f32; self.input_dim];
        let copy_len = values.len().min(self.input_dim);
        for i in 0..copy_len {
            if let Some(v) = values.get(i) {
                if let Some(slot) = padded.get_mut(i) {
                    *slot = *v as f32;
                }
            }
        }
        padded
    }
}

impl ModelInferenceAdapter for XlstmInferenceAdapter {
    fn model_name(&self) -> &str {
        "xLSTM"
    }

    fn predict(&self, features: &FeatureVector) -> MLResult<EnsemblePrediction> {
        let start = std::time::Instant::now();

        let padded = self.pad_features(&features.values);

        // Buffer phase (short-lived lock)
        let buffer_ready = {
            let mut buf = self
                .buffer
                .lock()
                .map_err(|e| MLError::LockError(format!("xLSTM buffer lock poisoned: {e}")))?;
            buf.push_back(padded);
            while buf.len() > self.sequence_length {
                buf.pop_front();
            }
            buf.len() >= self.sequence_length
        };

        if !buffer_ready {
            let latency_us = start.elapsed().as_micros() as u64;
            return Ok(EnsemblePrediction {
                model_name: "xLSTM".to_owned(),
                direction: 0.0,
                confidence: 0.0,
                metadata: PredictionMeta {
                    latency_us,
                    quantiles: None,
                    attention_weights: None,
                    q_values: None,
                },
            });
        }

        // Build [1, seq_len, input_dim] tensor from buffer snapshot
        let flat_data = {
            let buf = self
                .buffer
                .lock()
                .map_err(|e| MLError::LockError(format!("xLSTM buffer lock poisoned: {e}")))?;
            let mut data = Vec::with_capacity(self.sequence_length * self.input_dim);
            for frame in buf.iter() {
                data.extend_from_slice(frame);
            }
            data
        };

        let stream = extract_cuda_stream(&self.device)?;
        let gpu_input = GpuTensor::from_vec(
            flat_data,
            &[1, self.sequence_length, self.input_dim],
            &stream,
        )?;

        let model = self
            .model
            .lock()
            .map_err(|e| MLError::LockError(format!("xLSTM model lock poisoned: {e}")))?;
        let gpu_output = model.forward(&gpu_input)?;
        let output_host = gpu_output.to_vec() // gpu-exit: scalar prediction
            .map_err(|e| MLError::ModelError(format!("xLSTM output readback: {e}")))?;
        drop(model);

        let raw_val = output_host.first().copied().unwrap_or(0.0) as f64;

        // Sigmoid -> direction + confidence
        let prob = 1.0 / (1.0 + (-raw_val).exp());
        let direction = (2.0 * prob - 1.0).clamp(-1.0, 1.0);
        let confidence = ((prob - 0.5).abs() * 2.0).clamp(0.0, 1.0);

        let latency_us = start.elapsed().as_micros() as u64;

        Ok(EnsemblePrediction {
            model_name: "xLSTM".to_owned(),
            direction,
            confidence,
            metadata: PredictionMeta {
                latency_us,
                quantiles: None,
                attention_weights: None,
                q_values: None,
            },
        })
    }

    fn predict_raw(&self, features: &FeatureVector) -> MLResult<RawPrediction> {
        let padded = self.pad_features(&features.values);

        // Buffer phase (short-lived lock)
        let buffer_ready = {
            let mut buf = self
                .buffer
                .lock()
                .map_err(|e| MLError::LockError(format!("xLSTM buffer lock poisoned: {e}")))?;
            buf.push_back(padded);
            while buf.len() > self.sequence_length {
                buf.pop_front();
            }
            buf.len() >= self.sequence_length
        };

        if !buffer_ready {
            return Ok(RawPrediction {
                direction_scalar: 0.0,
                confidence: 0.0,
                logits: None,
            });
        }

        // Build [1, seq_len, input_dim] GpuTensor from buffer snapshot
        let flat_data = {
            let buf = self
                .buffer
                .lock()
                .map_err(|e| MLError::LockError(format!("xLSTM buffer lock poisoned: {e}")))?;
            let mut data = Vec::with_capacity(self.sequence_length * self.input_dim);
            for frame in buf.iter() {
                data.extend_from_slice(frame);
            }
            data
        };

        let stream = extract_cuda_stream(&self.device)?;
        let gpu_input = GpuTensor::from_vec(
            flat_data,
            &[1, self.sequence_length, self.input_dim],
            &stream,
        )?;

        let model = self
            .model
            .lock()
            .map_err(|e| MLError::LockError(format!("xLSTM model lock poisoned: {e}")))?;
        let gpu_output = model.forward(&gpu_input)?;
        let output_host = gpu_output.to_vec() // gpu-exit: scalar prediction
            .map_err(|e| MLError::ModelError(format!("xLSTM output readback: {e}")))?;
        drop(model);

        let raw_f32 = output_host.first().copied().unwrap_or(0.0);
        let prob = 1.0 / (1.0 + (-raw_f32 as f64).exp());
        let confidence = ((prob - 0.5).abs() * 2.0).clamp(0.0, 1.0);

        Ok(RawPrediction {
            direction_scalar: 0.0,
            confidence,
            logits: Some(vec![raw_f32]),
        })
    }

    fn is_ready(&self) -> bool {
        let buf_ok = self
            .buffer
            .lock()
            .map(|buf| buf.len() >= self.sequence_length)
            .unwrap_or(false);
        let model_ok = self.model.lock().is_ok();
        buf_ok && model_ok
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_config() -> XLSTMConfig {
        XLSTMConfig {
            input_dim: 51,
            hidden_dim: 32,
            num_blocks: 2,
            num_heads: 2,
            output_dim: 1,
            dropout: 0.0,
            ..Default::default()
        }
    }

    const TEST_SEQ_LEN: usize = 4;

    #[test]
    fn test_xlstm_adapter_creation() {
        let adapter = XlstmInferenceAdapter::new(test_config(), TEST_SEQ_LEN);
        assert!(
            adapter.is_ok(),
            "xLSTM adapter creation failed: {:?}",
            adapter.err()
        );
        let adapter = match adapter {
            Ok(a) => a,
            Err(_) => return,
        };
        assert_eq!(adapter.model_name(), "xLSTM");
        assert!(!adapter.is_ready(), "Should not be ready with empty buffer");
    }

    #[test]
    fn test_xlstm_adapter_buffers_and_predicts() {
        let adapter = match XlstmInferenceAdapter::new(test_config(), TEST_SEQ_LEN) {
            Ok(a) => a,
            Err(e) => {
                panic!("Failed to create xLSTM adapter: {e:?}");
            }
        };

        for i in 0..TEST_SEQ_LEN {
            let fv = FeatureVector {
                values: vec![0.1 * (i as f64 + 1.0); 51],
                timestamp: 1_700_000_000_000_000 + i as i64,
            };
            let pred = match adapter.predict(&fv) {
                Ok(p) => p,
                Err(e) => {
                    panic!("predict failed at step {i}: {e:?}");
                }
            };
            if i < TEST_SEQ_LEN - 1 {
                assert_eq!(pred.direction, 0.0, "Neutral direction while buffering");
                assert_eq!(pred.confidence, 0.0, "Zero confidence while buffering");
            } else {
                assert!(
                    pred.direction >= -1.0 && pred.direction <= 1.0,
                    "direction {} out of [-1,1]",
                    pred.direction
                );
                assert!(
                    pred.confidence >= 0.0 && pred.confidence <= 1.0,
                    "confidence {} out of [0,1]",
                    pred.confidence
                );
            }
        }
        assert!(adapter.is_ready());
    }

    #[test]
    fn test_xlstm_adapter_valid_outputs() {
        // xLSTM is stateful (LSTM hidden state mutates on each prediction),
        // so consecutive predictions with the same input will NOT be identical.
        // This test validates that outputs are finite and bounded.
        let adapter = match XlstmInferenceAdapter::new(test_config(), TEST_SEQ_LEN) {
            Ok(a) => a,
            Err(e) => {
                panic!("Failed to create xLSTM adapter: {e:?}");
            }
        };

        // Fill buffer
        for _ in 0..TEST_SEQ_LEN {
            let fv = FeatureVector {
                values: vec![0.1; 51],
                timestamp: 1_700_000_000_000_000,
            };
            let _ = adapter.predict(&fv);
        }

        // Multiple predictions should all be finite and bounded
        for _ in 0..5 {
            let fv = FeatureVector {
                values: vec![0.1; 51],
                timestamp: 1_700_000_000_000_000,
            };
            let pred = match adapter.predict(&fv) {
                Ok(p) => p,
                Err(e) => {
                    panic!("predict failed: {e:?}");
                }
            };

            assert!(
                pred.direction >= -1.0 && pred.direction <= 1.0 && pred.direction.is_finite(),
                "direction {} not in [-1,1] or not finite",
                pred.direction
            );
            assert!(
                pred.confidence >= 0.0 && pred.confidence <= 1.0 && pred.confidence.is_finite(),
                "confidence {} not in [0,1] or not finite",
                pred.confidence
            );
        }
    }
}
