from abc import ABC
from abc import abstractmethod


class Provider(ABC):
    """Base class for providers"""

    @abstractmethod
    def ask(
        self,
        prompt: str,
        stream: bool = False,
        raw: bool = False,
        optimizer: str = None,
        conversationally: bool = False,
    ) -> dict:
        """Chat with AI

        Args:
            prompt (str): Prompt to be sent
            stream (bool, optional): Flag for streaming response. Defaults to False.
            raw (bool, optional): Stream back raw response as received
            optimizer (str, optional): Prompt optimizer name - `[code, shell_command]`
            conversationally (bool, optional): Chat conversationally when using optimizer. Defaults to False.
        Returns:
           dict : {}
        ```json
        {
            "completion": "\nNext: domestic cat breeds with short hair >>",
            "stop_reason": null,
            "truncated": false,
            "stop": null,
            "model": "llama-2-13b-chat",
            "log_id": "cmpl-3kYiYxSNDvgMShSzFooz6t",
            "exception": null
        }
        ```
        """
        raise NotImplementedError("Method needs to be implemented in subclass")

    @abstractmethod
    def chat(
        self,
        prompt: str,
        stream: bool = False,
        optimizer: str = None,
        conversationally: bool = False,
    ) -> str:
        """Generate response `str`
        Args:
            prompt (str): Prompt to be sent
            stream (bool, optional): Flag for streaming response. Defaults to False.
            optimizer (str, optional): Prompt optimizer name - `[code, shell_command]`
            conversationally (bool, optional): Chat conversationally when using optimizer. Defaults to False.
        Returns:
            str: Response generated
        """
        raise NotImplementedError("Method needs to be implemented in subclass")

    @abstractmethod
    def get_message(self, response: dict) -> str:
        """Retrieves message only from response

        Args:
            response (dict): Response generated by `self.ask`

        Returns:
            str: Message extracted
        """
        raise NotImplementedError("Method needs to be implemented in subclass")


class AsyncProvider(ABC):
    """Asynchronous base class for providers"""

    @abstractmethod
    async def ask(
        self,
        prompt: str,
        stream: bool = False,
        raw: bool = False,
        optimizer: str = None,
        conversationally: bool = False,
    ) -> dict:
        """Asynchronously chat with AI

        Args:
            prompt (str): Prompt to be sent
            stream (bool, optional): Flag for streaming response. Defaults to False.
            raw (bool, optional): Stream back raw response as received
            optimizer (str, optional): Prompt optimizer name - `[code, shell_command]`
            conversationally (bool, optional): Chat conversationally when using optimizer. Defaults to False.
        Returns:
           dict : {}
        ```json
        {
            "completion": "\nNext: domestic cat breeds with short hair >>",
            "stop_reason": null,
            "truncated": false,
            "stop": null,
            "model": "llama-2-13b-chat",
            "log_id": "cmpl-3kYiYxSNDvgMShSzFooz6t",
            "exception": null
        }
        ```
        """
        raise NotImplementedError("Method needs to be implemented in subclass")

    @abstractmethod
    async def chat(
        self,
        prompt: str,
        stream: bool = False,
        optimizer: str = None,
        conversationally: bool = False,
    ) -> str:
        """Asynchronously generate response `str`
        Args:
            prompt (str): Prompt to be sent
            stream (bool, optional): Flag for streaming response. Defaults to False.
            optimizer (str, optional): Prompt optimizer name - `[code, shell_command]`
            conversationally (bool, optional): Chat conversationally when using optimizer. Defaults to False.
        Returns:
            str: Response generated
        """
        raise NotImplementedError("Method needs to be implemented in subclass")

    @abstractmethod
    async def get_message(self, response: dict) -> str:
        """Asynchronously retrieves message only from response

        Args:
            response (dict): Response generated by `self.ask`

        Returns:
            str: Message extracted
        """
        raise NotImplementedError("Method needs to be implemented in subclass")