"""
The module is primarily centered around the `FlowRetrieval` class which handles the
initialization, configuration validation, vector store management, and flow retrieval
logic. It integrates components for managing embeddings, vector stores, and
flow-specific templates, facilitating semantic search functionalities.
Key Features:
- Configurable embedding strategies for dialogue components.
- Seamless interaction with model storage and resource management.
- Supports dynamic loading and persistence of vector stores.
- Enables population of vector stores based on specified dialogue flows and domain
information.
- Implements flow retrieval mechanisms including semantic search based on dialogue
context.
Usage:
Interaction with this class typically involves creating an instance with a
configuration dict, model storage, and a resource reference, then using methods
like `populate`, `persist`, or dynamic retrieval methods to manage or utilize
flows within a conversational context.
"""

import importlib
from typing import Dict, Text, Any, List, Optional

import structlog
from jinja2 import Template
from langchain.docstore.document import Document
from langchain.schema.embeddings import Embeddings
from langchain.vectorstores.faiss import FAISS
from langchain.vectorstores.utils import DistanceStrategy
from rasa.engine.storage.resource import Resource
from rasa.engine.storage.storage import ModelStorage
from rasa.shared.core.domain import Domain
from rasa.shared.core.flows import FlowsList
from rasa.shared.core.trackers import DialogueStateTracker
from rasa.shared.nlu.constants import TEXT, FLOWS_FROM_SEMANTIC_SEARCH
from rasa.shared.nlu.training_data.message import Message
from rasa.shared.exceptions import ProviderClientAPIException
from rasa.shared.utils.llm import (
    tracker_as_readable_transcript,
    embedder_factory,
    DEFAULT_OPENAI_EMBEDDING_MODEL_NAME,
    USER,
    get_prompt_template,
    allowed_values_for_slot,
)

DEFAULT_FLOW_DOCUMENT_TEMPLATE = importlib.resources.read_text(
    "rasa.dialogue_understanding.generator", "flow_document_template.jinja2"
)

EMBEDDINGS_CONFIG_KEY = "embeddings"
DEFAULT_EMBEDDINGS_CONFIG = {
    "_type": "openai",
    "model": DEFAULT_OPENAI_EMBEDDING_MODEL_NAME,
}

MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY = "num_flows"
SHOULD_EMBED_SLOTS_KEY = "should_embed_slots"
TURNS_TO_EMBED_KEY = "turns_to_embed"

DEFAULT_MAX_FLOWS_FROM_SEMANTIC_SEARCH = 20
DEFAULT_TURNS_TO_EMBED = 1
DEFAULT_SHOULD_EMBED_SLOTS = True


structlogger = structlog.get_logger()


class FlowRetrieval:
    @classmethod
    def get_default_config(cls) -> Dict[str, Any]:
        """The default config for the flow retrieval."""
        return {
            EMBEDDINGS_CONFIG_KEY: DEFAULT_EMBEDDINGS_CONFIG,
            MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY: DEFAULT_MAX_FLOWS_FROM_SEMANTIC_SEARCH,
            TURNS_TO_EMBED_KEY: DEFAULT_TURNS_TO_EMBED,
            SHOULD_EMBED_SLOTS_KEY: DEFAULT_SHOULD_EMBED_SLOTS,
        }

    def __init__(
        self,
        config: Dict[str, Any],
        model_storage: ModelStorage,
        resource: Resource,
    ):
        config = {**self.get_default_config(), **config}
        self.config = self.validate_config(config)
        self.vector_store: Optional[FAISS] = None
        self.flow_document_template = get_prompt_template(
            None, DEFAULT_FLOW_DOCUMENT_TEMPLATE
        )
        self._model_storage = model_storage
        self._resource = resource

    @classmethod
    def validate_config(cls, config: Dict[Text, Any]) -> Dict[Text, Any]:
        if config[MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY] < 0:
            config[MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY] = (
                DEFAULT_MAX_FLOWS_FROM_SEMANTIC_SEARCH
            )
            structlogger.error(
                f"flow_retrieval.validate_config.{MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY}.set_as_negative",
                event_info=(
                    f"The `{MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY}` is a "
                    f"negative value. Setting it to the default value "
                    f"({DEFAULT_MAX_FLOWS_FROM_SEMANTIC_SEARCH})."
                ),
                old_value=config[MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY],
                new_value=DEFAULT_MAX_FLOWS_FROM_SEMANTIC_SEARCH,
            )

        if config[TURNS_TO_EMBED_KEY] < 1:
            config[TURNS_TO_EMBED_KEY] = DEFAULT_TURNS_TO_EMBED
            structlogger.error(
                f"flow_retrieval.validate_config.{TURNS_TO_EMBED_KEY}.less_than_one",
                event_info=(
                    f"The `{TURNS_TO_EMBED_KEY}` is less than 1."
                    f"Setting it to the default value ({DEFAULT_TURNS_TO_EMBED})."
                    f"Only the latest user utterance will be used."
                ),
                old_value=config[TURNS_TO_EMBED_KEY],
                new_value=DEFAULT_TURNS_TO_EMBED,
            )

        return config

    @classmethod
    def load(
        cls,
        config: Dict[Text, Any],
        model_storage: ModelStorage,
        resource: Resource,
        **kwargs: Any,
    ) -> "FlowRetrieval":
        """Load flow retrieval with previously populated FAISS vector store."""
        # initialize base flow retrieval
        flow_retrieval = FlowRetrieval(config, model_storage, resource)
        # load vector store
        vector_store = cls._load_vector_store(
            flow_retrieval.config, model_storage, resource
        )
        flow_retrieval.vector_store = vector_store
        return flow_retrieval

    @classmethod
    def _load_vector_store(
        cls, config: Dict[Text, Any], model_storage: ModelStorage, resource: Resource
    ) -> Optional[FAISS]:
        """Loads a FAISS vector store from a specified local path."""
        embeddings = cls._create_embedder(config)
        try:
            with model_storage.read_from(resource) as model_path:
                return FAISS.load_local(
                    folder_path=model_path,
                    embeddings=embeddings,
                    distance_strategy=DistanceStrategy.MAX_INNER_PRODUCT,
                )
        except Exception as e:
            structlogger.warning(
                "flow_retrieval.load_vector_store.failed",
                error=e,
                resource=resource.name,
            )
            return None

    @classmethod
    def _create_embedder(cls, config: Dict[Text, Any]) -> Embeddings:
        """Creates an embedder.

        Returns:
            The embedder.
        """
        return embedder_factory(
            config.get(EMBEDDINGS_CONFIG_KEY), DEFAULT_EMBEDDINGS_CONFIG
        )

    def persist(self) -> None:
        self._persist_vector_store()

    def _persist_vector_store(self) -> None:
        """Persists the FAISS vector store."""
        if self.vector_store is not None:
            with self._model_storage.write_to(self._resource) as model_path:
                self.vector_store.save_local(model_path)
        else:
            structlogger.warning(
                "flow_retrieval.persist_vector_store.not_initialized",
                event_info="Vector store is None, not persisted.",
            )

    def populate(self, flows: FlowsList, domain: Domain) -> None:
        """Populates the vector store with embeddings generated from
        documents based on the flow descriptions, and flow slots
        descriptions.

        Args:
            flows: List of flows to populate the vector store with.
            domain: The domain containing relevant slot information.
        """
        flows_to_embedd = flows.exclude_link_only_flows()
        embeddings = self._create_embedder(self.config)
        documents = self._generate_flow_documents(flows_to_embedd, domain)
        try:
            self.vector_store = FAISS.from_documents(
                documents=documents,
                embedding=embeddings,
                distance_strategy=DistanceStrategy.MAX_INNER_PRODUCT,
            )
        except Exception as e:
            error_type = e.__class__.__name__
            structlogger.error(
                "flow_retrieval.populate_vector_store.not_populated",
                event_info=(
                    "Failed to populate the FAISS store with the provided flows."
                ),
                error_type=error_type,
                error=e,
            )
            raise

    def _generate_flow_documents(
        self, flows: FlowsList, domain: Domain
    ) -> List[Document]:
        """Converts each flow in the provided list into an embeddable document. These
        documents include the following information for each flow: flow's name,
        flow's description, and associated slots and their descriptions and allowed
        values.

        Args:
            flows: List of flows.
            domain: The domain containing relevant slot information.

        Returns:
            List of documents, each representing the converted information of a single
            flow.
        """
        slots = {slot.name: slot for slot in domain.slots}
        flow_docs = []

        for flow in flows:
            flow_info: Dict[Text, Any] = {
                "flow": {
                    "name": flow.name,
                    "description": flow.description,
                    "slots": [],
                }
            }
            if self.config[SHOULD_EMBED_SLOTS_KEY]:
                flow_info["flow"]["slots"] = [
                    {
                        "name": q.collect,
                        "description": q.description,
                        "allowed_values": allowed_values_for_slot(slots[q.collect]),
                    }
                    for q in flow.get_collect_steps()
                ]

            flow_docs.append(
                Document(
                    page_content=Template(self.flow_document_template)
                    .render(flow_info)
                    .strip(),
                    metadata={"flow_id": flow.id},
                )
            )

        return flow_docs

    async def filter_flows(
        self, tracker: DialogueStateTracker, message: Message, flows: FlowsList
    ) -> FlowsList:
        """Filters the given flows.
        The filtered flows with the following rules:
        - Exclude flows that reachable only via link
        - Include flows set to be always included in the prompt
        - Include flows started during the conversation,
        - (Optionally) The rest of flows are filter so only the top 'k' with
          the highest similarity to the ongoing conversation are included.
          Set 'max_flows_from_semantic_search' to negative number to disable
          the option.

        Args:
            tracker: the tracker
            message: the user message
            flows: available flows

        Returns:
            List of flows that includes always-included flows, previously started flows
            and top `k` relevant flows for the current conversation.
        """
        # apply basic filtering
        flows = flows.exclude_link_only_flows()
        always_included_flows = flows.get_flows_always_included_in_prompt()
        previously_started_flows = tracker.get_previously_started_flows(flows)
        # apply semantic search filtering
        most_similar_flows = await self.find_most_similar_flows(tracker, message, flows)
        return FlowsList.from_multiple_flows_lists(
            always_included_flows,
            previously_started_flows,
            most_similar_flows,
        )

    async def find_most_similar_flows(
        self, tracker: DialogueStateTracker, message: Message, flows: FlowsList
    ) -> FlowsList:
        """Filters the given flows so only the top 'k' most similar
        flows to the current conversation are left.
        retrieved from the vector store.

        Args:
            tracker: the tracker
            message: the user message

        Returns:
            The most similar flows to the current conversation.
        """
        query = self._prepare_query(tracker, message)
        documents_with_scores = await self._query_vector_store(query)
        # filter out None i.e. more flows were embedded during training than are
        # available during prediction
        most_similar_flows_with_scores = [
            (flow, score)
            for doc, score in documents_with_scores
            if (flow := flows.flow_by_id(doc.metadata["flow_id"])) is not None
        ]
        # sort by score decreasing
        most_similar_flows_with_scores.sort(key=lambda x: x[1], reverse=True)
        # add the relevant flows to the message for evaluation purposes
        message.set(
            prop=FLOWS_FROM_SEMANTIC_SEARCH,
            info=[
                (flow.id, float(score))
                for flow, score in most_similar_flows_with_scores
            ],
            add_to_output=True,
        )
        return FlowsList([f for f, _ in most_similar_flows_with_scores])

    def _prepare_query(self, tracker: DialogueStateTracker, message: Message) -> Text:
        """Prepares the query for vector store. The query is composed
        of the conversation within the specified number of turns.

        Args:
            tracker: the tracker
            message: the user message

        Returns:
            The query for vector store.
        """
        if int(self.config[TURNS_TO_EMBED_KEY]) > 1:
            current_conversation = tracker_as_readable_transcript(
                tracker,
                human_prefix=USER,
                max_turns=int(self.config[TURNS_TO_EMBED_KEY]),
            )
            current_conversation += f"\n{USER}: {message.data[TEXT]}"
            return current_conversation

        return f"{message.data[TEXT]}"

    async def _query_vector_store(self, query: Text) -> List:
        """Compares the query with all flows using a vector store
        and returns the top k relevant flows for the current conversation.

        Args:
            query: query

        Returns:
            The top k documents with similarity scores.
        """
        if self.vector_store is None:
            structlogger.error(
                "flow_retrieval.query_vector_store.vector_store_not_configured",
                event_info="Vector store is not configured",
            )
            return []
        try:
            documents_with_scores = (
                await self.vector_store.asimilarity_search_with_score(
                    query, k=int(self.config[MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY])
                )
            )
            structlogger.debug(
                "flow_retrieval.query_vector_store.fetched",
                event_info=(
                    f"Fetched the top {self.config[MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY]}"
                    f"similar flows from the vector store"
                ),
                query=query,
                top_k=self.config[MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY],
                results=[
                    {
                        "flow_id": document.metadata["flow_id"],
                        "score": score,
                        "content": document.page_content,
                    }
                    for document, score in documents_with_scores
                ],
            )
            return documents_with_scores
        except Exception as e:
            error_type = e.__class__.__name__
            structlogger.error(
                "flow_retrieval.query_vector_store.error",
                event_info="Cannot fetch flows from vector store",
                error_type=error_type,
                error=e,
                query=query,
            )
            raise ProviderClientAPIException(
                message="Cannot fetch flows from vector store", original_exception=e
            )
