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

from openai import OpenAI

# 既存の基盤クラスをインポート
from pilot.generater.ai_base import AIBase
from pilot.generater.generation_config import GenerationConfigDTO


class AzureOpenAIAISingleton(AIBase):
    """
    Azure OpenAI Service 向けの Singleton クラス
    """
    _instance: Optional['AzureOpenAIAISingleton'] = None
    _lock = threading.Lock()

    def __new__(cls, deployment_name: str, endpoint_url: str):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super(AzureOpenAIAISingleton, cls).__new__(cls)
                    cls._instance._initialized = False
        return cls._instance

    def __init__(self, deployment_name: str, endpoint_url: str):
        if not self._initialized:
            with self._lock:
                if not self._initialized:
                    self.deployment_name = deployment_name
                    self.endpoint_url = endpoint_url.rstrip('/')

                    self.api_key = os.getenv("AZURE_OPENAI_API_KEY") or os.getenv("azureopenai_key")
                    if not self.api_key:
                        raise ValueError("環境変数 'AZURE_OPENAI_API_KEY' が設定されていません。")

                    base_url = f"{self.endpoint_url}/openai/v1/"

                    self.client = OpenAI(
                        api_key=self.api_key,
                        base_url=base_url,
                        default_headers={"api-key": self.api_key}
                    )

                    self._initialized = True

    def _build_generation_config(self, gen_config: GenerationConfigDTO) -> Dict[str, Any]:
        config: Dict[str, Any] = {
            "temperature": gen_config.temperature,
            #"top_p": gen_config.top_p,
            "max_completion_tokens": 128000 #gen_config.max_output_tokens
        }
        if gen_config.thinking_budget is not None:
            config["max_completion_tokens"] = gen_config.thinking_budget

        return {k: v for k, v in config.items() if v is not None}

    def _delete_response(self, response_id: Optional[str]):
        """生成完了後にサーバー側のレスポンス保存を削除する"""
        if not response_id:
            return
        try:
            self.client.responses.delete(response_id)
        except Exception as e:
            # 削除処理の失敗でメインの応答をブロックしないようにする
            print(f"[Warning] Failed to delete response {response_id}: {e}")

    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]:
        try:
            if gen_config is None:
                gen_config = GenerationConfigDTO()

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

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

            api_kwargs = self._build_generation_config(gen_config)
            if kwargs:
                api_kwargs.update(kwargs)

            req_endpoint = base_url.rstrip('/') if base_url else self.endpoint_url
            req_model = model if model else self.deployment_name

            if base_url:
                client_to_use = OpenAI(
                    api_key=self.api_key,
                    base_url=f"{req_endpoint}/openai/v1/",
                    default_headers={"api-key": self.api_key}
                )
            else:
                client_to_use = self.client

            response = client_to_use.chat.completions.create(
                model=req_model,
                messages=messages,
                stream=is_stream,
                timeout=gen_config.timeout,
                **api_kwargs
            )

            response_id = None

            if is_stream:
                full_content = ""
                for chunk in response:
                    # ストリームの最初のチャンクからIDを取得
                    if not response_id and hasattr(chunk, 'id'):
                        response_id = chunk.id

                    if chunk.choices and chunk.choices[0].delta.content:
                        full_content += chunk.choices[0].delta.content

                # 生成完了後に削除処理を実行
                #self._delete_response(response_id)

                return {
                    "prompt": prompt,
                    "response": self._remove_code_fence(full_content),
                    "success": True,
                    "error": None
                }
            else:
                #response_id = response.id
                content = response.choices[0].message.content

                # 生成完了後に削除処理を実行
                #self._delete_response(response_id)

                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 _AzureOpenAIChatSession(self)

    def count_tokens(self, text: str) -> int:
        return len(text) // 3

    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, deployment_name: str, endpoint_url: str) -> 'AzureOpenAIAISingleton':
        return cls(deployment_name, endpoint_url)


class _AzureOpenAIChatSession:
    def __init__(self, client: AzureOpenAIAISingleton):
        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

        if not self._messages and gen_config.system_instruction:
            self._messages.append({"role": "system", "content": gen_config.system_instruction})

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

        api_kwargs = self._client._build_generation_config(gen_config)

        response = self._client.client.chat.completions.create(
            model=self._client.deployment_name,
            messages=self._messages,
            stream=is_stream,
            timeout=gen_config.timeout,
            **api_kwargs
        )

        full_reply = ""
        response_id = None

        if is_stream:
            for chunk in response:
                if not response_id and hasattr(chunk, 'id'):
                    response_id = chunk.id
                if chunk.choices and chunk.choices[0].delta.content:
                    full_reply += chunk.choices[0].delta.content
        else:
            response_id = response.id
            full_reply = response.choices[0].message.content

        # ローカルのチャット履歴には追加
        self._messages.append({"role": "assistant", "content": full_reply})

        # サーバー側の保存データは削除
        self._client._delete_response(response_id)

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

        return _Resp(full_reply)