use serde::{Deserialize, Serialize};
use super::{StageResult, StageError};
use super::tokenization::Token;
use crate::types::{Dialect, Script};

/// Stage 4: Intermediate Canonical Representation (ICR)
/// Converts tokens to a unified canonical form
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ICR {
    /// Canonical phonemes/graphemes
    pub segments: Vec<ICRSegment>,
    /// Original dialect
    pub dialect: Dialect,
    /// Metadata
    pub metadata: ICRMetadata,
}

/// A segment in the ICR (represents a canonical unit)
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ICRSegment {
    /// Canonical representation (phonetic/phonemic)
    pub canonical: String,
    /// Original text
    pub original: String,
    /// Segment type
    pub segment_type: SegmentType,
    /// Confidence score (0.0 - 1.0)
    pub confidence: f64,
}

/// Types of ICR segments
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub enum SegmentType {
    /// Consonant
    Consonant,
    /// Vowel
    Vowel,
    /// Function word
    FunctionWord,
    /// Number
    Number,
    /// Punctuation
    Punctuation,
    /// Unknown
    Unknown,
}

/// Metadata about the ICR
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ICRMetadata {
    pub source_script: Script,
    pub total_segments: usize,
    pub confidence_avg: f64,
}

/// Builder for ICR
pub struct ICRBuilder {
    dialect: Dialect,
}

impl ICRBuilder {
    pub fn new(dialect: Dialect) -> Self {
        ICRBuilder { dialect }
    }

    /// Convert tokens to ICR
    pub fn build_from_tokens(&self, tokens: &[Token], source_script: Script) -> StageResult<ICR> {
        let mut segments = Vec::new();
        let mut total_confidence = 0.0;

        for token in tokens {
            let token_segments = self.tokenize_to_segments(&token.text, source_script)?;
            total_confidence += token_segments.iter().map(|s| s.confidence).sum::<f64>();
            segments.extend(token_segments);
        }

        let total_segments = segments.len();
        let confidence_avg = if total_segments > 0 {
            total_confidence / total_segments as f64
        } else {
            0.0
        };

        Ok(ICR {
            segments,
            dialect: self.dialect,
            metadata: ICRMetadata {
                source_script,
                total_segments,
                confidence_avg,
            },
        })
    }

    /// Convert a single token to ICR segments
    fn tokenize_to_segments(&self, text: &str, source_script: Script) -> StageResult<Vec<ICRSegment>> {
        match source_script {
            Script::Arabic => self.arabic_to_icr(text),
            Script::Latin => self.latin_to_icr(text),
            Script::Mixed => Err(StageError::ICRConversion(
                "Mixed script tokens should be processed separately".to_string()
            )),
            Script::Unknown => Err(StageError::ICRConversion(
                "Cannot convert unknown script to ICR".to_string()
            )),
        }
    }

    /// Convert Arabic text to ICR
    fn arabic_to_icr(&self, text: &str) -> StageResult<Vec<ICRSegment>> {
        let mut segments = Vec::new();

        for ch in text.chars() {
            let canonical = self.arabic_char_to_canonical(ch);
            let segment_type = self.classify_arabic_char(ch);

            segments.push(ICRSegment {
                canonical: canonical.to_string(),
                original: ch.to_string(),
                segment_type,
                confidence: 1.0, // High confidence for direct mapping
            });
        }

        Ok(segments)
    }

    /// Convert Latin text to ICR
    fn latin_to_icr(&self, text: &str) -> StageResult<Vec<ICRSegment>> {
        let mut segments = Vec::new();
        let chars: Vec<char> = text.chars().collect();
        let mut i = 0;

        while i < chars.len() {
            // Try to match multi-character sequences first
            let (canonical, original, advance) = if i + 1 < chars.len() {
                let two_char = format!("{}{}", chars[i], chars[i + 1]);
                if let Some(canon) = self.latin_sequence_to_canonical(&two_char) {
                    (canon, two_char, 2)
                } else {
                    let one_char = chars[i].to_string();
                    (self.latin_char_to_canonical(chars[i]), one_char, 1)
                }
            } else {
                let one_char = chars[i].to_string();
                (self.latin_char_to_canonical(chars[i]), one_char, 1)
            };

            let segment_type = self.classify_latin_char(chars[i]);

            segments.push(ICRSegment {
                canonical: canonical.to_string(),
                original,
                segment_type,
                confidence: 0.9, // Slightly lower confidence for Latin (more ambiguous)
            });

            i += advance;
        }

        Ok(segments)
    }

    /// Map Arabic character to canonical form
    fn arabic_char_to_canonical(&self, ch: char) -> &'static str {
        match ch {
            'ا' | 'أ' | 'إ' | 'آ' => "A",
            'ب' => "B",
            'ت' => "T",
            'ث' => "Ṯ", // th sound
            'ج' => "J",
            'ح' => "Ḥ", // strong h
            'خ' => "X", // kh sound
            'د' => "D",
            'ذ' => "Ḏ", // dh sound
            'ر' => "R",
            'ز' => "Z",
            'س' => "S",
            'ش' => "Š", // sh sound
            'ص' => "Ṣ", // emphatic s
            'ض' => "Ḍ", // emphatic d
            'ط' => "Ṭ", // emphatic t
            'ظ' => "Ẓ", // emphatic z
            'ع' => "ε", // ain
            'غ' => "Ġ", // gh sound
            'ف' => "F",
            'ق' => "Q",
            'ك' => "K",
            'ل' => "L",
            'م' => "M",
            'ن' => "N",
            'ه' => "H",
            'و' => "W",
            'ي' | 'ى' => "Y",
            'ة' => "T", // teh marbuta
            ' ' => " ",
            _ => "?",
        }
    }

    /// Map Latin character to canonical form
    fn latin_char_to_canonical(&self, ch: char) -> &'static str {
        match ch {
            'a' => "A",
            'b' => "B",
            't' => "T",
            'j' => "J",
            'd' => "D",
            'r' => "R",
            'z' => "Z",
            's' => "S",
            'f' => "F",
            'q' => "Q",
            'k' => "K",
            'l' => "L",
            'm' => "M",
            'n' => "N",
            'h' => "H",
            'w' => "W",
            'y' => "Y",
            'ε' | '3' => "ε", // ain
            'ḥ' | '7' => "Ḥ", // strong h
            'ġ' | 'g' => "Ġ", // gh sound
            'š' => "Š", // sh sound
            'ṣ' => "Ṣ", // emphatic s
            'ḍ' => "Ḍ", // emphatic d
            'ṭ' => "Ṭ", // emphatic t
            'ẓ' => "Ẓ", // emphatic z
            'x' => "X", // kh sound
            'e' | 'i' => "I", // vowel
            'o' | 'u' => "U", // vowel
            ' ' => " ",
            _ => "?",
        }
    }

    /// Map Latin sequences to canonical form
    fn latin_sequence_to_canonical(&self, seq: &str) -> Option<&'static str> {
        match seq {
            "sh" => Some("Š"),
            "kh" => Some("X"),
            "gh" => Some("Ġ"),
            "th" => Some("Ṯ"),
            "dh" => Some("Ḏ"),
            _ => None,
        }
    }

    /// Classify Arabic character
    fn classify_arabic_char(&self, ch: char) -> SegmentType {
        match ch {
            'ا' | 'و' | 'ي' | 'ى' => SegmentType::Vowel,
            ' ' => SegmentType::Punctuation,
            _ if ch.is_ascii_digit() => SegmentType::Number,
            _ => SegmentType::Consonant,
        }
    }

    /// Classify Latin character
    fn classify_latin_char(&self, ch: char) -> SegmentType {
        match ch {
            'a' | 'e' | 'i' | 'o' | 'u' => SegmentType::Vowel,
            ' ' => SegmentType::Punctuation,
            _ if ch.is_ascii_digit() => SegmentType::Number,
            _ => SegmentType::Consonant,
        }
    }
}

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

    #[test]
    fn test_arabic_to_icr() {
        let builder = ICRBuilder::new(Dialect::Moroccan);
        let segments = builder.arabic_to_icr("سلام").unwrap();
        assert_eq!(segments.len(), 4);
        assert_eq!(segments[0].canonical, "S");
        assert_eq!(segments[1].canonical, "L");
    }

    #[test]
    fn test_latin_to_icr() {
        let builder = ICRBuilder::new(Dialect::Moroccan);
        let segments = builder.latin_to_icr("salam").unwrap();
        assert_eq!(segments.len(), 5);
        assert_eq!(segments[0].canonical, "S");
        assert_eq!(segments[1].canonical, "A");
    }

    #[test]
    fn test_latin_sequences() {
        let builder = ICRBuilder::new(Dialect::Moroccan);
        let segments = builder.latin_to_icr("shal").unwrap();
        // "sh" should be recognized as single segment
        assert_eq!(segments[0].canonical, "Š");
    }

    #[test]
    fn test_special_chars() {
        let builder = ICRBuilder::new(Dialect::Moroccan);
        let segments = builder.latin_to_icr("3la").unwrap();
        // '3' should map to ain (ε)
        assert_eq!(segments[0].canonical, "ε");
    }
}
