"""A module for calling Groq's Cloud API."""
import datetime
import json
from typing import Any, AsyncGenerator, ClassVar, Generator, Optional, Type, Union

from groq import AsyncGroq, Groq
from groq.types.chat.completion_create_params import Message, ResponseFormat
from tenacity import AsyncRetrying, Retrying

from ..base import BaseCall, retry
from ..enums import MessageRole
from .tools import GroqTool
from .types import GroqCallParams, GroqCallResponse, GroqCallResponseChunk
from .utils import groq_api_calculate_cost

JSON_MODE_CONTENT = """
Extract a valid JSON object instance from to content using the following schema:

{schema}
""".strip()


def _json_mode_content(tool_type: Type[GroqTool]) -> str:
    """Returns the formatted `JSON_MODE_CONTENT` with the given tool type."""
    return JSON_MODE_CONTENT.format(
        schema=json.dumps(tool_type.model_json_schema(), indent=2)
    )


class GroqCall(BaseCall[GroqCallResponse, GroqCallResponseChunk, GroqTool]):
    """A base class for calling Groq's Cloud API.

    Example:

    ```python
    from mirascope.groq import GroqCall


    class BookRecommender(GroqCall):
        prompt_template = "Please recommend a {genre} book"

        genre: str

    response = BookRecommender(genre="fantasy").call()
    print(response.content)
    #> There are many great books to read, it ultimately depends...
    ```
    """

    call_params: ClassVar[GroqCallParams] = GroqCallParams()

    def messages(self) -> list[Message]:
        """Returns the template as a formatted list of messages."""
        return [
            Message(role=message["role"], content=message["content"])
            for message in self._parse_messages(
                [MessageRole.SYSTEM, MessageRole.USER, MessageRole.ASSISTANT]
            )
        ]

    @retry
    def call(
        self, retries: Union[int, Retrying] = 1, **kwargs: Any
    ) -> GroqCallResponse:
        """Makes a call to the model using this `GroqCall` instance.

        Args:
            **kwargs: Additional keyword arguments parameters to pass to the call. These
                will override any existing arguments in `call_params`.

        Returns:
            A `GroqCallResponse` instance.
        """
        kwargs, tool_types = self._setup_groq_kwargs(kwargs)
        client = Groq(api_key=self.api_key, base_url=self.base_url)
        if self.call_params.wrapper is not None:
            client = self.call_params.wrapper(client)
        create = client.chat.completions.create
        if self.call_params.weave is not None:
            create = self.call_params.weave(
                client.chat.completions.create
            )  # pragma: no cover
        if self.call_params.logfire:
            create = self.call_params.logfire(
                create, "groq", response_type=GroqCallResponse, tool_types=tool_types
            )  # pragma: no cover
        messages = self._update_messages_if_json(self.messages(), tool_types)
        start_time = datetime.datetime.now().timestamp() * 1000
        completion = create(messages=messages, stream=False, **kwargs)
        return GroqCallResponse(
            response=completion,
            tool_types=tool_types,
            start_time=start_time,
            end_time=datetime.datetime.now().timestamp() * 1000,
            response_format=self.call_params.response_format,
            cost=groq_api_calculate_cost(completion.usage, completion.model),
        )

    @retry
    async def call_async(
        self, retries: Union[int, AsyncRetrying] = 1, **kwargs: Any
    ) -> GroqCallResponse:
        """Makes an asynchronous call to the model using this `GroqCall`.

        Args:
            **kwargs: Additional keyword arguments parameters to pass to the call. These
                will override any existing arguments in `call_params`.

        Returns:
            An `GroqCallResponse` instance.
        """
        kwargs, tool_types = self._setup_groq_kwargs(kwargs)
        client = AsyncGroq(api_key=self.api_key, base_url=self.base_url)
        if self.call_params.wrapper_async is not None:
            client = self.call_params.wrapper_async(client)
        create = client.chat.completions.create
        if self.call_params.weave is not None:
            create = self.call_params.weave(
                client.chat.completions.create
            )  # pragma: no cover
        if self.call_params.logfire_async:
            create = self.call_params.logfire_async(
                create, "groq", response_type=GroqCallResponse, tool_types=tool_types
            )  # pragma: no cover
        messages = self._update_messages_if_json(self.messages(), tool_types)
        start_time = datetime.datetime.now().timestamp() * 1000
        completion = await create(messages=messages, stream=False, **kwargs)
        return GroqCallResponse(
            response=completion,
            tool_types=tool_types,
            start_time=start_time,
            end_time=datetime.datetime.now().timestamp() * 1000,
            response_format=self.call_params.response_format,
            cost=groq_api_calculate_cost(completion.usage, completion.model),
        )

    @retry
    def stream(
        self, retries: Union[int, Retrying] = 1, **kwargs: Any
    ) -> Generator[GroqCallResponseChunk, None, None]:
        """Streams the response for a call using this `GroqCall`.

        Args:
            **kwargs: Additional keyword arguments parameters to pass to the call. These
                will override any existing arguments in `call_params`.

        Yields:
            A `GroqCallResponseChunk` for each chunk of the response.
        """
        kwargs, tool_types = self._setup_groq_kwargs(kwargs)
        client = Groq(api_key=self.api_key, base_url=self.base_url)
        if self.call_params.wrapper is not None:
            client = self.call_params.wrapper(client)
        messages = self._update_messages_if_json(self.messages(), tool_types)
        create = client.chat.completions.create
        if self.call_params.logfire:
            create = self.call_params.logfire(
                create, "groq", response_chunk_type=GroqCallResponseChunk
            )  # pragma: no cover
        stream = create(messages=messages, stream=True, **kwargs)
        for completion in stream:
            yield GroqCallResponseChunk(
                chunk=completion,
                tool_types=tool_types,
                response_format=self.call_params.response_format,
            )

    @retry
    async def stream_async(
        self, retries: Union[int, AsyncRetrying] = 1, **kwargs: Any
    ) -> AsyncGenerator[GroqCallResponseChunk, None]:
        """Streams the response for an asynchronous call using this `GroqCall`.

        Args:
            **kwargs: Additional keyword arguments parameters to pass to the call. These
                will override any existing arguments in `call_params`.

        Yields:
            A `GroqCallResponseChunk` for each chunk of the response.
        """
        kwargs, tool_types = self._setup_groq_kwargs(kwargs)
        client = AsyncGroq(api_key=self.api_key, base_url=self.base_url)
        if self.call_params.wrapper_async is not None:
            client = self.call_params.wrapper_async(client)
        messages = self._update_messages_if_json(self.messages(), tool_types)
        create = client.chat.completions.create
        if self.call_params.logfire_async:  # pragma: no cover
            create = self.call_params.logfire_async(
                create, "groq", response_chunk_type=GroqCallResponseChunk
            )
            stream = create(messages=messages, stream=True, **kwargs)
        else:
            stream = await create(messages=messages, stream=True, **kwargs)
        async for completion in stream:
            yield GroqCallResponseChunk(
                chunk=completion,
                tool_types=tool_types,
                response_format=self.call_params.response_format,
            )

    ############################## PRIVATE METHODS ###################################

    def _setup_groq_kwargs(
        self,
        kwargs: dict[str, Any],
    ) -> tuple[
        dict[str, Any],
        Optional[list[Type[GroqTool]]],
    ]:
        """Overrides the `BaseCall._setup` for Groq specific setup."""
        kwargs, tool_types = self._setup(kwargs, GroqTool)
        if (
            self.call_params.response_format == ResponseFormat(type="json_object")
            and tool_types
        ):
            kwargs.pop("tools")
        return kwargs, tool_types

    def _update_messages_if_json(
        self,
        messages: list[Message],
        tool_types: Optional[list[type[GroqTool]]],
    ) -> list[Message]:
        if (
            self.call_params.response_format == ResponseFormat(type="json_object")
            and tool_types
        ):
            messages.append(
                Message(
                    role="user", content=_json_mode_content(tool_type=tool_types[0])
                )
            )
        return messages
