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

import requests

from pilot.generater.ai_base import AIBase
from pilot.generater.generation_config import GenerationConfigDTO


class lmstudioAISingleton(AIBase):
    _instance: Optional['lmstudioAISingleton'] = 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(lmstudioAISingleton, 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._session = requests.Session()
                    self._initialized = True

    def _build_generation_config(self, gen_config: GenerationConfigDTO) -> Dict[str, Any]:
        """
        [内部方法] 将通用的 GenerationConfigDTO 转换为 OpenAI 兼容 API 特有的 JSON 结构
        """
        config: Dict[str, Any] = {}
        if gen_config.temperature is not None:
            config["temperature"] = gen_config.temperature
        if gen_config.top_p is not None:
            config["top_p"] = gen_config.top_p
        if gen_config.max_output_tokens is not None:
            config["max_tokens"] = gen_config.max_output_tokens
        return config

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

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

            req_base_url = base_url.rstrip('/') if base_url else self.base_url
            req_model = model if model else self.model_name

            messages = []
            if gen_config.system_instruction:
                messages.append({"role": "system", "content": gen_config.system_instruction})
            messages.append({"role": "user", "content": prompt})

            payload: Dict[str, Any] = {
                "model": req_model,
                "messages": messages,
                "stream": is_stream
            }

            payload.update(self._build_generation_config(gen_config))

            if kwargs:
                payload.update(kwargs)

            # ==========================================
            # 流式处理分支 (SSE 解析)
            # ==========================================
            if is_stream:
                resp = self._session.post(
                    f"{req_base_url}/chat/completions",
                    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)
                                choices = chunk_data.get("choices", [])
                                if choices:
                                    delta = choices[0].get("delta", {})
                                    chunk_text = delta.get("content", "")
                                    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:
                resp = self._session.post(
                    f"{req_base_url}/chat/completions",
                    json=payload,
                    timeout=gen_config.timeout
                )
                resp.raise_for_status()
                data = resp.json()

                content = data["choices"][0]["message"]["content"]

                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 _LMStudioChatSession(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) -> 'lmstudioAISingleton':
        return cls(model_name, base_url)


class _LMStudioChatSession:
    """
    ChatSession：维护历史上下文的多轮对话类
    """
    def __init__(self, client: lmstudioAISingleton):
        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

        messages_to_send = []
        if gen_config.system_instruction:
            messages_to_send.append({"role": "system", "content": gen_config.system_instruction})

        self._messages.append({"role": "user", "content": message})
        messages_to_send.extend(self._messages)

        payload: Dict[str, Any] = {
            "model": self._client.model_name,
            "messages": messages_to_send,
            "stream": is_stream
        }

        payload.update(self._client._build_generation_config(gen_config))

        if kwargs:
            payload.update(kwargs)

        # ==========================================
        # Chat Session 流式处理分支
        # ==========================================
        if is_stream:
            resp = self._client._session.post(
                f"{self._client.base_url}/chat/completions",
                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)
                            choices = chunk_data.get("choices", [])
                            if choices:
                                delta = choices[0].get("delta", {})
                                chunk_text = delta.get("content", "")
                                if chunk_text:
                                    full_reply += chunk_text
                        except json.JSONDecodeError:
                            continue

            self._messages.append({"role": "assistant", "content": full_reply})

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

            return _Resp(full_reply)

        # ==========================================
        # Chat Session 非流式处理分支
        # ==========================================
        else:
            resp = self._client._session.post(
                f"{self._client.base_url}/chat/completions",
                json=payload,
                timeout=gen_config.timeout
            )
            resp.raise_for_status()
            data = resp.json()

            reply = data["choices"][0]["message"]["content"]
            self._messages.append({"role": "assistant", "content": reply})

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

            return _Resp(reply)