from pathlib import Path
from typing import Any, List, Optional, Union
from uuid import uuid4

from PIL import Image

from byaldi.colpali import ColPaliModel
from byaldi.objects import Result


class RAGMultiModalModel:
    """
    Wrapper class for a pretrained RAG multi-modal model, and all the associated utilities.
    Allows you to load a pretrained model from disk or from the hub, build or query an index.

    ## Usage

    Load a pre-trained checkpoint:

    ```python
    from byaldi import RAGMultiModalModel

    RAG = RAGMultiModalModel.from_pretrained("vidore/colpali")
    ```

    Both methods will load a fully initialised instance of ColPali, which you can use to build and query indexes.

    ```python
    RAG.search("How many people live in France?")
    ```
    """

    model: Optional[ColPaliModel] = None

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: Union[str, Path],
        device: str = "cuda",
    ):
        """Load a ColPali model from a pre-trained checkpoint.

        Parameters:
            pretrained_model_name_or_path (str): Local path or huggingface model name.
            device (str): The device to load the model on. Default is "cuda".

        Returns:
            cls (RAGMultiModalModel): The current instance of RAGMultiModalModel, with the model initialised.
        """
        instance = cls()
        instance.model = ColPaliModel.from_pretrained(
            pretrained_model_name_or_path, device=device
        )
        return instance

    @classmethod
    def from_index(cls, index_path: Union[str, Path], index_root: str = ".byaldi", device: str = "cuda"):
        """Load an Index and the associated ColPali model from an existing document index.

        Parameters:
            index_path (Union[str, Path]): Path to the index.
            device (str): The device to load the model on. Default is "cuda".

        Returns:
            cls (RAGMultiModalModel): The current instance of RAGMultiModalModel, with the model and index initialised.
        """
        instance = cls()
        index_path = Path(index_path)
        instance.model = ColPaliModel.from_index(index_path, index_root=index_root, device=device)

        return instance

    def index(
        self,
        input_path: Union[str, Path],
        index_name: Optional[str] = None,
        doc_ids: Optional[List[Union[str, int]]] = None,
        store_collection_with_index: bool = False,
        override: bool = False,
    ):
        """Build an index from input documents.

        Parameters:
            input_path (Union[str, Path]): Path to the input documents.
            index_name (Optional[str]): The name of the index that will be built.
            doc_ids (Optional[List[Union[str, int]]]): List of document IDs.
            store_collection_with_index (bool): Whether to store the collection with the index.

        Returns:
            None
        """
        self.model.index(input_path, index_name, doc_ids, store_collection_with_index, override=override)

    def add_to_index(
        self,
        input_item: Union[str, Path, Image.Image],
        store_collection_with_index: bool,
        doc_id: Union[str, int],
    ):
        """Add an item to an existing index.

        Parameters:
            input_item (Union[str, Path, Image.Image]): The item to add to the index.
            store_collection_with_index (bool): Whether to store the collection with the index.
            doc_id (Union[str, int]): The document ID for the item being added.

        Returns:
            None
        """
        self.model.add_to_index(input_item, store_collection_with_index, doc_id)

    def search(
        self,
        query: Union[str, List[str]],
        k: int = 10,
        return_base64_results: Optional[bool] = None,
    ) -> Union[List[Result], List[List[Result]]]:
        """Query an index.

        Parameters:
            query (Union[str, List[str]]): The query or queries to search for.
            k (int): The number of results to return. Default is 10.
            return_base64_results (Optional[bool]): Whether to return base64-encoded image results.

        Returns:
            Union[List[Result], List[List[Result]]]: A list of Result objects or a list of lists of Result objects.
        """
        return self.model.search(query, k, return_base64_results)
