use unicode_segmentation::UnicodeSegmentation;
use super::{StageResult, StageError};
use crate::types::{Dialect, Script};

/// Stage 3: Darija-aware Tokenization
/// Tokenizes text while respecting Darija-specific patterns
pub struct Tokenizer {
    dialect: Dialect,
}

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

    /// Tokenize text into words
    pub fn tokenize(&self, text: &str, script: Script) -> StageResult<Vec<Token>> {
        if text.is_empty() {
            return Err(StageError::Tokenization("Empty input".to_string()));
        }

        let mut tokens = Vec::new();
        let mut current_pos = 0;

        // First pass: split by whitespace
        let words = text.unicode_words();

        for word in words {
            let start = text[current_pos..].find(word).map(|p| current_pos + p).unwrap_or(current_pos);
            let end = start + word.len();

            // Detect if this is a special Darija construct
            let token_type = self.classify_token(word, script);

            tokens.push(Token {
                text: word.to_string(),
                token_type,
                start_pos: start,
                end_pos: end,
                script,
            });

            current_pos = end;
        }

        // Second pass: handle clitics and prefixes for Darija
        tokens = self.handle_clitics(tokens, script)?;

        Ok(tokens)
    }

    /// Classify a token based on its characteristics
    fn classify_token(&self, word: &str, script: Script) -> TokenType {
        // Check for numbers
        if word.chars().all(|c| c.is_ascii_digit()) {
            return TokenType::Number;
        }

        // Check for punctuation
        if word.chars().all(|c| c.is_ascii_punctuation()) {
            return TokenType::Punctuation;
        }

        // Check for common Darija function words (Articles, prepositions)
        if self.is_function_word(word, script) {
            return TokenType::FunctionWord;
        }

        // Check for Darija-specific constructs
        if self.is_darija_specific(word, script) {
            return TokenType::DarijaSpecific;
        }

        TokenType::Word
    }

    /// Check if a word is a function word (article, preposition, etc.)
    fn is_function_word(&self, word: &str, script: Script) -> bool {
        match script {
            Script::Arabic => {
                matches!(word,
                    "من" | "في" | "على" | "إلى" | "مع" | "عن" |
                    "ال" | "و" | "ف" | "ب" | "ل" | "ك"
                )
            }
            Script::Latin => {
                matches!(word.to_lowercase().as_str(),
                    "f" | "men" | "m3a" | "3la" | "l" | "b" | "w" | "li" | "dial" | "dyal"
                )
            }
            _ => false,
        }
    }

    /// Check if a word is Darija-specific
    fn is_darija_specific(&self, word: &str, script: Script) -> bool {
        match self.dialect {
            Dialect::Moroccan => {
                match script {
                    Script::Arabic => {
                        matches!(word, "غير" | "بزاف" | "شوية" | "واخا" | "كيفاش" | "شنو" | "فين")
                    }
                    Script::Latin => {
                        matches!(word.to_lowercase().as_str(),
                            "ghir" | "bzaf" | "shwiya" | "wakha" | "kifash" | "shno" | "fin" | "feen"
                        )
                    }
                    _ => false,
                }
            }
        }
    }

    /// Handle Darija clitics (prefixes and suffixes that attach to words)
    fn handle_clitics(&self, tokens: Vec<Token>, script: Script) -> StageResult<Vec<Token>> {
        let mut result = Vec::new();

        for token in tokens {
            // Only process word tokens
            if !matches!(token.token_type, TokenType::Word | TokenType::DarijaSpecific) {
                result.push(token);
                continue;
            }

            // Split clitics if needed
            let split_tokens = self.split_clitics(&token, script);
            result.extend(split_tokens);
        }

        Ok(result)
    }

    /// Split clitics from a token
    fn split_clitics(&self, token: &Token, script: Script) -> Vec<Token> {
        let word = &token.text;

        match script {
            Script::Arabic => {
                // Handle Arabic prefixes: و، ف، ب، ل، ال
                if word.starts_with("وال") && word.chars().count() > 3 {
                    let rest: String = word.chars().skip(3).collect();
                    return self.create_split_tokens(token, vec!["و", "ال"], &rest);
                } else if word.starts_with('و') && word.chars().count() > 1 {
                    let rest: String = word.chars().skip(1).collect();
                    return self.create_split_tokens(token, vec!["و"], &rest);
                } else if word.starts_with("ال") && word.chars().count() > 2 {
                    let rest: String = word.chars().skip(2).collect();
                    return self.create_split_tokens(token, vec!["ال"], &rest);
                }
            }
            Script::Latin => {
                // Handle Latin prefixes: l-, f-, b-, w-
                if word.starts_with("l") && word.len() > 1 {
                    let rest = &word[1..];
                    if rest.chars().next().map(|c| c.is_alphabetic()).unwrap_or(false) {
                        return self.create_split_tokens(token, vec!["l"], rest);
                    }
                } else if word.starts_with("f") && word.len() > 1 {
                    let rest = &word[1..];
                    if rest.chars().next().map(|c| c.is_alphabetic()).unwrap_or(false) {
                        return self.create_split_tokens(token, vec!["f"], rest);
                    }
                }
            }
            _ => {}
        }

        // No split needed
        vec![token.clone()]
    }

    /// Create split tokens from clitics
    fn create_split_tokens(&self, original: &Token, prefixes: Vec<&str>, stem: &str) -> Vec<Token> {
        let mut tokens = Vec::new();
        let mut pos = original.start_pos;

        for prefix in prefixes {
            tokens.push(Token {
                text: prefix.to_string(),
                token_type: TokenType::Clitic,
                start_pos: pos,
                end_pos: pos + prefix.len(),
                script: original.script,
            });
            pos += prefix.len();
        }

        tokens.push(Token {
            text: stem.to_string(),
            token_type: original.token_type,
            start_pos: pos,
            end_pos: original.end_pos,
            script: original.script,
        });

        tokens
    }
}

/// Represents a token in the text
#[derive(Debug, Clone, PartialEq)]
pub struct Token {
    pub text: String,
    pub token_type: TokenType,
    pub start_pos: usize,
    pub end_pos: usize,
    pub script: Script,
}

/// Types of tokens
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum TokenType {
    Word,
    Number,
    Punctuation,
    FunctionWord,
    DarijaSpecific,
    Clitic,
}

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

    #[test]
    fn test_basic_tokenization() {
        let tokenizer = Tokenizer::new(Dialect::Moroccan);
        let tokens = tokenizer.tokenize("ana kan3raf darija", Script::Latin).unwrap();
        assert_eq!(tokens.len(), 3);
        assert_eq!(tokens[0].text, "ana");
        assert_eq!(tokens[1].text, "kan3raf");
        assert_eq!(tokens[2].text, "darija");
    }

    #[test]
    fn test_function_word_detection() {
        let tokenizer = Tokenizer::new(Dialect::Moroccan);
        let tokens = tokenizer.tokenize("men l mghrib", Script::Latin).unwrap();
        assert!(matches!(tokens[0].token_type, TokenType::FunctionWord));
    }

    #[test]
    fn test_darija_specific_detection() {
        let tokenizer = Tokenizer::new(Dialect::Moroccan);
        let tokens = tokenizer.tokenize("bzaf ghir shwiya", Script::Latin).unwrap();
        assert!(matches!(tokens[0].token_type, TokenType::DarijaSpecific));
    }

    #[test]
    fn test_clitic_splitting_latin() {
        let tokenizer = Tokenizer::new(Dialect::Moroccan);
        let tokens = tokenizer.tokenize("lmghrib", Script::Latin).unwrap();
        // Should split into "l" + "mghrib"
        assert!(tokens.len() >= 2);
    }
}
