import warnings
from typing import List, Union, Dict, Optional

from openai import OpenAI, Stream
from openai.types import Embedding
from openai.types.chat import ChatCompletionChunk, ChatCompletion
from pydantic import BaseModel

from adprompt.rerank import RerankScore, RerankResponse


class OpenAIConfig(BaseModel):
    base_url: str
    api_key: str
    model: Optional[str] = None
    client_id: Optional[str] = None
    default_kwargs: Optional[dict] = None


class AIClient:

    def __init__(self, openai_config: Union[OpenAIConfig, List[OpenAIConfig]]):
        self.client_map: Dict[str, OpenAI] = {}
        self.client_default_model: Dict[str, str] = {}
        self.client_default_kwargs: Dict[str, dict] = {}

        if not isinstance(openai_config, list):
            arr = [openai_config]
        else:
            arr = openai_config
        for i in arr:
            client = OpenAI(
                api_key=i.api_key,
                base_url=i.base_url,
            )
            client_id = i.client_id or 'default'
            default_model = i.model
            self.client_map[client_id] = client
            self.client_default_model[client_id] = default_model
            self.client_default_kwargs[client_id] = i.default_kwargs or {}

    def chat(
            self,
            messages: List[dict],
            client_id: str = None,
            **kwargs,
    ) -> str:
        completion = self._chat(
            messages=messages,
            client_id=client_id,
            **kwargs,
        )
        msg = completion.choices[0].message
        if msg.content is not None:
            return msg.content
        if hasattr(msg, "reasoning") and msg.reasoning is not None:
            return msg.reasoning
        elif hasattr(msg, "reasoning_content") and msg.reasoning_content is not None:
            return msg.reasoning_content
        return ''

    def _chat(
            self,
            messages: List[dict],
            client_id: str = None,
            **kwargs,
    ) -> ChatCompletion:
        if kwargs.get('stream'):
            warnings.warn('use "stream_chat" instead')
            kwargs.pop('stream')
        client_id = client_id or 'default'
        if 'model' not in kwargs:
            kwargs['model'] = self.client_default_model[client_id]
        client = self.client_map[client_id]
        for k, v in self.client_default_kwargs[client_id].items():
            if k not in kwargs:
                kwargs[k] = v
        completion = client.chat.completions.create(
            messages=messages,
            **kwargs,
        )
        return completion

    def stream_chat(
            self,
            messages: List[dict],
            client_id: str = None,
            **kwargs,
    ) -> Stream[ChatCompletionChunk]:
        client_id = client_id or 'default'
        if 'stream' not in kwargs:
            kwargs['stream'] = True
        if 'model' not in kwargs:
            kwargs['model'] = self.client_default_model[client_id]
        client = self.client_map[client_id]
        for k, v in self.client_default_kwargs[client_id].items():
            if k not in kwargs:
                kwargs[k] = v
        completion = client.chat.completions.create(
            messages=messages,
            **kwargs,
        )
        return completion

    def embed(
            self,
            input: List[str],
            client_id: str = None,
            **kwargs,
    ) -> List[Embedding]:
        client_id = client_id or 'default'
        if 'model' not in kwargs:
            kwargs['model'] = self.client_default_model[client_id]
        client = self.client_map[client_id]
        for k, v in self.client_default_kwargs[client_id].items():
            if k not in kwargs:
                kwargs[k] = v
        resp = client.embeddings.create(
            input=input,
            **kwargs,
        ).data
        return resp

    def rerank(
            self,
            query: str,
            text: List[str],
            client_id: str = None,
            model: str = None,
    ) -> List[RerankScore]:
        client_id = client_id or 'default'
        if model is None:
            model = self.client_default_model[client_id]
        client = self.client_map[client_id]
        resp = client.post(
            '/score',
            cast_to=RerankResponse,
            body=dict(
                model=model,
                text_1=query,
                text_2=text,
            )
        )
        result = resp.data
        return result
