import os
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 AzureOpenAIAISingleton(AIBase):
    """
    Azure OpenAI Service (GPT-5-mini等) 向けの Singleton クラス
    """
    _instance: Optional['AzureOpenAIAISingleton'] = None
    _lock = threading.Lock()

    def __new__(cls, deployment_name: str, endpoint_url: str, api_version: 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, api_version: str):
        if not self._initialized:
            with self._lock:
                if not self._initialized:
                    # Azure OpenAI の設定
                    self.deployment_name = deployment_name
                    self.endpoint_url = endpoint_url.rstrip('/')
                    self.api_version = api_version  # 引数から設定

                    # 認証情報の取得
                    self.api_key = os.environ.get("azureopenai_key")
                    if not self.api_key:
                        raise ValueError(
                            "環境変数 'azureopenai_key' が設定されていません。")

                    self._session = requests.Session()
                    self._session.headers.update({
                        "Content-Type": "application/json",
                        "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_tokens": gen_config.max_output_tokens
        }
        if gen_config.thinking_budget is not None:
            config["max_completion_tokens"] = 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]:
        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})

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

            payload.update(self._build_generation_config(gen_config))
            if kwargs:
                payload.update(kwargs)

            # 引数で受け取った api_version を使用
            api_endpoint = f"{self.endpoint_url}/openai/deployments/{self.deployment_name}/chat/completions?api-version={self.api_version}"

            if is_stream:
                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').replace('data: ', '')
                        if decoded_line.strip() == "[DONE]":
                            break
                        try:
                            chunk_data = json.loads(decoded_line)
                            delta = chunk_data.get("choices", [{}])[0].get("delta", {})
                            if "content" in delta:
                                full_content += delta["content"]
                        except json.JSONDecodeError:
                            continue

                return {
                    "prompt": prompt,
                    "response": self._remove_code_fence(full_content),
                    "success": True,
                    "error": None
                }
            else:
                resp = self._session.post(api_endpoint, 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 _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, api_version: str) -> 'AzureOpenAIAISingleton':
        return cls(deployment_name, endpoint_url, api_version)


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_version を使用
        api_endpoint = f"{self._client.endpoint_url}/openai/deployments/{self._client.deployment_name}/chat/completions?api-version={self._client.api_version}"

        payload: Dict[str, Any] = {
            "messages": self._messages,
            "stream": is_stream
        }
        payload.update(self._client._build_generation_config(gen_config))

        full_reply = ""
        if is_stream:
            resp = self._client._session.post(api_endpoint, json=payload, stream=True, timeout=gen_config.timeout)
            resp.raise_for_status()
            for line in resp.iter_lines():
                if line:
                    decoded_line = line.decode('utf-8').replace('data: ', '')
                    if decoded_line.strip() == "[DONE]": break
                    try:
                        chunk = json.loads(decoded_line)
                        content = chunk.get("choices", [{}])[0].get("delta", {}).get("content", "")
                        full_reply += content
                    except:
                        continue
        else:
            resp = self._client._session.post(api_endpoint, json=payload, timeout=gen_config.timeout)
            resp.raise_for_status()
            full_reply = resp.json()["choices"][0]["message"]["content"]

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

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

        return _Resp(full_reply)