//! Universal tokenization trainer supporting multiple algorithms
//!
//! This module provides a unified training interface for different tokenization
//! algorithms including BPE, `SentencePiece` unigram, and vocabulary-based approaches.

use crate::{
    Result, TokenizationError, Tokenizer,
    bpe::{BpeConfig, BpeTokenizer, BpeTrainer},
    multimodal::{MultimodalConfig, MultimodalTokenizer},
    sentencepiece::{SentencePieceConfig, SentencePieceTokenizer},
};
use serde::{Deserialize, Serialize};
use std::path::Path;

/// Training algorithm types
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum TrainingAlgorithm {
    /// Byte-Pair Encoding
    Bpe(BpeConfig),
    /// `SentencePiece` Unigram
    SentencePieceUnigram(SentencePieceConfig),
    /// `SentencePiece` BPE
    SentencePieceBpe(SentencePieceConfig),
    /// Multimodal tokenization
    Multimodal(MultimodalConfig),
}

/// Training data source
#[derive(Debug, Clone)]
pub enum TrainingData {
    /// Text data from string
    Text(String),
    /// Text data from file path
    TextFile(String),
    /// Image data for multimodal training
    Images(Vec<Vec<u8>>),
    /// Audio data for multimodal training
    Audio(Vec<Vec<f32>>),
    /// Mixed multimodal data
    Multimodal {
        /// Text training data
        text: Vec<String>,
        /// Image training data (raw bytes)
        images: Vec<Vec<u8>>,
        /// Audio training data (waveform samples)
        audio: Vec<Vec<f32>>,
    },
}

/// Training configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingConfig {
    /// Training algorithm and its configuration
    pub algorithm: TrainingAlgorithm,
    /// Number of training iterations
    pub max_iterations: usize,
    /// Early stopping patience
    pub patience: usize,
    /// Validation split ratio
    pub validation_split: f32,
    /// Random seed for reproducibility
    pub seed: Option<u64>,
    /// Enable progress logging
    pub verbose: bool,
    /// Save intermediate models
    pub save_checkpoints: bool,
}

impl Default for TrainingConfig {
    fn default() -> Self {
        Self {
            algorithm: TrainingAlgorithm::Bpe(BpeConfig::default()),
            max_iterations: 1000,
            patience: 10,
            validation_split: 0.1,
            seed: None,
            verbose: true,
            save_checkpoints: false,
        }
    }
}

/// Training metrics and statistics
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct TrainingMetrics {
    /// Training loss history
    pub loss_history: Vec<f64>,
    /// Validation loss history
    pub validation_loss: Vec<f64>,
    /// Vocabulary growth over iterations
    pub vocab_growth: Vec<usize>,
    /// Training time in seconds
    pub training_time_seconds: f64,
    /// Final vocabulary size
    pub final_vocab_size: usize,
    /// Number of iterations completed
    pub iterations_completed: usize,
    /// Whether training converged
    pub converged: bool,
}

/// Universal tokenization trainer
pub struct Trainer {
    config: TrainingConfig,
    metrics: TrainingMetrics,
}

impl Trainer {
    /// Create a new trainer with configuration
    #[must_use]
    pub fn new(config: TrainingConfig) -> Self {
        // Set random seed if provided
        if let Some(seed) = config.seed {
            use rand::{SeedableRng, rngs::StdRng};
            let _rng = StdRng::seed_from_u64(seed);
        }

        Self {
            config,
            metrics: TrainingMetrics::default(),
        }
    }

    /// Train a tokenizer with the provided data
    pub async fn train(&mut self, data: TrainingData) -> Result<TrainedTokenizer> {
        let start_time = std::time::Instant::now();

        if self.config.verbose {
            tracing::info!(
                "Starting tokenization training with algorithm: {:?}",
                std::mem::discriminant(&self.config.algorithm)
            );
        }

        let tokenizer = match &self.config.algorithm {
            TrainingAlgorithm::Bpe(bpe_config) => self.train_bpe(data, bpe_config.clone()).await?,
            TrainingAlgorithm::SentencePieceUnigram(sp_config) => {
                self.train_sentencepiece_unigram(data, sp_config.clone())
                    .await?
            }
            TrainingAlgorithm::SentencePieceBpe(sp_config) => {
                self.train_sentencepiece_bpe(data, sp_config.clone())
                    .await?
            }
            TrainingAlgorithm::Multimodal(mm_config) => {
                self.train_multimodal(data, mm_config.clone()).await?
            }
        };

        self.metrics.training_time_seconds = start_time.elapsed().as_secs_f64();
        self.metrics.converged = true;

        if self.config.verbose {
            tracing::info!(
                "Training completed in {:.2}s with final vocab size: {}",
                self.metrics.training_time_seconds,
                self.metrics.final_vocab_size
            );
        }

        Ok(tokenizer)
    }

    /// Train BPE tokenizer
    async fn train_bpe(
        &mut self,
        data: TrainingData,
        config: BpeConfig,
    ) -> Result<TrainedTokenizer> {
        let mut trainer = BpeTrainer::new(config);

        let texts = self.extract_text_data(data)?;

        // Add training data with progress tracking
        for (i, text) in texts.iter().enumerate() {
            trainer.add_text(text)?;

            if self.config.verbose && i % 1000 == 0 {
                tracing::debug!("Processed {} texts", i);
            }
        }

        let tokenizer = trainer.train()?;
        self.metrics.final_vocab_size = tokenizer.vocab_size();

        Ok(TrainedTokenizer::Bpe(tokenizer))
    }

    /// Train `SentencePiece` Unigram tokenizer
    async fn train_sentencepiece_unigram(
        &mut self,
        _data: TrainingData,
        config: SentencePieceConfig,
    ) -> Result<TrainedTokenizer> {
        // This would require a full SentencePiece implementation
        // For now, return a basic SentencePiece tokenizer
        let tokenizer = SentencePieceTokenizer::new(config);
        self.metrics.final_vocab_size = tokenizer.vocab_size();

        Ok(TrainedTokenizer::SentencePiece(tokenizer))
    }

    /// Train `SentencePiece` BPE tokenizer
    async fn train_sentencepiece_bpe(
        &mut self,
        _data: TrainingData,
        config: SentencePieceConfig,
    ) -> Result<TrainedTokenizer> {
        let tokenizer = SentencePieceTokenizer::new(config);
        self.metrics.final_vocab_size = tokenizer.vocab_size();

        Ok(TrainedTokenizer::SentencePiece(tokenizer))
    }

    /// Train multimodal tokenizer
    async fn train_multimodal(
        &mut self,
        _data: TrainingData,
        config: MultimodalConfig,
    ) -> Result<TrainedTokenizer> {
        let tokenizer = MultimodalTokenizer::new(config);
        self.metrics.final_vocab_size = tokenizer.vocab_size();

        Ok(TrainedTokenizer::Multimodal(tokenizer))
    }

    /// Extract text data from various data sources
    fn extract_text_data(&self, data: TrainingData) -> Result<Vec<String>> {
        match data {
            TrainingData::Text(text) => Ok(vec![text]),
            TrainingData::TextFile(path) => {
                let content = std::fs::read_to_string(&path).map_err(TokenizationError::IoError)?;
                Ok(vec![content])
            }
            TrainingData::Images(_) => Err(TokenizationError::InvalidInput(
                "Cannot extract text from image data".to_string(),
            )),
            TrainingData::Audio(_) => Err(TokenizationError::InvalidInput(
                "Cannot extract text from audio data".to_string(),
            )),
            TrainingData::Multimodal { text, .. } => Ok(text),
        }
    }

    /// Get current training metrics
    #[must_use]
    pub fn metrics(&self) -> &TrainingMetrics {
        &self.metrics
    }

    /// Save trained model to file
    pub async fn save_model<P: AsRef<Path>>(
        &self,
        tokenizer: &TrainedTokenizer,
        path: P,
    ) -> Result<()> {
        let serialized = bincode::serialize(tokenizer)
            .map_err(|e| TokenizationError::InvalidInput(format!("Serialization error: {e}")))?;

        tokio::fs::write(path, serialized)
            .await
            .map_err(TokenizationError::IoError)?;

        Ok(())
    }

    /// Load trained model from file
    pub async fn load_model<P: AsRef<Path>>(path: P) -> Result<TrainedTokenizer> {
        let data = tokio::fs::read(path)
            .await
            .map_err(TokenizationError::IoError)?;

        let tokenizer = bincode::deserialize(&data)
            .map_err(|e| TokenizationError::InvalidInput(format!("Deserialization error: {e}")))?;

        Ok(tokenizer)
    }
}

/// Container for different types of trained tokenizers
#[derive(Debug, Serialize, Deserialize)]
pub enum TrainedTokenizer {
    /// Trained BPE tokenizer
    Bpe(BpeTokenizer),
    /// Trained `SentencePiece` tokenizer
    SentencePiece(SentencePieceTokenizer),
    /// Trained multimodal tokenizer
    Multimodal(MultimodalTokenizer),
}

impl TrainedTokenizer {
    /// Get vocabulary size
    #[must_use]
    pub fn vocab_size(&self) -> usize {
        match self {
            Self::Bpe(tokenizer) => tokenizer.vocab_size(),
            Self::SentencePiece(tokenizer) => tokenizer.vocab_size(),
            Self::Multimodal(tokenizer) => tokenizer.vocab_size(),
        }
    }
}

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

    #[test]
    fn training_config_default() {
        let config = TrainingConfig::default();
        assert_eq!(config.max_iterations, 1000);
        assert_eq!(config.patience, 10);
        assert_eq!(config.validation_split, 0.1);
        assert!(config.verbose);
        assert!(!config.save_checkpoints);
    }

    #[test]
    fn training_data_creation() {
        let text_data = TrainingData::Text("hello world".to_string());
        match text_data {
            TrainingData::Text(text) => assert_eq!(text, "hello world"),
            _ => panic!("Expected text data"),
        }

        let file_data = TrainingData::TextFile("test.txt".to_string());
        match file_data {
            TrainingData::TextFile(path) => assert_eq!(path, "test.txt"),
            _ => panic!("Expected file data"),
        }
    }

    #[test]
    fn training_algorithm_types() {
        let bpe_alg = TrainingAlgorithm::Bpe(BpeConfig::default());
        match bpe_alg {
            TrainingAlgorithm::Bpe(_) => (),
            _ => panic!("Expected BPE algorithm"),
        }

        let sp_config = SentencePieceConfig {
            vocab_size: 8000,
            model_type: SentencePieceModelType::Unigram,
            special_tokens: vec!["[UNK]".to_string()],
            character_coverage_scaled: 9995, // 0.9995 * 10000
            input_sentence_size: 10000,
            normalize_text: true,
            add_dummy_prefix: true,
            min_frequency: 2,
            max_subword_length: 16,
            unk_surface: "▁".to_string(),
        };
        let sp_alg = TrainingAlgorithm::SentencePieceUnigram(sp_config);
        match sp_alg {
            TrainingAlgorithm::SentencePieceUnigram(_) => (),
            _ => panic!("Expected SentencePiece Unigram algorithm"),
        }
    }

    #[tokio::test]
    async fn trainer_creation_and_basic_training() {
        let config = TrainingConfig {
            algorithm: TrainingAlgorithm::Bpe(BpeConfig {
                vocab_size: 100,
                min_frequency: 1,
                ..BpeConfig::default()
            }),
            max_iterations: 10,
            verbose: false,
            ..TrainingConfig::default()
        };

        let mut trainer = Trainer::new(config);
        let data = TrainingData::Text("hello world test hello".to_string());

        let trained_tokenizer = trainer.train(data).await.unwrap();
        assert!(trained_tokenizer.vocab_size() > 0);

        let metrics = trainer.metrics();
        assert!(metrics.training_time_seconds > 0.0);
        assert!(metrics.converged);
    }

    #[test]
    fn training_metrics_default() {
        let metrics = TrainingMetrics::default();
        assert_eq!(metrics.loss_history.len(), 0);
        assert_eq!(metrics.validation_loss.len(), 0);
        assert_eq!(metrics.vocab_growth.len(), 0);
        assert_eq!(metrics.training_time_seconds, 0.0);
        assert_eq!(metrics.final_vocab_size, 0);
        assert_eq!(metrics.iterations_completed, 0);
        assert!(!metrics.converged);
    }
}