"""
模型适配器
"""
from .client import OpenAIClient
from openai import OpenAI
import os

class ModelAdapter:
    def __init__(self):
        self.api_key = None
        self.model_name = None
        self.temperature = 0.7

    def set_model(self, model_name: str):
        self.model_name = model_name

    def product(self, prompt: str) -> str:
        raise NotImplementedError("This method should be implemented by subclasses")

    def chat(self, messages: list) -> str:
        raise NotImplementedError("This method should be implemented by subclasses")


class BianXieAdapter(ModelAdapter):
    """BianXie格式的适配器
    """
    def __init__(self, api_key: str = None, api_base: str = "https://api.bianxie.ai/v1/chat/completions"):
        """初始化

        Args:
            api_key (str): API key for authentication.
            api_base (str): Base URL for the API endpoint.
        """
        super().__init__()
        self.client = OpenAIClient(api_key=api_key or os.getenv('BIANXIE_API_KEY') , api_base=api_base)

    def get_model(self):
        return ["详见 看见"]

    def product(self, prompt: str) -> str:
        """Generate a response from the model based on a single prompt.

        Args:
            prompt (str): The input text prompt to generate a response for.

        Returns:
            str: The response generated by the model.
        """
        data = {
            'model': self.model_name,
            'messages': [{'role': 'user', 'content': prompt}],
            'temperature': self.temperature
        }
        return self.client.request(data).get('choices')[0].get('message').get('content')

    def chat(self, messages: list) -> str:
        """Engage in a conversation with the model using a list of messages.

        Args:
            messages (list): A list of message dictionaries, each containing a role and content.

        Returns:
            str: The response generated by the model for the conversation.
        """
        data = {
            'model': self.model_name,
            'messages': messages,
            'temperature': self.temperature
        }
        return self.client.request(data).get('choices')[0].get('message').get('content')

class KimiAdapter(ModelAdapter):
    """Kimi格式的适配器

    """
    def __init__(self, api_key: str = None, api_base: str = "https://api.moonshot.cn/v1",):
        """初始化

        Args:
            api_key (str): API key for authentication.
            api_base (str): Base URL for the API endpoint.
        """
        super().__init__()
        
        self.client = OpenAI(api_key=api_key or os.getenv('MOONSHOT_API_KEY') , base_url=api_base)

    def get_model(self):
        return ["moonshot-v1-128k","moonshot-v1-128k","moonshot-v1-128k"]

    def product(self, prompt: str) -> str:
        """Generate a response from the model based on a single prompt.

        Args:
            prompt (str): The input text prompt to generate a response for.

        Returns:
            str: The response generated by the model.
        """

        data = {
            'model': self.model_name,
            'messages': [{'role': 'user', 'content': prompt}],
            'temperature': self.temperature
        }
        return self.client.chat.completions.create(**data).choices[0].message.content

    def chat(self, messages: list) -> str:
        """Engage in a conversation with the model using a list of messages.

        Args:
            messages (list): A list of message dictionaries, each containing a role and content.

        Returns:
            str: The response generated by the model for the conversation.
        """
        data = {
            'model': self.model_name,
            'messages': messages,
            'temperature': self.temperature
        }
        return self.client.chat.completions.create(**data).choices[0].message.content

class ArkAdapter(ModelAdapter):
    def __init__(self, api_key: str = None, api_base: str = None,):
        """初始化

        Args:
            api_key (str): API key for authentication.
            api_base (str): Base URL for the API endpoint.
        """
        super().__init__()

        from volcenginesdkarkruntime import Ark

        self.client = Ark(api_key=api_key or os.getenv('ARK_API_KEY'))

    def get_model(self):
        return ["doubao-1-5-pro-256k-250115"]

    def product(self, prompt: str) -> str:
        """Generate a response from the model based on a single prompt.

        Args:
            prompt (str): The input text prompt to generate a response for.

        Returns:
            str: The response generated by the model.
        """
        data = {
            'model': self.model_name,
            'messages': [{'role': 'user', 'content': prompt}],
            'temperature': self.temperature
        }
        return self.client.chat.completions.create(**data).choices[0].message.content

    def chat(self, messages: list) -> str:
        """Engage in a conversation with the model using a list of messages.

        Args:
            messages (list): A list of message dictionaries, each containing a role and content.

        Returns:
            str: The response generated by the model for the conversation.
        """
        data = {
            'model': self.model_name,
            'messages': messages,
            'temperature': self.temperature
        }
        return self.client.chat.completions.create(**data).choices[0].message.content