"""
Trie (Prefix Tree) - Cấu trúc dữ liệu cho prefix search

Production-ready Trie với:
- Word-start phrase indexing: Tìm được keyword ở giữa câu
- Top-K limiting: Giới hạn kết quả để tránh DFS chậm
- Normalize tiếng Việt: Bỏ dấu để tìm kiếm linh hoạt

Ví dụ cấu trúc:
         root
        /
       l
      / \
     à   a
     |   |
     m   u

"làm", "lau" được lưu trong cây trên.
"""

import unicodedata
from typing import List, Set, Tuple
from heapq import nsmallest


class TrieNode:
    """
    Node trong cây Trie
    
    Attributes:
        children (dict): Dictionary các node con {ký_tự: TrieNode}
        is_end (bool): Đánh dấu node này kết thúc một từ
        documents (list): Danh sách (doc_id, score) của documents kết thúc tại node này
        top_k (list): Top K documents phổ biến nhất (cache để tránh DFS)
    """
    
    __slots__ = ['children', 'is_end', 'documents', 'top_k']  # Tối ưu memory
    
    def __init__(self):
        self.children: dict = {}
        self.is_end: bool = False
        self.documents: List[Tuple[int, float]] = []  # [(doc_id, score), ...]
        self.top_k: List[Tuple[int, float]] = []  # Cache top K


class Trie:
    """
    Production-ready Trie cho autocomplete search
    
    Features:
        - Prefix search: O(m) với m là độ dài query
        - Word-start indexing: Tìm được keyword ở giữa câu
        - Top-K limiting: Giới hạn kết quả để performance ổn định
        - Normalize: Bỏ dấu tiếng Việt
    
    Cách hoạt động:
        Câu "làm sạch nhà vệ sinh" được index thành:
        - "làm sạch nhà vệ sinh" (prefix: làm)
        - "sạch nhà vệ sinh"     (prefix: sạch)
        - "nhà vệ sinh"           (prefix: nhà)
        - "vệ sinh"               (prefix: vệ)
        
        → User gõ "nhà" vẫn tìm được câu gốc!
    
    Example:
        >>> trie = Trie(max_results=10)
        >>> trie.insert("làm sạch nhà vệ sinh", 0, index_word_starts=True)
        >>> results = trie.search_prefix("nhà")  # Tìm được!
        >>> print(results)  # [(0, score)]
    """
    
    def __init__(self, max_results: int = 20):
        """
        Khởi tạo Trie
        
        Args:
            max_results: Số kết quả tối đa trả về (để giới hạn DFS)
        """
        self.root = TrieNode()
        self.max_results = max_results
    
    def insert(
        self, 
        text: str, 
        doc_id: int,
        score: float = 1.0,
        index_word_starts: bool = True,
        normalize: bool = True
    ) -> None:
        """
        Thêm text vào Trie với word-start indexing
        
        Args:
            text: Chuỗi cần thêm
            doc_id: ID của document
            score: Điểm ưu tiên (cao hơn = ưu tiên hơn)
            index_word_starts: Nếu True, index cả các phrase bắt đầu từ mỗi từ
            normalize: Nếu True, normalize (bỏ dấu)
        
        Example:
            >>> trie.insert("làm sạch wc", 0, index_word_starts=True)
            # Index:
            # - "làm sạch wc"
            # - "sạch wc"
            # - "wc"
        """
        if normalize:
            text = self._normalize(text)
        
        text_lower = text.lower().strip()
        
        if index_word_starts:
            # Index tất cả các phrase bắt đầu từ mỗi từ
            words = text_lower.split()
            for i in range(len(words)):
                phrase = ' '.join(words[i:])
                # Từ đầu câu (i=0) có score cao hơn từ giữa câu
                phrase_score = score if i == 0 else score * 0.5
                self._insert_phrase(phrase, doc_id, phrase_score)
        else:
            # Chỉ index câu gốc
            self._insert_phrase(text_lower, doc_id, score)
    
    def _insert_phrase(self, phrase: str, doc_id: int, score: float) -> None:
        """
        Insert một phrase vào Trie
        
        Process:
            1. Duyệt từng ký tự, tạo node nếu chưa có
            2. Đánh dấu is_end và lưu (doc_id, score)
            3. Update top_k cache trên đường đi
        """
        node = self.root
        
        for char in phrase:
            if char not in node.children:
                node.children[char] = TrieNode()
            node = node.children[char]
            
            # Update top_k cache tại mỗi node trên đường đi
            self._update_top_k(node, doc_id, score)
        
        # Đánh dấu kết thúc phrase
        node.is_end = True
        node.documents.append((doc_id, score))
        self._update_top_k(node, doc_id, score)
    
    def _update_top_k(self, node: TrieNode, doc_id: int, score: float) -> None:
        """
        Update top_k cache tại node
        
        Giữ tối đa max_results documents có score cao nhất
        """
        # Thêm document mới
        node.top_k.append((doc_id, score))
        
        # Giữ unique doc_ids (ưu tiên score cao hơn)
        seen = {}
        for did, s in node.top_k:
            if did not in seen or s > seen[did]:
                seen[did] = s
        
        # Sắp xếp theo score giảm dần, giữ top K
        node.top_k = sorted(
            [(did, s) for did, s in seen.items()],
            key=lambda x: x[1],
            reverse=True
        )[:self.max_results]
    
    def search_prefix(self, prefix: str, normalize: bool = True) -> List[Tuple[int, float]]:
        """
        Tìm documents có prefix này (NHANH với top_k cache)
        
        Args:
            prefix: Chuỗi prefix
            normalize: Có normalize không
        
        Returns:
            List [(doc_id, score), ...] đã sắp xếp theo score
        
        Time Complexity: O(m) với m là độ dài prefix (không cần DFS!)
        
        Example:
            >>> results = trie.search_prefix("làm")
            [(0, 1.0), (1, 0.9), ...]
        """
        if normalize:
            prefix = self._normalize(prefix)
        
        prefix = prefix.lower().strip()
        
        if not prefix:
            return []
        
        node = self.root
        
        # Duyệt đến node cuối của prefix
        for char in prefix:
            if char not in node.children:
                return []
            node = node.children[char]
        
        # Trả về top_k đã được cache (KHÔNG cần DFS!)
        return node.top_k[:self.max_results]
    
    def _normalize(self, text: str) -> str:
        """
        Normalize text: Bỏ dấu tiếng Việt
        
        Example:
            >>> self._normalize("Làm sạch WC")
            "lam sach wc"
        """
        # NFD: tách dấu ra khỏi chữ
        text = unicodedata.normalize('NFD', text)
        
        # Loại bỏ dấu (Mark, Nonspacing)
        text = ''.join(
            char for char in text 
            if unicodedata.category(char) != 'Mn'
        )
        
        return text.lower()
    
    def has_prefix(self, prefix: str, normalize: bool = True) -> bool:
        """
        Kiểm tra có tồn tại prefix này không
        
        Args:
            prefix: Chuỗi prefix
            normalize: Có normalize không
        
        Returns:
            True nếu tồn tại, False nếu không
        """
        if normalize:
            prefix = self._normalize(prefix)
        
        prefix = prefix.lower().strip()
        node = self.root
        
        for char in prefix:
            if char not in node.children:
                return False
            node = node.children[char]
        
        return True
    
    def count_documents(self, prefix: str, normalize: bool = True) -> int:
        """
        Đếm số documents có prefix này
        
        Args:
            prefix: Chuỗi prefix
            normalize: Có normalize không
        
        Returns:
            Số lượng documents
        """
        results = self.search_prefix(prefix, normalize)
        return len(results)
    
    def get_stats(self) -> dict:
        """
        Lấy thống kê về Trie
        
        Returns:
            Dict chứa thống kê
        """
        total_nodes = self._count_nodes_iterative(self.root)
        
        return {
            'total_nodes': total_nodes,
            'max_results': self.max_results,
            'memory_optimized': True,
            'features': ['word_start_indexing', 'top_k_cache', 'normalize']
        }
    
    def _count_nodes_iterative(self, root: TrieNode) -> int:
        """
        Đếm tổng số nodes trong Trie (iterative để tránh recursion error)
        
        Dùng stack thay vì đệ quy
        """
        if not root:
            return 0
        
        count = 0
        stack = [root]
        
        while stack:
            node = stack.pop()
            count += 1
            
            # Thêm tất cả node con vào stack
            for child in node.children.values():
                stack.append(child)
        
        return count

