"""Abstract base class for augmentations in the litemind agentic framework.

This module defines ``AugmentationBase``, the interface that all augmentation
implementations must satisfy. Augmentations provide additional context to
large language models by retrieving relevant pieces of information based on
a query, enabling Retrieval-Augmented Generation (RAG) workflows.
"""

from abc import ABC, abstractmethod
from typing import Iterator, List, Optional, Union

from pydantic import BaseModel

from litemind.agent.augmentations.information.information import Information
from litemind.media.media_base import MediaBase


class AugmentationBase(ABC):
    """
    Abstract base class defining the interface for all augmentations.

    Augmentations provide additional context to the LLM by retrieving relevant
    informations based on a query.
    """

    def __init__(self, name: str, description: Optional[str] = None):
        """
        Create a new augmentation.

        Parameters
        ----------
        name: str
            The name of the augmentation.
        description: Optional[str]
            A description of the augmentation.
        """
        self.name = name
        self.description = description or f"{name} Augmentation"

    @abstractmethod
    def get_relevant_informations(
        self,
        query: Union[str, BaseModel, MediaBase, Information],
        k: int = 5,
        threshold: float = 0.0,
    ) -> List[MediaBase]:
        """
        Retrieve informations relevant to the query.

        Parameters
        ----------
        query : Union[str, BaseModel, MediaBase, Information]
            The query to retrieve informations for. Can be a string, a
            Pydantic model, a MediaBase object, or an Information object.
        k : int
            The maximum number of informations to retrieve.
        threshold : float
            Minimum relevance score. Results below this threshold are
            excluded.

        Returns
        -------
        List[InformationBase]
            A list of relevant informations, ordered by relevance.
        """
        pass

    def get_relevant_informations_iterator(
        self,
        query: Union[str, BaseModel, MediaBase, Information],
        k: int = 5,
        threshold: float = 0.0,
    ) -> Iterator[MediaBase]:
        """
        Retrieve informations relevant to the query as an iterator.

        The default implementation calls ``get_relevant_informations``
        and yields from the result.

        Parameters
        ----------
        query : Union[str, BaseModel, MediaBase, Information]
            The query to retrieve informations for.
        k : int
            The maximum number of informations to retrieve.
        threshold : float
            Minimum relevance score. Results below this threshold are
            excluded.

        Returns
        -------
        Iterator[InformationBase]
            An iterator over relevant informations.
        """
        pass

    def __repr__(self):
        """Return a detailed string representation of the augmentation.

        Returns
        -------
        str
            A string of the form ``ClassName(name='augmentation_name')``.
        """
        return f"{self.__class__.__name__}(name='{self.name}')"

    def __str__(self):
        """Return the name of the augmentation.

        Returns
        -------
        str
            The augmentation name.
        """
        return self.name
