"""A pipeline module for Retrieval Augmented Generation (RAG)"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/04_pipelines.rag.ipynb.

# %% auto 0
__all__ = ['DEFAULT_QA_PROMPT', 'DEFAULT_ROUTER_PROMPT', 'SUBQUESTION_PROMPT', 'FOLLOWUP_PROMPT', 'RAGPipeline',
           'CategorySelection', 'KVRouter']

# %% ../../nbs/04_pipelines.rag.ipynb 3
from typing import Optional, Dict, List, Any
from langchain_core.documents import Document
from ..utils import format_string, SafeFormatter
from ..llm import helpers
from pydantic import BaseModel, Field

# %% ../../nbs/04_pipelines.rag.ipynb 4
DEFAULT_QA_PROMPT = """Use the following pieces of context delimited by three backticks to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.

```{context}```

Question: {question}
Helpful Answer:"""

DEFAULT_ROUTER_PROMPT = """Given the following query/question, select the most appropriate category that would contain the relevant information.

Query: {question}

Available categories:
{categories}

Select the best category from the list above, or 'none' if no category is appropriate.
Do not provide an explanation for the categorization."""

# Question decomposition prompt template
SUBQUESTION_PROMPT = """\
Given a user question, output a list of relevant sub-questions \
in json markdown that when composed can help answer the full user question.
Only return the JSON response, with no additional text or explanations.
Generate between 2-5 sub-questions.  Do not exceed 5 sub-questions.

# Example 1
<User Question>
Compare and contrast the revenue growth and EBITDA of Uber and Lyft for year 2021

<Output>
```json
{
    "items": [
        {
            "sub_question": "What is the revenue growth of Uber",
        },
        {
            "sub_question": "What is the EBITDA of Uber",
        },
        {
            "sub_question": "What is the revenue growth of Lyft",
        },
        {
            "sub_question": "What is the EBITDA of Lyft",
        }
    ]
}
```

# Example 2
<User Question>
{query_str}

<Output>
"""

# Followup question determination prompt template
FOLLOWUP_PROMPT = """\
Given a question, answer "yes" only if the question is complex and follow-up questions are needed or "no" if not.
Always respond with "no" for short questions that are less than 8 words.
Answer only with either "yes" or "no" with no additional text or explanations.

# Example 1
<User Question>
Compare and contrast the revenue growth and EBITDA of Uber and Lyft for year 2021

<Output>
yes

# Example 2
<User Question>
How is the Coast Guard using artificial intelligence?

<Output>
No

# Example 3
<User Question
What is AutoGluon?

<Output>
No

# Example 4
<User Question>
{query_str}

<Output>
"""

# %% ../../nbs/04_pipelines.rag.ipynb 5
class RAGPipeline:
    """
    Retrieval-Augmented Generation pipeline for answering questions based on source documents.
    """
    
    def __init__(self, llm, qa_template: str = DEFAULT_QA_PROMPT):
        """
        Initialize RAG pipeline.
        
        Args:
            llm: The language model instance (LLM object)
            qa_template: Question-answering prompt template
        """
        self.llm = llm
        self.qa_template = qa_template
    
    def semantic_search(self,
                       query: str, # search query as string
                       limit: int = 4, # number of sources to retrieve
                       score_threshold: float = 0.0, # minimum threshold for score
                       filters: Optional[Dict[str, str]] = None, # metadata filters
                       where_document = None, # filter search results based syntax of underlying store
                       folders: Optional[list] = None, # list of folders to consider
                       **kwargs) -> List[Document]:
        """
        Perform a semantic search of the vector DB.

        The `where_document` parameter varies depending on the value of `LLM.store_type`.
        If `LLM.store_type` is 'dense', then `where_document` should be a dictionary in Chroma syntax (e.g., {"$contains": "Canada"})
        to filter results.
        If `LLM.store_type` is 'sparse', then `where_document` should be a boolean search string to filter query in Lucene syntax.
        """
        import os
        store = self.llm.load_vectorstore()
        if folders:
            folders = [folders] if isinstance(folders, str) else folders
            # This is needed because only the where argument supports the $like operator
            # and Langchain does not properly forward the where parameter to Chroma
            n_candidates = store.get_size() if store.get_size() < 10000 else 10000
            results = store.semantic_search(query, 
                                            filters=filters,
                                            where_document=where_document,
                                            limit = n_candidates, **kwargs)
            # Handle path separator differences between Windows and Unix
            if os.name == 'nt':  # Windows
                # Normalize paths for case-insensitive comparison on Windows
                normalized_folders = [os.path.normpath(f).lower().replace('\\', '/') for f in folders]
                results = [d for d in results if any(os.path.normpath(d.metadata['source']).lower().replace('\\', '/').startswith(nf) for nf in normalized_folders)]
            else:
                # On Unix systems, use direct path comparison
                results = [d for d in results if any(d.metadata['source'].startswith(f) for f in folders)]
            results = results[:limit]
            
        else:
            results = store.semantic_search(query, 
                                            filters=filters,
                                            where_document=where_document,
                                            limit = limit, **kwargs)

        return [d for d in results if d.metadata['score'] >= score_threshold]

    def _retrieve_documents(self, 
                          question: str,
                          filters: Optional[Dict[str, str]] = None,
                          where_document = None,
                          folders: Optional[list] = None,
                          limit: int = 4,
                          score_threshold: float = 0.0,
                          table_k: int = 1,
                          table_score_threshold: float = 0.35) -> List[Document]:
        """
        Retrieve relevant documents from vector database.
        """
        docs = self.semantic_search(
            question, 
            filters=filters, 
            where_document=where_document, 
            folders=folders,
            limit=limit,
            score_threshold=score_threshold
        )
        
        # Add table documents if requested
        if table_k > 0:
            table_filters = filters.copy() if filters else {}
            table_filters = dict(table_filters, table=True)
            table_docs = self.semantic_search(
                f'{question} (table)', 
                filters=table_filters, 
                where_document=where_document,
                folders=folders,
                limit=table_k,
                score_threshold=table_score_threshold
            )
            if table_docs:
                docs.extend(table_docs[:limit])
        
        return docs
    
    def _generate_answer(self, question: str, context: str, **kwargs) -> str:
        """
        Generate answer using the language model.
        """
        prompt = format_string(
            self.qa_template,
            question=question,
            context=context
        )
        return self.llm.prompt(prompt, **kwargs)
    
    def decompose_question(self, question: str, parse=True, **kwargs):
        """
        Decompose a question into subquestions
        """
        prompt = SafeFormatter({'query_str': question}).format(SUBQUESTION_PROMPT)
        json_string = self.llm.prompt(prompt)
        json_dict = helpers.parse_json_markdown(json_string)
        subquestions = [d['sub_question'] for d in json_dict['items']]
        return subquestions


    def needs_followup(self, question: str, parse=True, **kwargs):
        """
        Decide if follow-up questions are needed
        """
        prompt = SafeFormatter({'query_str': question}).format(FOLLOWUP_PROMPT)
        output = self.llm.prompt(prompt)
        return "yes" in output.lower()
    

    def ask(self,
            question: str, # question as string
            contexts: Optional[list] = None, # optional list of contexts to answer question. If None, retrieve from vectordb.
            qa_template: Optional[str] = None, # question-answering prompt template to use
            filters: Optional[Dict[str, str]] = None, # filter sources by metadata values using Chroma metadata syntax (e.g., {'table':True})
            where_document = None, # filter sources by document content (syntax varies by store type)
            folders: Optional[list] = None, # folders to search (needed because LangChain does not forward "where" parameter)
            limit: Optional[int] = None, # Number of sources to consider. If None, use `LLM.rag_num_source_docs`.
            score_threshold: Optional[float] = None, # minimum similarity score of source. If None, use `LLM.rag_score_threshold`.
            table_k: int = 1, # maximum number of tables to consider when generating answer
            table_score_threshold: float = 0.35, # minimum similarity score for table to be considered in answer
            selfask: bool = False, # If True, use an agentic Self-Ask prompting strategy.
            router = None, # Optional KVRouter instance for automatic filtering
            **kwargs) -> Dict[str, Any]:
        """
        Answer a question using RAG approach.
        Additional kwargs arguments passed to LLM.prompt
        Returns dictionary with keys: answer, source_documents, question.
        """
        template = qa_template or self.qa_template
        limit = limit if limit is not None else self.llm.rag_num_source_docs
        score_threshold = score_threshold if score_threshold is not None else self.llm.rag_score_threshold
        
        if selfask and self.needs_followup(question):
            return self._ask_with_decomposition(
                question, template, filters, where_document, folders,
                limit, score_threshold, table_k, table_score_threshold, router, **kwargs
            )
        else:
            return self._ask_direct(
                question, contexts, template, filters, where_document, folders,
                limit, score_threshold, table_k, table_score_threshold, router, **kwargs
            )
    
    def _ask_direct(self,
                   question: str,
                   contexts: Optional[list],
                   qa_template: str,
                   filters: Optional[Dict[str, str]],
                   where_document,
                   folders: Optional[list],
                   limit: int,
                   score_threshold: float,
                   table_k: int,
                   table_score_threshold: float,
                   router = None,
                   **kwargs) -> Dict[str, Any]:
        """Direct RAG without decomposition."""
        # Apply router if provided
        if router and not filters:
            router_filters = router.route(question)
            if router_filters:
                filters = router_filters
        elif router and filters:
            # Merge router filters with existing filters
            router_filters = router.route(question)
            if router_filters:
                combined_filters = filters.copy()
                combined_filters.update(router_filters)
                filters = combined_filters
        
        if contexts is None:
            docs = self._retrieve_documents(
                question, filters, where_document, folders,
                limit, score_threshold, table_k, table_score_threshold
            )
            context = '\n\n'.join([d.page_content for d in docs])
        else:
            docs = [Document(page_content=c, metadata={'source': '<SUBANSWER>'}) for c in contexts]
            context = "\n\n".join(contexts)
        
        answer = self._generate_answer(question, context, **kwargs)
        
        return {
            'question': question,
            'answer': answer,
            'source_documents': docs
        }
    
    def _ask_with_decomposition(self,
                               question: str,
                               qa_template: str,
                               filters: Optional[Dict[str, str]],
                               where_document,
                               folders: Optional[list],
                               limit: int,
                               score_threshold: float,
                               table_k: int,
                               table_score_threshold: float,
                               router = None,
                               **kwargs) -> Dict[str, Any]:
        """RAG with question decomposition (Self-Ask)."""
        subquestions = self.decompose_question(question)
        subanswers = []
        sources = []
        
        for q in subquestions:
            res = self._ask_direct(
                q, None, qa_template, filters, where_document, folders,
                limit, score_threshold, table_k, table_score_threshold, router, **kwargs
            )
            subanswers.append(res['answer'])
            for doc in res['source_documents']:
                doc.metadata = dict(doc.metadata, subquestion=q)
            sources.extend(res['source_documents'])
        
        # Generate final answer based on subanswers
        res = self._ask_direct(
            question, subanswers, qa_template, filters, where_document, folders,
            limit, score_threshold, table_k, table_score_threshold, None, **kwargs
        )
        res['source_documents'] = sources
        
        return res




# %% ../../nbs/04_pipelines.rag.ipynb 10
class CategorySelection(BaseModel):
    """Pydantic model for category selection response."""
    category: str = Field(description="Selected category value or 'none' if no appropriate category")


class KVRouter:
    """
    Key-Value Router for intelligent filtering based on query content.
    
    Uses an LLM to select the most appropriate field value for filtering
    based on the query/question content.
    """
    
    def __init__(self, 
                 field_name: str,
                 field_descriptions: Dict[str, str],
                 llm,
                 router_prompt: str = DEFAULT_ROUTER_PROMPT):
        """
        Initialize KV Router.
        
        Args:
            field_name: The metadata field name to filter on (e.g., 'folder')
            field_descriptions: Dict mapping field values to descriptions
                               (e.g., {'sotu': "Biden's State of the Union Address"})
            llm: The LLM instance to use for routing decisions
            router_prompt: Template for the routing prompt
        """
        self.field_name = field_name
        self.field_descriptions = field_descriptions
        self.llm = llm
        self.router_prompt = router_prompt
    
    def _format_categories(self) -> str:
        """Format field descriptions for the prompt."""
        categories = []
        for value, description in self.field_descriptions.items():
            categories.append(f"- {value}: {description}")
        return "\n".join(categories)
    

    def route(self, question: str, **kwargs) -> Optional[Dict[str, str]]:
        """
        Select the best field value for the given question.
        Extra **kwargs supplied to LLM.prompt.
        
        Args:
            question: The user's question/query
            
        Returns:
            Dictionary for filters parameter, or None if no appropriate category
            Example: {'folder': 'sotu'} or None
        """
        # Format the prompt
        categories_text = self._format_categories()
        prompt = format_string(
            self.router_prompt,
            question=question,
            categories=categories_text
        )
        
        # Use response_format for structured output
        try:
            response = self.llm.prompt(
                prompt, 
                response_formatl=CategorySelection,
                **kwargs
            )
            
            selected_category = response.category.strip().lower()
            
            # Check if it's a valid category (case-insensitive)
            valid_values = {v.lower(): v for v in self.field_descriptions.keys()}
            
            if selected_category in valid_values:
                return {self.field_name: valid_values[selected_category]}
            elif selected_category == 'none':
                return None
            else:
                # Fuzzy matching: check if any valid value appears in the selected_category
                for valid_key, original_value in valid_values.items():
                    if valid_key in selected_category:
                        return {self.field_name: original_value}
                
                # Fallback: if response doesn't match, return None
                return None
                
        except Exception as e:
            # Fallback to None if pydantic parsing fails
            return None
    
    def route_and_search(self, 
                        query: str,
                        rag_pipeline,
                        **search_kwargs) -> List[Document]:
        """
        Convenience method that routes and performs semantic search.
        
        Args:
            query: The search query
            rag_pipeline: RAGPipeline instance to search with
            **search_kwargs: Additional arguments passed to semantic_search
            
        Returns:
            List of Document objects
        """
        # Get routing decision
        route_filters = self.route(query)
        
        # Merge with existing filters if any
        existing_filters = search_kwargs.get('filters', {})
        if route_filters:
            if existing_filters:
                # Combine filters
                combined_filters = existing_filters.copy()
                combined_filters.update(route_filters)
                search_kwargs['filters'] = combined_filters
            else:
                search_kwargs['filters'] = route_filters
        
        # Perform search
        return rag_pipeline.semantic_search(query, **search_kwargs)
