//! Experience dataset serialization for offline RL
//!
//! Supports saving and loading fixed datasets of (state, action, reward, `next_state`, done)
//! tuples for offline CQL/IQL training. Datasets are serialized as bincode for compact
//! storage and fast I/O.

use std::io::{BufReader, BufWriter};
use std::path::Path;

use serde::{Deserialize, Serialize};

use super::experience::Experience;

/// A fixed dataset of experiences for offline RL training.
///
/// Contains a flat vector of experience tuples plus metadata about how the dataset
/// was collected (behavior policy, symbol, timestamp). This is the primary interchange
/// format between online data collection and offline training (CQL/IQL).
#[derive(Debug, Serialize, Deserialize)]
pub struct ExperienceDataset {
    /// All experiences in the dataset
    pub experiences: Vec<Experience>,
    /// Metadata about how the dataset was collected
    pub metadata: DatasetMetadata,
}

/// Metadata describing the provenance of an experience dataset.
#[derive(Debug, Serialize, Deserialize)]
pub struct DatasetMetadata {
    /// Number of experiences
    pub num_experiences: usize,
    /// State dimension
    pub state_dim: usize,
    /// Model that generated the data (behavior policy)
    pub behavior_policy: String,
    /// Symbol traded
    pub symbol: String,
    /// Collection timestamp (RFC 3339)
    pub collected_at: String,
}

impl ExperienceDataset {
    /// Create a new dataset from a vector of experiences.
    ///
    /// Automatically populates metadata from the first experience's state dimension.
    pub fn new(experiences: Vec<Experience>, behavior_policy: &str, symbol: &str) -> Self {
        let state_dim = experiences.first().map_or(0, |e| e.state.len());
        Self {
            metadata: DatasetMetadata {
                num_experiences: experiences.len(),
                state_dim,
                behavior_policy: behavior_policy.to_owned(),
                symbol: symbol.to_owned(),
                collected_at: chrono::Utc::now().to_rfc3339(),
            },
            experiences,
        }
    }

    /// Save dataset to a bincode file at the given path.
    ///
    /// Uses buffered I/O for performance on large datasets.
    ///
    /// # Errors
    ///
    /// Returns an error if the file cannot be created or serialization fails.
    pub fn save<P: AsRef<Path>>(&self, path: P) -> Result<(), Box<dyn std::error::Error>> {
        let file = std::fs::File::create(path.as_ref())?;
        let writer = BufWriter::new(file);
        bincode::serialize_into(writer, self)?;
        tracing::info!(
            "Saved dataset: {} experiences, state_dim={}, policy={}",
            self.metadata.num_experiences,
            self.metadata.state_dim,
            self.metadata.behavior_policy
        );
        Ok(())
    }

    /// Load dataset from a bincode file at the given path.
    ///
    /// Uses buffered I/O for performance on large datasets.
    ///
    /// # Errors
    ///
    /// Returns an error if the file cannot be opened or deserialization fails.
    pub fn load<P: AsRef<Path>>(path: P) -> Result<Self, Box<dyn std::error::Error>> {
        let file = std::fs::File::open(path.as_ref())?;
        let reader = BufReader::new(file);
        let dataset: Self = bincode::deserialize_from(reader)?;
        tracing::info!(
            "Loaded dataset: {} experiences, state_dim={}, policy={}",
            dataset.metadata.num_experiences,
            dataset.metadata.state_dim,
            dataset.metadata.behavior_policy
        );
        Ok(dataset)
    }
}

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

    #[test]
    fn test_dataset_roundtrip() {
        let experiences = vec![
            Experience::new(vec![1.0, 2.0, 3.0], 0, 0.5, vec![1.1, 2.1, 3.1], false),
            Experience::new(vec![2.0, 3.0, 4.0], 1, -0.2, vec![2.1, 3.1, 4.1], true),
        ];
        let dataset = ExperienceDataset::new(experiences, "dqn_baseline", "ES.FUT");

        assert_eq!(dataset.metadata.num_experiences, 2);
        assert_eq!(dataset.metadata.state_dim, 3);
        assert_eq!(dataset.metadata.behavior_policy, "dqn_baseline");
        assert_eq!(dataset.metadata.symbol, "ES.FUT");

        // Save and reload via tempfile
        let dir = tempfile::tempdir().ok();
        if let Some(ref d) = dir {
            let path = d.path().join("test_dataset.bin");
            dataset.save(&path).ok();
            let loaded = ExperienceDataset::load(&path);
            if let Ok(loaded) = loaded {
                assert_eq!(loaded.metadata.num_experiences, 2);
                assert_eq!(loaded.experiences.len(), 2);
                assert_eq!(loaded.experiences.first().map(|e| e.action), Some(0));
            }
        }
    }

    #[test]
    fn test_empty_dataset() {
        let dataset = ExperienceDataset::new(Vec::new(), "none", "");
        assert_eq!(dataset.metadata.num_experiences, 0);
        assert_eq!(dataset.metadata.state_dim, 0);
    }
}
