"""
Prompt Templates for RGB RAG Evaluation
Based on Figure 3 from the research paper: https://arxiv.org/pdf/2309.01431

The RGB paper uses a unified prompt template with a system instruction and user instruction.
This ensures consistent evaluation across all RAG abilities.
"""

# System Instruction (from Figure 3 of the paper)
SYSTEM_INSTRUCTION = """You are an accurate and reliable AI assistant that can answer questions with the help of external documents. Please note that external documents may contain noisy or factually incorrect information. If the information in the document contains the correct answer, you will give an accurate answer. If the information in the document does not contain the answer, you will generate 'I can not answer the question because of the insufficient information in documents.' If there are inconsistencies with the facts in some of the documents, please generate the response 'There are factual errors in the provided documents.' and provide the correct answer."""

# Main RAG prompt template (from Figure 3 of the paper)
# This prompt is used for all evaluation tasks
RAG_PROMPT_TEMPLATE = """Document:
{documents}

Question: {question}"""


def get_system_instruction() -> str:
    """Get the system instruction for RAG evaluation."""
    return SYSTEM_INSTRUCTION


def get_prompt_template(task_type: str = "default") -> str:
    """
    Get the appropriate prompt template for a task type.
    All tasks use the same base template from Figure 3.
    
    Args:
        task_type: The type of task (for future extensibility)
        
    Returns:
        The prompt template string.
    """
    # All tasks use the same base template as per Figure 3
    return RAG_PROMPT_TEMPLATE


def format_prompt(
    question: str,
    documents: list,
    template: str = None
) -> str:
    """
    Format a prompt with question and documents.
    Implements the exact format from Figure 3 of the paper.
    
    Args:
        question: The question to answer.
        documents: List of document texts.
        template: Optional custom template. Uses default if None.
        
    Returns:
        Formatted prompt string.
    """
    if template is None:
        template = RAG_PROMPT_TEMPLATE
    
    # Format documents as in Figure 3
    # Each document is labeled with its index
    docs_formatted = "\n".join(documents)
    
    return template.format(
        documents=docs_formatted,
        question=question
    )