import os
import threading
from typing import Dict, Any, Optional

# 导入官方 SDK
import google.generativeai as genai

# 导入基类和统一的 DTO
from pilot.generater.ai_base import AIBase
from pilot.generater.generation_config import GenerationConfigDTO


class GeminiAISingleton(AIBase):
    _instance: Optional['GeminiAISingleton'] = None
    _lock = threading.Lock()

    def __new__(cls, model_name: str, base_url: str = ""):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super(GeminiAISingleton, cls).__new__(cls)
                    cls._instance._initialized = False
        return cls._instance

    def __init__(self, model_name: str, base_url: str = ""):
        if not self._initialized:
            with self._lock:
                if not self._initialized:
                    self.model_name = model_name
                    # 官方 SDK 通常不需要 base_url，但如果需要走企业版代理可以配置
                    self.base_url = base_url.rstrip('/')

                    # 支持读取 gemini_key 或标准的 GEMINI_API_KEY
                    self.api_key = os.environ.get("gemini_key") or os.environ.get("GEMINI_API_KEY")
                    if not self.api_key:
                        raise ValueError("环境变量 'gemini_key' 或 'GEMINI_API_KEY' 未设置。请在运行前配置。")

                    # 初始化并配置官方 SDK
                    client_options = {}
                    if self.base_url and self.base_url != ".":
                        client_options["api_endpoint"] = self.base_url

                    genai.configure(
                        api_key=self.api_key,
                        client_options=client_options if client_options else None
                    )

                    self._initialized = True

    def _build_generation_config(self, gen_config: GenerationConfigDTO) -> genai.types.GenerationConfig:
        """
        [内部方法] 将通用的 DTO 转换为 Google SDK 专属的 GenerationConfig 对象
        """
        # 注意：SDK 中使用的是 snake_case，例如 top_p 和 max_output_tokens
        config_args = {
            "temperature": gen_config.temperature,
            "top_p": gen_config.top_p,
            "max_output_tokens": gen_config.max_output_tokens
        }

        # 目前官方 Python SDK 正在逐步完善对 thinkingConfig 的支持。
        # 如果当前版本 SDK 支持 kwargs 透传，我们可以直接放入字典。
        # (如果运行时因为 thinkingConfig 报错，可在此处将其注释掉)
        #if gen_config.thinking_budget is not None:
        #    config_args["thinking_config"] = {
        #        "include_thoughts": gen_config.include_thoughts,
        #        "thinking_budget": gen_config.thinking_budget
        #    }

        return genai.types.GenerationConfig(**config_args)

    def generate_content(
            self,
            prompt: str,
            gen_config: Optional[GenerationConfigDTO] = None,
            stream: Optional[bool] = None,
            **kwargs
    ) -> Dict[str, Any]:
        """
        统一调用接口：通过 SDK 生成 AI 回复
        """
        try:
            if gen_config is None:
                gen_config = GenerationConfigDTO()

            is_stream = stream if stream is not None else gen_config.stream

            # 在 SDK 中，system_instruction 是在实例化 Model 时传入的
            model_kwargs = {"model_name": self.model_name}
            if gen_config.system_instruction:
                model_kwargs["system_instruction"] = gen_config.system_instruction

            # 动态实例化模型 (在 SDK 中这是轻量级操作)
            model = genai.GenerativeModel(**model_kwargs)

            # 提取 API 配置
            sdk_gen_config = self._build_generation_config(gen_config)

            # 发起调用
            response = model.generate_content(
                contents=prompt,
                generation_config=sdk_gen_config,
                stream=is_stream,
                request_options={"timeout": gen_config.timeout},
                **kwargs
            )

            # ==========================================
            # 流式处理分支
            # ==========================================
            if is_stream:
                full_content = ""
                # SDK 内部已经封装好了 SSE 解析，直接迭代 response 即可
                for chunk in response:
                    if chunk.text:
                        full_content += chunk.text

                return {
                    "prompt": prompt,
                    "response": self._remove_code_fence(full_content),
                    "success": True,
                    "error": None
                }

            # ==========================================
            # 非流式处理分支
            # ==========================================
            else:
                return {
                    "prompt": prompt,
                    "response": self._remove_code_fence(response.text),
                    "success": True,
                    "error": None
                }

        except Exception as e:
            return {
                "prompt": prompt,
                "response": None,
                "success": False,
                "error": str(e)
            }

    def start_chat(self):
        """
        开启多轮对话会话
        """
        return _GeminiChatSession(self)

    def count_tokens(self, text: str) -> int:
        """
        官方 SDK 提供了精确的 token 计算方法
        """
        try:
            model = genai.GenerativeModel(self.model_name)
            response = model.count_tokens(text)
            return response.total_tokens
        except Exception:
            return 1

    def _remove_code_fence(self, text: str) -> str:
        lines = text.splitlines()
        if lines and lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].startswith("```"):
            lines = lines[:-1]
        return "\n".join(lines)

    @classmethod
    def get_instance(cls, model_name: str, base_url: str) -> 'GeminiAISingleton':
        return cls(model_name, base_url)


class _GeminiChatSession:
    """
    ChatSession：利用官方 SDK 原生支持的格式维护历史上下文
    """

    def __init__(self, client: GeminiAISingleton):
        self._client = client
        # 以 SDK 原生接受的字典格式存储历史
        self._messages = []

    def send_message(
            self,
            message: str,
            gen_config: Optional[GenerationConfigDTO] = None,
            stream: Optional[bool] = None,
            **kwargs
    ):
        if gen_config is None:
            gen_config = GenerationConfigDTO()

        is_stream = stream if stream is not None else gen_config.stream

        # 追加用户消息
        self._messages.append({
            "role": "user",
            "parts": [message]
        })

        model_kwargs = {"model_name": self._client.model_name}
        if gen_config.system_instruction:
            model_kwargs["system_instruction"] = gen_config.system_instruction

        model = genai.GenerativeModel(**model_kwargs)
        sdk_gen_config = self._client._build_generation_config(gen_config)

        # 发起调用，将全部历史记录作为 contents 传入
        response = model.generate_content(
            contents=self._messages,
            generation_config=sdk_gen_config,
            stream=is_stream,
            request_options={"timeout": gen_config.timeout},
            **kwargs
        )

        # ==========================================
        # 流式处理分支
        # ==========================================
        if is_stream:
            full_reply = ""
            for chunk in response:
                if chunk.text:
                    full_reply += chunk.text

            # 追加模型回复到历史
            self._messages.append({
                "role": "model",
                "parts": [full_reply]
            })

            class _Resp:
                def __init__(self, text):
                    self.text = text

            return _Resp(full_reply)

        # ==========================================
        # 非流式处理分支
        # ==========================================
        else:
            reply = response.text
            self._messages.append({
                "role": "model",
                "parts": [reply]
            })

            class _Resp:
                def __init__(self, text):
                    self.text = text

            return _Resp(reply)