//! Generation quality features and filters

pub mod diversity;
pub mod repetition;
pub mod safety;

use crate::Result;
use serde::{Deserialize, Serialize};

/// Quality filter interface
pub trait QualityFilter: Send + Sync {
    /// Check if text passes the quality filter
    fn passes_filter(&self, text: &str) -> Result<bool>;

    /// Get filter score (0.0 to 1.0)
    fn filter_score(&self, text: &str) -> Result<f32>;

    /// Get filter name for logging
    fn name(&self) -> &str;
}

/// Repetition penalty configuration and implementation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RepetitionPenalty {
    pub penalty_factor: f32,
    pub ngram_size: usize,
    pub max_repetitions: usize,
}

impl Default for RepetitionPenalty {
    fn default() -> Self {
        Self {
            penalty_factor: 1.2,
            ngram_size: 3,
            max_repetitions: 2,
        }
    }
}

impl RepetitionPenalty {
    pub fn new(penalty_factor: f32) -> Self {
        Self {
            penalty_factor,
            ..Default::default()
        }
    }

    pub fn calculate_penalty(&self, tokens: &[u32]) -> f32 {
        if tokens.len() < self.ngram_size {
            return 0.0;
        }

        let mut ngram_counts = std::collections::HashMap::new();
        let mut total_penalty = 0.0;

        // Count n-grams
        for window in tokens.windows(self.ngram_size) {
            *ngram_counts.entry(window.to_vec()).or_insert(0) += 1;
        }

        // Apply penalties for repeated n-grams
        for (_, count) in ngram_counts {
            if count > self.max_repetitions {
                total_penalty += (count - self.max_repetitions) as f32 * self.penalty_factor;
            }
        }

        total_penalty
    }
}

impl QualityFilter for RepetitionPenalty {
    fn passes_filter(&self, text: &str) -> Result<bool> {
        let score = self.filter_score(text)?;
        Ok(score < 0.5) // Pass if low repetition
    }

    fn filter_score(&self, text: &str) -> Result<f32> {
        // Mock tokenization for demonstration
        let words: Vec<&str> = text.split_whitespace().collect();
        let tokens: Vec<u32> = words.iter().enumerate().map(|(i, _)| i as u32).collect();

        let penalty = self.calculate_penalty(&tokens);
        Ok((penalty / words.len() as f32).clamp(0.0, 1.0))
    }

    fn name(&self) -> &'static str {
        "repetition_penalty"
    }
}

/// Diversity promoter for encouraging varied generation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DiversityPromoter {
    pub min_unique_ratio: f32,
    pub min_entropy: f32,
}

impl Default for DiversityPromoter {
    fn default() -> Self {
        Self {
            min_unique_ratio: 0.7,
            min_entropy: 2.0,
        }
    }
}

impl QualityFilter for DiversityPromoter {
    fn passes_filter(&self, text: &str) -> Result<bool> {
        let unique_ratio = self.calculate_unique_ratio(text);
        let entropy = self.calculate_entropy(text);

        Ok(unique_ratio >= self.min_unique_ratio && entropy >= self.min_entropy)
    }

    fn filter_score(&self, text: &str) -> Result<f32> {
        let unique_ratio = self.calculate_unique_ratio(text);
        let entropy = self.calculate_entropy(text);

        let diversity_score = (unique_ratio * 0.5) + ((entropy / 5.0).clamp(0.0, 1.0) * 0.5);
        Ok(diversity_score)
    }

    fn name(&self) -> &'static str {
        "diversity_promoter"
    }
}

impl DiversityPromoter {
    fn calculate_unique_ratio(&self, text: &str) -> f32 {
        let words: Vec<&str> = text.split_whitespace().collect();
        if words.is_empty() {
            return 0.0;
        }

        let unique_words: std::collections::HashSet<_> = words.iter().collect();
        unique_words.len() as f32 / words.len() as f32
    }

    fn calculate_entropy(&self, text: &str) -> f32 {
        let words: Vec<&str> = text.split_whitespace().collect();
        if words.is_empty() {
            return 0.0;
        }

        let mut word_counts = std::collections::HashMap::new();
        for word in &words {
            *word_counts.entry(*word).or_insert(0) += 1;
        }

        let total_words = words.len() as f32;
        let entropy: f32 = word_counts
            .values()
            .map(|&count| {
                let p = count as f32 / total_words;
                -p * p.log2()
            })
            .sum();

        entropy
    }
}

/// Toxicity filter using simple keyword detection
#[derive(Debug, Clone)]
pub struct ToxicityFilter {
    toxic_keywords: Vec<String>,
    threshold: f32,
}

impl Default for ToxicityFilter {
    fn default() -> Self {
        Self {
            toxic_keywords: vec![
                "hate".to_string(),
                "violence".to_string(),
                // Would include more comprehensive list
            ],
            threshold: 0.5,
        }
    }
}

impl ToxicityFilter {
    pub fn new(threshold: f32) -> Self {
        Self {
            threshold,
            ..Default::default()
        }
    }
}

impl QualityFilter for ToxicityFilter {
    fn passes_filter(&self, text: &str) -> Result<bool> {
        let score = self.filter_score(text)?;
        Ok(score < self.threshold)
    }

    fn filter_score(&self, text: &str) -> Result<f32> {
        let text_lower = text.to_lowercase();
        let mut toxicity_score = 0.0f32;

        for keyword in &self.toxic_keywords {
            if text_lower.contains(keyword) {
                toxicity_score += 0.3; // Simple scoring
            }
        }

        Ok(toxicity_score.clamp(0.0, 1.0))
    }

    fn name(&self) -> &'static str {
        "toxicity_filter"
    }
}

/// Combined quality checker
pub struct QualityChecker {
    filters: Vec<Box<dyn QualityFilter>>,
}

impl Default for QualityChecker {
    fn default() -> Self {
        Self::new()
    }
}

impl QualityChecker {
    pub fn new() -> Self {
        Self {
            filters: Vec::new(),
        }
    }

    pub fn add_filter(mut self, filter: Box<dyn QualityFilter>) -> Self {
        self.filters.push(filter);
        self
    }

    pub fn check_quality(&self, text: &str) -> Result<QualityReport> {
        let mut passed_filters = Vec::new();
        let mut failed_filters = Vec::new();
        let mut scores = std::collections::HashMap::new();

        for filter in &self.filters {
            let score = filter.filter_score(text)?;
            let passes = filter.passes_filter(text)?;

            scores.insert(filter.name().to_string(), score);

            if passes {
                passed_filters.push(filter.name().to_string());
            } else {
                failed_filters.push(filter.name().to_string());
            }
        }

        let overall_score = if scores.is_empty() {
            1.0
        } else {
            scores.values().sum::<f32>() / scores.len() as f32
        };

        Ok(QualityReport {
            overall_score,
            passed_filters,
            failed_filters: failed_filters.clone(),
            scores,
            passes_all: failed_filters.is_empty(),
        })
    }
}

#[derive(Debug, Clone)]
pub struct QualityReport {
    pub overall_score: f32,
    pub passed_filters: Vec<String>,
    pub failed_filters: Vec<String>,
    pub scores: std::collections::HashMap<String, f32>,
    pub passes_all: bool,
}

// Re-exports
pub use diversity::*;
pub use repetition::*;
pub use safety::*;

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

    #[test]
    fn test_repetition_penalty() -> Result<()> {
        let penalty = RepetitionPenalty::default();

        let repetitive_text = "hello hello hello world world world";
        let normal_text = "this is a normal sentence without repetition";

        let rep_score = penalty.filter_score(repetitive_text)?;
        let normal_score = penalty.filter_score(normal_text)?;

        assert!(rep_score > normal_score);

        Ok(())
    }

    #[test]
    fn test_diversity_promoter() -> Result<()> {
        let promoter = DiversityPromoter::default();

        let diverse_text = "the quick brown fox jumps over lazy dog";
        let monotonous_text = "same same same same same same same same";

        let diverse_score = promoter.filter_score(diverse_text)?;
        let mono_score = promoter.filter_score(monotonous_text)?;

        assert!(diverse_score > mono_score);

        Ok(())
    }

    #[test]
    fn test_quality_checker() -> Result<()> {
        let checker = QualityChecker::new()
            .add_filter(Box::new(RepetitionPenalty::default()))
            .add_filter(Box::new(DiversityPromoter::default()));

        let text = "This is a diverse and interesting text without repetition";
        let report = checker.check_quality(text)?;

        assert!(report.overall_score >= 0.0 && report.overall_score <= 1.0);
        assert_eq!(report.scores.len(), 2);

        Ok(())
    }
}