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

import requests

# 导入基类和统一的 DTO (请根据您的实际目录结构调整导入路径)
from pilot.generater.ai_base import AIBase
from pilot.generater.generation_config import GenerationConfigDTO


class VectorEngineAISingleton(AIBase):
    _instance: Optional['VectorEngineAISingleton'] = 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(VectorEngineAISingleton, 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
                    self.base_url = base_url.rstrip('/')

                    self.api_key = os.environ.get("vector_key")
                    if not self.api_key:
                        raise ValueError("环境变量 'vector_key' 未设置。请在运行前配置。")

                    self._session = requests.Session()
                    self._session.headers.update({
                        "Content-Type": "application/json",
                        "Authorization": f"Bearer {self.api_key}"
                    })

                    self._initialized = True

    def _build_generation_config(self, gen_config: GenerationConfigDTO) -> Dict[str, Any]:
        """
        [内部方法] 将通用的 GenerationConfigDTO 转换为 Vector Engine (Gemini格式) 特有的 JSON 结构
        """
        config: Dict[str, Any] = {
            "temperature": gen_config.temperature,
            "topP": gen_config.top_p,
            "maxOutputTokens": gen_config.max_output_tokens
        }
        # 如果设置了 thinking_budget，则组装 thinkingConfig
        if gen_config.thinking_budget is not None:
            config["thinkingConfig"] = {
                "includeThoughts": gen_config.include_thoughts,
                "thinkingBudget": gen_config.thinking_budget
            }
        return config

    def generate_content(
            self,
            prompt: str,
            gen_config: Optional[GenerationConfigDTO] = None,
            stream: Optional[bool] = None,
            **kwargs
    ) -> Dict[str, Any]:
        """
        统一调用接口：生成 AI 回复
        支持多线程安全调用，完美继承 AIBase 签名。
        """
        try:
            if gen_config is None:
                gen_config = GenerationConfigDTO()

            # 独立的 stream 参数优先于 DTO 中的设置
            is_stream = stream if stream is not None else gen_config.stream

            payload: Dict[str, Any] = {
                "contents": [
                    {
                        "role": "user",
                        "parts": [{"text": prompt}]
                    }
                ]
            }

            # 设定系统指令 (Persona / System Prompt)
            if gen_config.system_instruction:
                payload["systemInstruction"] = {
                    "parts": [{"text": gen_config.system_instruction}]
                }

            # 调用私有方法，将 DTO 映射为 API 所需的特定格式
            payload["generationConfig"] = self._build_generation_config(gen_config)

            # 透传特定 API 可能需要的额外参数 (保持接口横向扩展性)
            if kwargs:
                payload.update(kwargs)

            # ==========================================
            # 流式处理分支 (SSE 解析)
            # ==========================================
            if is_stream:
                api_endpoint = f"{self.base_url}/v1beta/models/{self.model_name}:streamGenerateContent?alt=sse"

                resp = self._session.post(
                    api_endpoint,
                    json=payload,
                    stream=True,
                    timeout=gen_config.timeout  # 配置文件传导下来的超时时间
                )
                resp.raise_for_status()

                full_content = ""
                for line in resp.iter_lines():
                    if line:
                        decoded_line = line.decode('utf-8')
                        if decoded_line.startswith("data: "):
                            data_str = decoded_line[6:]
                            if data_str.strip() == "[DONE]":
                                break
                            try:
                                chunk_data = json.loads(data_str)
                                candidates = chunk_data.get("candidates", [])
                                if candidates:
                                    parts = candidates[0].get("content", {}).get("parts", [])
                                    chunk_text = "".join([part.get("text", "") for part in parts])
                                    if chunk_text:
                                        full_content += chunk_text
                            except json.JSONDecodeError:
                                continue

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

            # ==========================================
            # 非流式处理分支
            # ==========================================
            else:
                api_endpoint = f"{self.base_url}/v1beta/models/{self.model_name}:generateContent"

                resp = self._session.post(
                    api_endpoint,
                    json=payload,
                    timeout=gen_config.timeout  # 配置文件传导下来的超时时间
                )
                resp.raise_for_status()
                data = resp.json()

                content_parts = data["candidates"][0]["content"]["parts"]
                content = "".join([part.get("text", "") for part in content_parts])

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

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

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

    def count_tokens(self, text: str) -> int:
        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) -> 'VectorEngineAISingleton':
        return cls(model_name, base_url)


class _VectorEngineChatSession:
    """
    ChatSession：维护历史上下文的多轮对话类
    """

    def __init__(self, client: VectorEngineAISingleton):
        self._client = client
        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": [{"text": message}]
        })

        payload: Dict[str, Any] = {
            "contents": self._messages
        }

        if gen_config.system_instruction:
            payload["systemInstruction"] = {
                "parts": [{"text": gen_config.system_instruction}]
            }

        # 调用客户端绑定的映射方法
        payload["generationConfig"] = self._client._build_generation_config(gen_config)

        if kwargs:
            payload.update(kwargs)

        # ==========================================
        # Chat Session 流式处理分支
        # ==========================================
        if is_stream:
            api_endpoint = f"{self._client.base_url}/v1beta/models/{self._client.model_name}:streamGenerateContent?alt=sse"

            resp = self._client._session.post(
                api_endpoint,
                json=payload,
                stream=True,
                timeout=gen_config.timeout
            )
            resp.raise_for_status()

            full_reply = ""
            for line in resp.iter_lines():
                if line:
                    decoded_line = line.decode('utf-8')
                    if decoded_line.startswith("data: "):
                        data_str = decoded_line[6:]
                        if data_str.strip() == "[DONE]":
                            break
                        try:
                            chunk_data = json.loads(data_str)
                            candidates = chunk_data.get("candidates", [])
                            if candidates:
                                parts = candidates[0].get("content", {}).get("parts", [])
                                chunk_text = "".join([part.get("text", "") for part in parts])
                                if chunk_text:
                                    full_reply += chunk_text
                        except json.JSONDecodeError:
                            continue

            self._messages.append({
                "role": "model",
                "parts": [{"text": full_reply}]
            })

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

            return _Resp(full_reply)

        # ==========================================
        # Chat Session 非流式处理分支
        # ==========================================
        else:
            api_endpoint = f"{self._client.base_url}/v1beta/models/{self._client.model_name}:generateContent"

            resp = self._client._session.post(
                api_endpoint,
                json=payload,
                timeout=gen_config.timeout
            )
            resp.raise_for_status()
            data = resp.json()

            content_parts = data["candidates"][0]["content"]["parts"]
            reply = "".join([part.get("text", "") for part in content_parts])

            self._messages.append({
                "role": "model",
                "parts": [{"text": reply}]
            })

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

            return _Resp(reply)