"""Base LLM abstract class module"""

from abc import ABC, abstractmethod
from typing import Dict, List, Union

from pydantic import BaseModel

from synthgenai.schemas.config import LLMConfig
from synthgenai.schemas.messages import InputMessage


class BaseLLM(ABC):
    """Base abstract class for LLM implementations"""

    def __init__(self, llm_config: LLMConfig):
        """
        Initialize the base LLM class

        Args:
            llm_config (LLMConfig): The configuration for the LLM
        """
        self.model = llm_config.model
        self.temperature = llm_config.temperature
        self.top_p = llm_config.top_p
        self.max_tokens = llm_config.max_tokens
        self.api_base = (
            str(llm_config.api_base) if llm_config.api_base is not None else None
        )
        self.api_key = llm_config.api_key
        self.response_format = None

    @abstractmethod
    def _check_allowed_models(self) -> None:
        """Check if the model is allowed"""
        pass

    @abstractmethod
    def _check_llm_api_keys(self) -> None:
        """Check if the required API keys are set"""
        pass

    @abstractmethod
    def set_temperature(self, temperature: Union[float, None]) -> None:
        """
        Set the temperature for the LLM

        Args:
            temperature (Union[float, None]): The temperature to set
        """
        pass

    @abstractmethod
    def get_temperature(self) -> Union[float, None]:
        """
        Get the temperature for the LLM

        Returns:
            Union[float, None]: The temperature if set, None otherwise
        """
        pass

    @abstractmethod
    def get_model(self) -> str:
        """
        Get the model name

        Returns:
            str: The model name
        """
        pass

    @abstractmethod
    def set_response_format(self, response_format: Union[Dict, BaseModel]) -> None:
        """
        Set the response format for the LLM

        Args:
            response_format (Union[dict, BaseModel]): The response format
                to set
        """
        pass

    @abstractmethod
    def check_response_format(self) -> bool:
        """
        Check if the response format is supported by the LLM model

        Returns:
            bool: True if the response format is supported, False otherwise
        """
        pass

    @abstractmethod
    def generate(self, messages: List[InputMessage]) -> str:
        """
        Generate completions using the LLM API

        Args:
            messages (List[InputMessage]): List of messages to generate
                completions for using the LLM API

        Returns:
            str: The completion generated by the LLM API
        """
        pass

    @abstractmethod
    async def agenerate(self, messages: List[InputMessage]) -> str:
        """
        Generate completions using the LLM API asynchronously.

        Args:
            messages (List[InputMessage]): List of messages to generate
                completions for using the LLM API

        Returns:
            str: The completion generated by the LLM API
        """
        pass
