
from typing import List, Optional, Dict
from mlx_vlm.utils import load_image

# 外部モジュールのインポート
from .mlx_llm_interface import MlxLLMInterface
from .prompt_builder import PromptBuilder

#--------------------------------------------------------------------------------------------
# LLM推論を制御するメインクラス
#--------------------------------------------------------------------------------------------
class MlxLLM:

    def __init__(
            self,
            model_path: str, 
            use_vision: bool = False, 
            temp: float = 0.0,
            top_k: int = -1, 
            top_p: float = 1.0, 
            min_p: float = 0.0, 
            max_tokens: int = 8192
    ):
        # LLMの初期化
        self.llm = MlxLLMInterface(
            model_path=model_path,
            use_vision=use_vision,
            temp=temp,
            top_k=top_k,
            top_p=top_p,
            min_p=min_p,
            max_tokens=max_tokens
        )
    
    #-----------------------------------------------------------------------------
    # MLX/VLM形式のユーザーメッセージ構造を構築
    #-----------------------------------------------------------------------------
    def build_user_message(self, text: str, images: Optional[List] = None) -> Dict:
        
        # 画像があり、かつVLMインターフェースが有効な場合
        if images and self.llm.use_vision:
            content = [{"type": "text", "text": text}]
            for _ in images:
                content.append({"type": "image"})
            return {"role": "user", "content": content}
        
        # テキストのみの場合
        else:
            return {"role": "user", "content": text}

    #-----------------------------------------------------------------------------
    # ユーザー入力に対してツール実行と回答生成を行うメインメソッド
    #-----------------------------------------------------------------------------
    def respond(
            self,
            user_text: str, 
            user_images: List = None,
            system_prompt: str = None,
            stream: bool = False):
        
        # 初期化
        user_images = user_images or []
        
        # システムプロンプト
        prompt_builder = PromptBuilder(system_prompt_text=system_prompt)
        system_message = prompt_builder.build_system_message(
            use_vision=self.llm.use_vision
        )
        
        # ユーザー画像の処理
        final_user_image_paths = user_images
        pil_final_user_images = [load_image(str(p)) for p in final_user_image_paths]
        
        # 最終的なユーザーメッセージ
        user_message = self.build_user_message(text=user_text, images=pil_final_user_images)
        
        # すべてを結合
        messages = [system_message] + [user_message]

        # 4. 全画像リストの整理（プロンプト内での出現順序を維持）
        all_pils = []
        if self.llm.use_vision:
            # システムメッセージに画像が含まれる場合
            if "pil_images" in system_message:
                all_pils.extend(system_message["pil_images"])
            # ユーザーからの画像
            all_pils.extend(pil_final_user_images)

        # 5. 回答生成フェーズ
        full_reply = ""
        generator = self.llm.generate(
            messages, 
            images=all_pils if all_pils else None, 
            stream=stream
        )

        # ストリーミング応答
        if stream:
            for response in generator:
                chunk = response.text
                full_reply += chunk
                yield chunk
        
        # 一括応答
        else:
            for response in generator:
                full_reply += response.text
            yield full_reply