from typing import Dict, List

from sparkdesk.authentication import Authentication
from sparkdesk.chat_utils import limit_chat_history
from sparkdesk.completion import ChatCompletion


class Chat:
    """
    Class for creating chat
    """

    def __init__(
        self,
        authentication: Authentication,
        app_id: str,
        user_id="default_user",
        domain="general",
        temperature=0.5,
        max_tokens=2048,
    ):
        """
        Constructor for the class.

        Args:
            authentication (Authentication):
                The authentication object used to authenticate the API calls.
            app_id (str):
                The ID of the application.
            user_id (str, optional):
                The ID of the user. Defaults to "default_user".
            domain (str, optional):
                The domain of the user. Defaults to "general".
            temperature (float, optional):
                The temperature to control the randomness of the generated output.
                Defaults to 0.5.
            max_tokens (int, optional):
                The maximum number of tokens to generate. Defaults to 2048.
        """
        self.chat_completion = ChatCompletion(
            authentication, app_id, user_id, domain, temperature, max_tokens
        )
        self.chat_history: List[Dict[str, str]] = []

    def reset_chat_context(self):
        self.chat_history = []

    def ask(self, question: str):
        """
        Appends the user's question to the chat context,
        generates responses using the chat_completion model,
        and adds the assistant's response to the chat context.

        Args:
            question (str): The user's question.

        Yields:
            dict: A dictionary representing each response generated by
                the chat_completion model.
            reference: https://www.xfyun.cn/doc/spark/Web.html
        """

        self.chat_history.append({"role": "user", "content": question})
        self.chat_history = limit_chat_history(self.chat_history)

        response_messages = []
        responses = self.chat_completion.create(self.chat_history)
        try:
            for response in responses:
                response_messages.extend(
                    [e["content"] for e in response["payload"]["choices"]["text"]]
                )
                yield response
        finally:  # GeneratorExit
            self.chat_history.append(
                {"role": "assistant", "content": "".join(response_messages)}
            )
