"""Class for interacting with OpenAI through Chat APIs."""
import datetime
import logging
import warnings
from typing import Any, Callable, Generator, Optional, Type, TypeVar, Union

from openai import OpenAI
from pydantic import BaseModel, ValidationError

from ...base import BasePrompt, convert_base_model_to_tool
from ..tools import OpenAITool
from ..types import OpenAIChatCompletion, OpenAIChatCompletionChunk
from ..utils import (
    convert_tools_list_to_openai_tools,
    patch_openai_kwargs,
)

warnings.filterwarnings("always", category=DeprecationWarning, module="mirascope")
logger = logging.getLogger("mirascope")
BaseModelT = TypeVar("BaseModelT", bound=BaseModel)


class OpenAIChat:
    '''A convenience wrapper for the OpenAI Chat client.

    The Mirascope convenience wrapper for OpenAI provides a more user-friendly interface
    for interacting with their API. For more usage details, check out our examples.

    Example:

    ```python
    import os

    from mirascope import OpenAIChat, BasePrompt

    os.environ["OPENAI_API_KEY"] = "YOUR_API_KEY"


    class BookRecommendationPrompt(BasePrompt):
        """
        Can you recommend some books on {topic}?
        """

        topic: str


    prompt = BookRecommendationPrompt(topic="how to bake a cake")
    model = OpenAIChat()
    completion = model.create(prompt)

    print(completion)
    #> Certinly! Here are some books on how to bake a cake:
    #  1. "The Cake Bible" by Rose Levy Beranbaum
    #  2. "Joy of Baking" by Irma S Rombauer and Marion Rombauer Becker
    #  ...
    ```

    '''

    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        client_wrapper: Optional[Callable] = None,
        **kwargs: Any,
    ):
        """Initializes an instance of `OpenAIChat."""
        warnings.warn(
            "`OpenAIChat` is deprecated. Use `OpenAIPrompt` instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        if "model" in kwargs:
            self.model = kwargs.pop("model")
            self.model_is_set = True
        else:
            self.model = "gpt-3.5-turbo"
            self.model_is_set = False
        self.client = OpenAI(api_key=api_key, base_url=base_url, **kwargs)
        if client_wrapper is not None:
            self.client = client_wrapper(self.client)

    def create(
        self,
        prompt: Optional[Union[BasePrompt, str]] = None,
        tools: Optional[list[Union[Callable, Type[OpenAITool]]]] = None,
        **kwargs: Any,
    ) -> OpenAIChatCompletion:
        """Makes a call to the model using `prompt`.

        Args:
            prompt: The prompt to use for the call. This can either be a `BasePrompt`
                instance, a raw string, or `None`. If `prompt` is a `BasePrompt` instance,
                then the call will use the `OpenAICallParams` in the prompt before
                anything else. If `prompt` is `None`, then the call will attempt to use
                the `messages` keyword argument.
            tools: (Deprecated) A list of `OpenAITool` types or `Callable` functions
                that the creation call can decide to use. If `tools` is provided,
                `tool_choice` will be set to `auto` unless manually specified.
            **kwargs: Additional keyword arguments to pass to the API call. You can
                find available keyword arguments here:
                https://platform.openai.com/docs/api-reference/chat/create

        Returns:
            A `OpenAIChatCompletion` instance.

        Raises:
            ValueError: if neither `prompt` nor `messages` are provided.
            OpenAIError: raises any OpenAI errors, see:
                https://platform.openai.com/docs/guides/error-codes/api-errors
        """
        extract = kwargs.pop("extract") if "extract" in kwargs else False
        if isinstance(prompt, BasePrompt):
            if self.model_is_set:
                warnings.warn(
                    "The `model` parameter will be ignored when `prompt` is of type "
                    "`BasePrompt` in favor of `OpenAICallParams.model` field inside of "
                    "`prompt`; version>=0.3.0. Use `OpenAICallParams` inside of your "
                    "`BasePrompt` instead.",
                    DeprecationWarning,
                    stacklevel=2,
                )
            self.model = prompt.call_params.model

            if tools is not None and not extract:
                warnings.warn(
                    "The `tools` parameter is deprecated; version>=0.3.0. "
                    "Use `OpenAICallParams` inside of your `BasePrompt` instead.",
                    DeprecationWarning,
                    stacklevel=2,
                )
            if prompt.call_params.tools is not None:
                tools = prompt.call_params.tools

        start_time = datetime.datetime.now().timestamp() * 1000
        openai_tools = convert_tools_list_to_openai_tools(tools)
        patch_openai_kwargs(kwargs, prompt, openai_tools)
        return OpenAIChatCompletion(
            completion=self.client.chat.completions.create(
                model=self.model,
                stream=False,
                **kwargs,
            ),
            tool_types=openai_tools if tools else None,
            start_time=start_time,
            end_time=datetime.datetime.now().timestamp() * 1000,
        )

    def stream(
        self,
        prompt: Optional[Union[BasePrompt, str]] = None,
        tools: Optional[list[Union[Callable, Type[OpenAITool]]]] = None,
        **kwargs: Any,
    ) -> Generator[OpenAIChatCompletionChunk, None, None]:
        """Streams the response for a call to the model using `prompt`.

        Args:
            prompt: The `BasePrompt` to use for the call.
            tools: A list of `OpenAITool` types or `Callable` functions that the
                creation call can decide to use. If `tools` is provided, `tool_choice`
                will be set to `auto` unless manually specified.
            **kwargs: Additional keyword arguments to pass to the API call. You can
                find available keyword arguments here:
                https://platform.openai.com/docs/api-reference/chat/create

        Yields:
            A `OpenAIChatCompletionChunk` for each chunk of the response.

        Raises:
            ValueError: if neither `prompt` nor `messages` are provided.
            OpenAIError: raises any OpenAI errors, see:
                https://platform.openai.com/docs/guides/error-codes/api-errors
        """
        if isinstance(prompt, BasePrompt):
            if self.model_is_set:
                warnings.warn(
                    "The `model` parameter will be ignored when `prompt` is of type "
                    "`BasePrompt` in favor of `OpenAICallParams.model` field inside of "
                    "`prompt`; version>=0.3.0. Use `OpenAICallParams` inside of your "
                    "`BasePrompt` instead.",
                    DeprecationWarning,
                    stacklevel=2,
                )
            self.model = prompt.call_params.model

            if tools is not None:
                warnings.warn(
                    "The `tools` parameter is deprecated; version>=0.3.0. "
                    "Use `OpenAICallParams` inside of your `BasePrompt` instead.",
                    DeprecationWarning,
                    stacklevel=2,
                )
            if prompt.call_params.tools is not None:
                tools = prompt.call_params.tools

        openai_tools = convert_tools_list_to_openai_tools(tools)
        patch_openai_kwargs(kwargs, prompt, openai_tools)

        completion_stream = self.client.chat.completions.create(
            model=self.model,
            stream=True,
            **kwargs,
        )

        for chunk in completion_stream:
            yield OpenAIChatCompletionChunk(
                chunk=chunk,
                tool_types=openai_tools if tools else None,
            )

    def extract(
        self,
        schema: Type[BaseModelT],
        prompt: Optional[Union[BasePrompt, str]] = None,
        retries: int = 0,
        **kwargs: Any,
    ) -> BaseModelT:
        """Extracts the given schema from the response of a chat `create` call.

        The given schema is converted into an `OpenAITool`, complete with a description
        of the tool, all of the fields, and their types. This allows us to take
        advantage of OpenAI's tool/function calling functionality to extract information
        from a prompt according to the context provided by the `BaseModel` schema.

        Args:
            schema: The `BaseModel` schema to extract from the completion.
            prompt: The prompt from which the schema will be extracted. If `prompt` is
                a `BasePrompt` instance, then the call will use the `OpenAICallParams` in
                the prompt.
            retries: The maximum number of times to retry the query on validation error.
            **kwargs: Additional keyword arguments to pass to the API call. You can
                find available keyword arguments here:
                https://platform.openai.com/docs/api-reference/chat/create

        Returns:
            The `Schema` instance extracted from the completion.

        Raises:
            ValidationError: if the schema cannot be instantiated from the completion.
            OpenAIError: raises any OpenAI errors, see:
                https://platform.openai.com/docs/guides/error-codes/api-errors
        """
        tool = convert_base_model_to_tool(schema, OpenAITool)
        completion = self.create(
            prompt,
            tools=[tool],
            tool_choice={
                "type": "function",
                "function": {"name": tool.__name__},
            },
            extract=True,
            **kwargs,
        )

        try:
            model = schema(**completion.tool.model_dump())  # type: ignore
            model._completion = completion
            return model
        except (AttributeError, ValidationError) as e:
            if retries > 0:
                logging.info(f"Retrying due to exception: {e}")
                # TODO: update this to include failure history once prompts can handle
                # chat history properly.
                return self.extract(schema, prompt, retries - 1)
            raise  # re-raise if we have no retries left
