import math
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import TypedDict, Any

from think.ai import ask
from think.llm.base import LLM


class RagDocument(TypedDict):
    """
    Document structure for RAG systems.

    A typed dictionary representing a document in the RAG index,
    containing a unique identifier and the document text content.

    Attributes:
        id: Unique identifier for the document
        text: The textual content of the document
    """

    id: str
    text: str


@dataclass
class RagResult:
    """
    Result from a RAG document retrieval operation.

    Represents a single document retrieved from the RAG system
    along with its relevance score for the given query.

    Attributes:
        doc: The retrieved document with id and text
        score: Relevance score (typically 0.0 to 1.0, higher is more relevant)
    """

    doc: RagDocument
    score: float


# Default Jinja2 template for generating answers from retrieved context.
#
# This template is used by RAG systems to generate answers based on retrieved
# documents. It formats the retrieved documents with their relevance scores
# and prompts the LLM to answer the query using only the provided context.
#
# Template variables:
#     results: List of RagResult objects containing retrieved documents and scores
#     query: The original user query to answer
#
# The template formats each document with its relevance score and separates
# them with horizontal rules. It instructs the LLM to base its answer solely
# on the provided context and to avoid mentioning the contextual nature of
# the response to maintain natural flow.
BASE_ANSWER_PROMPT = """Based ONLY on the provided context:

{% for item in results %}
{{ item.doc.text }}
Score: {{ item.score | round(3) }}
{% if not loop.last %}
---
{% endif %}
{% endfor %}

Answer the question:

{{query}}

(Note: don't say "based on provided context" in the output, it's confusing for the reader.)
"""


class RAG(ABC):
    """
    Abstract base class for Retrieval-Augmented Generation (RAG) systems.

    This class defines the common interface for RAG implementations that can
    index documents, perform semantic search, and generate answers based on
    retrieved context. Different providers can be plugged in by subclassing
    this class and implementing the abstract methods.

    The RAG pipeline consists of several stages:
    1. Document indexing via add_documents()
    2. Query preparation via prepare_query()
    3. Document retrieval via fetch_results()
    4. Optional result reranking via rerank()
    5. Answer generation via get_answer()

    Supported providers include TxtAI, ChromaDB, and Pinecone, each with
    their own specific implementations and capabilities.

    Class Attributes:
        PROVIDERS: List of supported RAG provider names
        QUERY_PROMPT: Optional template for query enhancement
        ANSWER_PROMPT: Jinja2 template for answer generation

    Example usage:
        rag = RAG.for_provider("txtai")(llm)
        await rag.add_documents([{"id": "1", "text": "content"}])
        answer = await rag("What is the content about?")
    """

    PROVIDERS = ["txtai", "chroma", "pinecone"]
    QUERY_PROMPT: str | None = None
    ANSWER_PROMPT: str = BASE_ANSWER_PROMPT

    def __init__(
        self,
        llm: LLM,
        **kwargs: Any,
    ):
        """
        Initialize the RAG system.

        :param llm: The LLM instance to use for query processing and answer generation
        :param kwargs: Additional keyword arguments for provider-specific configuration
        """
        self.llm = llm

    @abstractmethod
    async def add_documents(self, documents: list[RagDocument]):
        """
        Add documents to the RAG index.

        :param documents: Documents to add.
        """

    @abstractmethod
    async def remove_documents(self, ids: list[str]):
        """
        Remove documents from the RAG index.

        :param ids: Document IDs to remove.
        """

    async def prepare_query(self, query: str) -> str:
        """
        Process user input into query suitable for semantic search.

        :param query: User input.
        :return: Query suitable for semantic search.
        """
        if self.QUERY_PROMPT is None:
            return query
        return await ask(self.llm, self.QUERY_PROMPT, query=query)

    @abstractmethod
    async def fetch_results(
        self,
        user_query: str,
        prepared_query: str,
        limit: int,
    ) -> list[RagResult]:
        """
        Use the provided and processed query to search for relevant context.

        :param user_query: Unprocessed user input.
        :param prepared_query: Processed user input.
        :param limit: Maximum number of search results to return.
        :return: Search results to include in the context.
        """

    async def get_answer(self, query: str, results: list[RagResult]) -> str:
        """
        Ask the LLM to provide the answer based on the retrieved context
        and the user's original query.

        :param query: User input.
        :param results: Search results.
        :return: Answer to the user query.
        """
        return await ask(self.llm, self.ANSWER_PROMPT, results=results, query=query)

    async def rerank(self, results: list[RagResult]) -> list[RagResult]:
        """
        Rerank the search results based on the user query.

        :param results: Search results.
        :return: Reranked search results
        """
        return results

    async def __call__(self, query: str, limit: int = 10) -> str:
        """
        Execute the complete RAG pipeline for a query.

        This method orchestrates the full RAG process: query preparation,
        document retrieval, result reranking, and answer generation.

        :param query: The user's query string
        :param limit: Maximum number of documents to retrieve (default 10)
        :return: Generated answer based on retrieved context
        """
        prepared_query = await self.prepare_query(query)
        results = await self.fetch_results(query, prepared_query, limit)
        reranked_results = await self.rerank(results)
        return await self.get_answer(query, reranked_results)

    @abstractmethod
    async def calculate_similarity(self, query: str, docs: list[str]) -> list[float]:
        """
        Calculate the similarity between the query and the text.

        :param query: User input.
        :param text: Text to compare.
        :return: Similarity score.
        """

    @abstractmethod
    async def count(self) -> int:
        """
        Get the number of documents in the RAG index.

        :return: The number of documents in the RAG index.
        """

    @classmethod
    def for_provider(cls, provider: str) -> type["RAG"]:
        """
        Get the RAG class for the specified provider/engine.

        :param provider: The provider name
        :return The RAG class for the provider

        Raises a ValueError if the provider is not supported.
        The list of supported providers is available in the
        PROVIDERS class attribute.
        """
        if provider == "txtai":
            from .txtai_rag import TxtAiRag

            return TxtAiRag

        elif provider == "chroma":
            from .chroma_rag import ChromaRag

            return ChromaRag

        elif provider == "pinecone":
            from .pinecone_rag import PineconeRag

            return PineconeRag

        else:
            raise ValueError(f"Unknown provider: {provider}")

    @staticmethod
    def _cosine_similarity(a: list[float], b: list[float]):
        """
        Compute cosine similarity between two vectors of equal length.

        :param a: First vector
        :param b: Second vector
        :return: Cosine similarity between the two vectors
        """
        if len(a) != len(b):
            raise ValueError("Vectors must be of equal length")

        # Compute dot product
        dot_product = sum(x * y for x, y in zip(a, b))

        # Compute magnitudes
        magnitude1 = math.sqrt(sum(x * x for x in a))
        magnitude2 = math.sqrt(sum(x * x for x in b))

        # Prevent division by zero
        if magnitude1 == 0 or magnitude2 == 0:
            return 0.0

        # Compute cosine similarity
        return dot_product / (magnitude1 * magnitude2)
