package/word_processor.py

from typing import List, Tuple, Any, Optional
from dataclasses import dataclass
from enum import Enum
import os
import math
from word_counter import WordCounter
from config import Config


class WordState(Enum):
    """単語の状態"""
    INCOMPLETE = "incomplete"  # 未完成
    COMPLETE = "complete"      # 完成
    TRIGGER = "trigger"        # トリガー（次語の開始）

class KList:
    def __init__(self, num: int):
        self.num = num
        self.list: List[Any] = []

    def check_k(self) -> None:
        if len(self.list) >= self.num:
            self.list.sort(key=lambda x: x.probability, reverse=True)
            self.list = self.list[:self.num]
        else:
            self.list.sort(key=lambda x: x.probability, reverse=True)

    def add(self, piece_word: Any) -> None:
        # 重複チェック: 同じテキストのピースが既に存在するか確認
        new_text = piece_word.get_full_text()
        for existing_piece in self.list:
            if existing_piece.get_full_text() == new_text:
                # 既存のピースに確率を足す
                existing_piece.probability += piece_word.probability
                # 確率を更新したので、ソートし直す
                self.check_k()
                return
        
        # 重複がない場合は追加
        self.list.append(piece_word)
        self.check_k()

    def pop(self) -> Any:
        if self.list:
            return self.list.pop(0)
        raise IndexError("List is empty")

    def empty(self) -> bool:
        return len(self.list) == 0
        
@dataclass
class WordPiece:
    """単語のピース（部分）"""
    text: str              # ピースのテキスト
    probability: float     # 確率
    next_tokens: Optional[List[Tuple[str, float]]] = None  # 次のトークン候補
    parent: Optional['WordPiece'] = None  # 親ピース
    children: List['WordPiece'] = None   # 子ピース
    
    def __post_init__(self):
        if self.children is None:
            self.children = []
    
    def get_full_text(self) -> str:
        """ルートからこのピースまでの完全なテキストを取得"""
        pieces = []
        current = self
        while current is not None:
            if current.text:
                pieces.append(current.text)
            current = current.parent
        return "".join(reversed(pieces))
    
    def get_full_word(self) -> str:
        """ルートの次語からこのピースまでの完全な単語を取得"""
        pieces = []
        current = self
        while current is not None:
            if current.text:
                pieces.append(current.text)
            current = current.parent
        reversed_pieces = reversed(pieces[:-1])
        return "".join(reversed_pieces)

    def add_child(self, text: str, probability: float, next_tokens: Optional[List[Tuple[str, float]]] = None) -> 'WordPiece':
        """子ピースを追加"""
        child = WordPiece(
            text=text,
            probability=probability,
            next_tokens=next_tokens,
            parent=self
        )
        self.children.append(child)
        return child
    
    def is_leaf(self) -> bool:
        """葉ノードかどうか"""
        return len(self.children) == 0
    
    def get_depth(self) -> int:
        """ルートからの深さを取得"""
        depth = 0
        current = self.parent
        while current is not None:
            depth += 1
            current = current.parent
        return depth


class WordDeterminer:
    """単語確定システム（ストリーミング向けリアルタイムアルゴリズム）"""
    
    def __init__(self, word_counter: WordCounter = None):
        """
        初期化
        
        Args:
            word_counter: WordCounterインスタンス（Noneの場合はデフォルトを使用）
        """
        self.word_counter = word_counter or WordCounter()
    
    def is_boundary_char(self, char: str) -> bool:
        """境界文字かどうかを判定（fugashi使用）"""
        if not char:
            return False
        
        # 空白文字
        if char.isspace():
            return True
        
        # 句読点
        punctuation = ",，.．。！？!?:；;、\n\t"
        return char in punctuation
    
    def is_word_boundary(self, text: str, position: int) -> bool:
        """
        WordCounterを使用して単語境界を判定
        
        Args:
            text: テキスト
            position: 位置（負の値で末尾から指定可能）
            
        Returns:
            bool: 単語境界かどうか
        """
        return self.word_counter.is_word_boundary(text, position)
    
    def check_word_completion(self, piece: WordPiece, root_count: int, model: Any = None) -> Tuple[WordState, Optional[Any]]:
        """
        ストリーミング向けリアルタイム単語決定アルゴリズム
        
        戦略:
        1. 確率エントロピー: 次のトークンの不確実性を測定
        2. 確率重み付き境界検出: 高確率トークンの挙動を重視
        3. 信頼度ベース判定: 高確率トークンが明確に境界を示す場合のみ確定
        
        アルゴリズム:
        - エントロピーが低い（確率が集中）→ 単語継続の可能性が高い
        - エントロピーが高い（確率が分散）→ 単語境界の可能性
        - 高確率トークンが境界を示す → 確定
        - 低確率トークンだけが境界を示す → 無視
        
        Args:
            piece: チェックするピース
            root_count: ルートテキストの単語数
            model: LLMモデル（必要に応じて）
            
        Returns:
            Tuple[WordState, Optional[Any]]: (状態, ペイロード)
        """
        full_text = piece.get_full_text()
        
        # next_tokensを取得
        if not piece.next_tokens:
            if model:
                piece.next_tokens = self._get_next_tokens_from_model(model, full_text)
            else:
                return (WordState.COMPLETE, None)
        
        if not piece.next_tokens:
            return (WordState.COMPLETE, None)
        
        # 確率順にソート（念のため）
        sorted_tokens = sorted(piece.next_tokens, key=lambda x: x[1], reverse=True)
        # sorted_tokens = piece.next_tokens
        if piece.get_full_word()[-1] in ["(","「","（","【","〈","《","［","｛","｟","《","［","｛","｟","《","［","｛","｟","《","［","｛","｟","《","［","｛","｟"]:
            return (WordState.INCOMPLETE, None)
        if piece.get_full_word()[-1] in [")","]","}","》","〉","》","］","｝","｠","》","］","｝","｠","》","］","｝","｠","》","］","｝","｠","》","］","｝","｠"]:
            return (WordState.COMPLETE, None)
        
        # 2.全トークンの挙動を確認
        count = max(1, len(sorted_tokens) )
        tokens = sorted_tokens[:count]
        
        boundary_prob = 0.0  # 境界を示すトークンの確率合計
        continuation_prob = 0.0  # 継続を示すトークンの確率合計
        total = sum(prob for _, prob in tokens)
        
        for token, prob in tokens:
            test_text = full_text + token
            test_word_count = self._count_words(test_text)
            
            # 単語数がより多く増えた場合のみ境界と判定（まとまりを上げる）
            if test_word_count > root_count + 1:
                boundary_prob += prob
            else:
                continuation_prob += prob
        
        # 3. 判定ロジック
        if total > 0:
            boundary_ratio = boundary_prob / total
            
            # トークンの多くが境界を示す場合 → 確定（閾値を上げてまとまりを上げる）
            if boundary_ratio > 0.85:
                return (WordState.COMPLETE, None)
            
            # トークンの多くが継続を示す場合 → 継続（閾値を下げて継続しやすく）
            if boundary_ratio < 0.2:
                return (WordState.INCOMPLETE, None)
                # 1. 確率エントロピーを計算
        probs = [prob for _, prob in sorted_tokens]
        entropy = -sum(p * math.log(p + 1e-10) for p in probs if p > 0)
        max_entropy = math.log(len(sorted_tokens)) if len(sorted_tokens) > 1 else 1.0
        normalized_entropy = entropy / max_entropy if max_entropy > 0 else 0
        # 4. エントロピーベース判定
        # エントロピーが低い（確率が集中）→ 単語継続の可能性
        # エントロピーが高い（確率が分散）→ 単語境界の可能性

        return (WordState.INCOMPLETE, None)
    
    def _count_words(self, text: str) -> int:
        """
        WordCounterを使用してテキストの単語数をカウント
        
        Args:
            text: カウントするテキスト
            
        Returns:
            int: 単語数
        """
        return self.word_counter.count_words(text)
    
    def _get_next_tokens_from_model(self, model: Any, text: str, top_k: int = 5) -> List[Tuple[str, float]]:
        """
        モデルから次のトークン候補を取得（常駐AIモデルを使用）
        
        Args:
            model: LLMモデル（パス文字列またはモデルオブジェクト）
            text: 入力テキスト
            top_k: 取得する候補数
            
        Returns:
            List[Tuple[str, float]]: (トークン, 確率)のリスト
        """
        try:
            # AIクラスをインポート
            from ai import AI
            
            # モデルパスを取得
            if isinstance(model, str):
                model_path = model
            elif hasattr(model, 'model_path'):
                model_path = model.model_path
            else:
                # デフォルトモデルを使用
                model_path = None
            
            # 常駐AIモデルを使用
            ai_model = AI.get_model(model_path)
            return ai_model.get_token_probabilities(text, top_k)
                
        except Exception as e:
            print(f"モデルからのトークン取得に失敗: {e}")
        
        return []
    
    def expand_piece(self, piece: WordPiece, model: Any = None) -> List[WordPiece]:
        """
        ピースを展開して子ピースを生成
        
        Args:
            piece: 展開するピース
            model: LLMモデル
            
        Returns:
            List[WordPiece]: 生成された子ピースのリスト
        """
        children = []
        full_text = piece.get_full_text()
        #1#print(f"[WORD_PROCESSOR_STREAMING] expand_piece: '{full_text}'")
        
        if piece.next_tokens:
            # 既存のnext_tokensを使用
            #1#print(f"[WORD_PROCESSOR_STREAMING] Using existing next_tokens: {len(piece.next_tokens)}")
            for token, prob in piece.next_tokens:
                # 空文字列トークンを無視
                if not token:
                    continue
                child_prob = piece.probability * prob
                child = piece.add_child(token, child_prob)
                children.append(child)
        elif model:
            # モデルから次のトークンを取得
            #1#print(f"[WORD_PROCESSOR_STREAMING] Getting tokens from model for: '{full_text}'")
            next_tokens = self._get_next_tokens_from_model(model, full_text)
            #1#print(f"[WORD_PROCESSOR_STREAMING] Got {len(next_tokens)} tokens from model")
            
            if next_tokens:
                piece.next_tokens = next_tokens
                for token, prob in next_tokens:
                    # 空文字列トークンを無視
                    if not token:
                        continue
                    child_prob = piece.probability * prob
                    child = piece.add_child(token, child_prob)
                    children.append(child)
        else:
            print(f"[WORD_PROCESSOR_STREAMING] No model provided for expansion")
        
        #1#print(f"[WORD_PROCESSOR_STREAMING] Generated {len(children)} children")
        return children
    
    def build_word_tree(self, prompt_text: str, root_text: str, model: Any, top_k: int = 5, max_depth: int = 10) -> List[WordPiece]:
        """
        単語ツリーを構築
        
        Args:
            root_text: ルートテキスト
            model: LLMモデル
            top_k: 取得する候補数
            max_depth: 最大深さ
            
        Returns:
            List[WordPiece]: 完成した単語ピースのリスト
        """
        #1#print(f"[WORD_PROCESSOR_STREAMING] build_word_tree called: prompt='{prompt_text}', root='{root_text}', top_k={top_k}")
        
        # ルートピースを作成
        root = WordPiece(text=self.build_chat_prompt(prompt_text, )+root_text, probability=1.0)
        #1#print(f"[WORD_PROCESSOR_STREAMING] Root piece created: '{root.get_full_text()}'")
        
        # 優先度付きキュー（確率順）
        candidates = KList(2*top_k)
        completed = []
        iteration = 0
        max_iterations = 1000
        children = self.expand_piece(root, model)
        #1#print(f"[WORD_PROCESSOR_STREAMING] Initial children: {len(children)}")
        for child in children:
            candidates.add(child)
        while not candidates.empty() and iteration < max_iterations and len(completed) < top_k:
            iteration += 1
            
            # 最も確率の高い候補を取得
            current = candidates.pop()
            
            # # 深さ制限チェック
            # if current.get_depth() >= max_depth:
            #     completed.append(current)
            #     continue
            
            # 単語完成状態をチェック
            root_count = self._count_words(root.get_full_text())
            state, payload = self.check_word_completion(current, root_count, model)
            
            if state == WordState.COMPLETE:
                completed.append(current)
                # print(f"☆☆☆☆☆complete: {current.get_full_text()}")
            elif state == WordState.INCOMPLETE:
                # ピースを展開
                children = self.expand_piece(current, model)
                if len(children) == 0:
                    # 子が生成できない場合、ピースを完成として扱う（無限ループ防止）
                    print(f"[WORD_PROCESSOR_STREAMING] No children generated for '{current.get_full_text()}', marking as COMPLETE")
                    completed.append(current)
                else:
                    for child in children:
                        candidates.add(child)
        # print(f"☆☆☆☆☆while end{len(completed),candidates.empty(),iteration}")
        
        # 確率で正規化
        total_prob = sum(p.probability for p in completed)
        if total_prob > 0:
            for piece in completed:
                piece.probability = piece.probability / total_prob
        
        return completed[:top_k]
        
    def build_chat_prompt(self, user_content: str, 
        system_content: str = "あなたは親切で役に立つAIアシスタントです。簡潔な回答をしてください") -> str:
        """
        チャットプロンプトを構築
        
        注意: Rust側で既に整形済みのプロンプトが渡される場合は、そのまま返す
        後方互換性のため、単一のuser_contentが渡された場合は従来の形式で整形
        """
        # Rust側で既に整形済みのプロンプトが渡されている場合（複数行、ヘッダーを含む）
        # そのまま返す
        if "<|start_header_id|>" in user_content or "<|eot_id|>" in user_content:
            return user_content
        
        # 後方互換性: 単一のuser_contentが渡された場合の従来の形式
        prompt_text = (
            f"<|begin_of_text|>"
            f"<|start_header_id|>system<|end_header_id|>\n"
            f"{system_content}\n<|eot_id|>"
            f"<|start_header_id|>user<|end_header_id|>\n"
            f"{user_content}\n<|eot_id|>"
            f"<|start_header_id|>assistant<|end_header_id|>\n"
        )
        
        # BOS(<|begin_of_text|>) の重複を抑止: 先頭のBOSを全て除去
        # llama-cpp 側でBOSが自動付与されるため、ここでは付与しない
        BOS = "<|begin_of_text|>"
        s = prompt_text.lstrip()
        while s.startswith(BOS):
            s = s[len(BOS):]
        prompt_text = s
        return prompt_text
    
if __name__ == "__main__":
    """WordDeterminerのテスト（ストリーミング版）"""
    print("=== WordDeterminerテスト（ストリーミング版） ===")
    
    try:
        # WordDeterminerを初期化
        determiner = WordDeterminer()
        
        # プロンプト設定
        prompt_text = "電球を作ったのは誰？"
        root_text = ""
        
        print(f"プロンプト: '{prompt_text}'")
        print(f"ルートテキスト: '{root_text}'")

        print("\nAIモデルテスト:")
        prompt_text = "電球を作ったのは誰？"
        root_text = "電球を作ったのは候補1：トマス"
        try:
            from ai import AI
            
            # モデルを取得
            model = AI.get_model()
            print(f"モデル取得成功: {type(model)}")
            
            # トークン確率取得テスト
            test_text = prompt_text
            tokens = model.get_token_probabilities(test_text, k=5)
            print(f"トークン確率 ({test_text}): {tokens}")
            
            # 単語ツリー構築テスト
            print("\n単語ツリー構築テスト:")
            completed_pieces = determiner.build_word_tree(
                prompt_text=prompt_text,
                root_text=root_text,
                model=model,
                top_k=3,
                max_depth=5
            )
            
            print(f"完成したピース数: {len(completed_pieces)}")
            for i, piece in enumerate(completed_pieces):
                full_text = piece.get_full_text()
                print(f"  ピース{i+1}: '{full_text}' (確率: {piece.probability:.4f})")
            
        except Exception as e:
            print(f"AIモデルテスト失敗: {e}")

        # 単語数カウントテスト
        print("\n単語数カウントテスト:")
        test_texts = [
            "電球",
            "電球を作った",
            "電球を作ったのは",
            "電球を作ったのは誰",
            "電球を作ったのは誰？"
        ]
        
        for text in test_texts:
            word_count = determiner._count_words(text)
            tokens = determiner._get_next_tokens_from_model(model, text)
            print(f"  '{text}' → {word_count}語: {tokens}")
        
        # 単語確定テスト
        print("\n単語確定テスト:")
        test_sequence = ["電球", "電球を", "電球を作", "電球を作った", "電球を作ったの", "電球を作ったのは"]
        prev_count = 0
        
        for text in test_sequence:
            current_count = determiner._count_words(text)
            if current_count > prev_count:
                print(f"  '{text}' → {current_count}語 (確定!)")
                prev_count = current_count
            else:
                print(f"  '{text}' → {current_count}語 (継続)")
        
        # 境界文字テスト
        print("\n境界文字テスト:")
        test_chars = [" ", "？", "、", "。", "a", "1"]
        for char in test_chars:
            is_boundary = determiner.is_boundary_char(char)
            print(f"  '{char}': {is_boundary}")
        
        # ピース作成テスト
        print("\nピース作成テスト:")
        root = WordPiece(text="電球", probability=1.0)
        child1 = root.add_child("を", 0.6)
        child2 = root.add_child("の", 0.3)
        
        print(f"ルートテキスト: {root.get_full_text()}")
        print(f"子1テキスト: {child1.get_full_text()}")
        print(f"子2テキスト: {child2.get_full_text()}")
        
        print("\nテスト完了")
        
    except ImportError as e:
        print(f"必要なライブラリがインストールされていません: {e}")
    except Exception as e:
        print(f"テストエラー: {e}")