import json
from typing import List, Dict, Literal
import openai
from swarmauri.core.typing import SubclassUnion

from swarmauri.standard.messages.base.MessageBase import MessageBase
from swarmauri.standard.messages.concrete.AgentMessage import AgentMessage
from swarmauri.standard.llms.base.LLMBase import LLMBase


class DeepSeekModel(LLMBase):
    api_key: str
    allowed_models: List[str] = [
        "deepseek-chat",
        "deepseek-coder",
    ]
    name: str = "deepseek-chat"
    type: Literal["DeepSeekModel"] = "DeepSeekModel"

    def _format_messages(
        self, messages: List[SubclassUnion[MessageBase]]
    ) -> List[Dict[str, str]]:
        # Get only the properties that we require
        message_properties = ["content", "role"]

        # Exclude FunctionMessages
        formatted_messages = [
            message.model_dump(include=message_properties) for message in messages
        ]
        return formatted_messages

    def predict(
        self,
        conversation,
        temperature=0.7,
        max_tokens=256,
        frequency_penalty=0,
        presence_penalty=0,
        stop="\n",
        stream=False,
        top_p=1.0,
    ):

        # Create client
        client = openai.OpenAI(
            api_key=self.api_key, base_url="https://api.deepseek.com"
        )

        # Get system_context from last message with system context in it
        formatted_messages = self._format_messages(conversation.history)

        response = client.chat.completions.create(
            model=self.name,
            messages=formatted_messages,
            temperature=temperature,
            max_tokens=max_tokens,
            frequency_penalty=frequency_penalty,
            presence_penalty=presence_penalty,
            stop=stop,
            stream=stream,
            top_p=top_p,
        )

        message_content = response.choices[0].message.content
        conversation.add_message(AgentMessage(content=message_content))

        return conversation
