from typing import List, Dict, Optional


def get_generate_subquery_prompt(query: str, past_subqueries: List[str], past_subanswers: List[str], task_desc: str) -> List[Dict]:
    assert len(past_subqueries) == len(past_subanswers)
    past = ''
    for idx in range(len(past_subqueries)):
        past += f"""Intermediate query {idx+1}: {past_subqueries[idx]}
Intermediate answer {idx + 1}: {past_subanswers[idx]}\n"""
    past = past.strip()

    prompt = f"""You are using a search engine to answer the main query by iteratively searching the web. Given the following intermediate queries and answers, generate a new simple follow-up question that can help answer the main query. You may rephrase or decompose the main query when previous answers are not helpful. Ask simple follow-up questions only as the search engine may not understand complex questions.

## Previous intermediate queries and answers
{past or 'Nothing yet'}

## Task description
{task_desc}

## Main query to answer
{query}

Respond with a simple follow-up question that will help answer the main query, do not explain yourself or output anything else."""

    messages: List[Dict] = [
        {'role': 'user', 'content': prompt}
    ]
    return messages


def get_generate_intermediate_answer_prompt(subquery: str, documents: List[str]) -> List[Dict]:
    context = ''
    for idx, doc in enumerate(documents):
        context += f"""{doc}\n\n"""

    prompt = f"""Given the following documents, generate an appropriate answer for the query. DO NOT hallucinate any information, only use the provided documents to generate the answer. Respond "No relevant information found" if the documents do not contain useful information.

## Documents
{context.strip()}

## Query
{subquery}

Respond with a concise answer only, do not explain yourself or output anything else."""

    messages: List[Dict] = [
        {'role': 'user', 'content': prompt}
    ]
    return messages


def get_generate_final_answer_prompt(
        query: str, past_subqueries: List[str], past_subanswers: List[str], task_desc: str,
        documents: Optional[List[str]] = None
) -> List[Dict]:

    assert len(past_subqueries) == len(past_subanswers)
    past = ''
    for idx in range(len(past_subqueries)):
        past += f"""Intermediate query {idx+1}: {past_subqueries[idx]}
Intermediate answer {idx+1}: {past_subanswers[idx]}\n"""
    past = past.strip()

    context = ''
    if documents:
        for idx, doc in enumerate(documents):
            context += f"""Doc {idx}: {doc}\n\n"""

    prompt = f"""Given the following intermediate queries and answers, generate a final answer for the main query by combining relevant information. Note that intermediate answers are generated by an LLM and may not always be accurate.

## Documents
{context.strip()}

## Intermediate queries and answers
{past or 'Nothing yet'}

## Task description
{task_desc}

## Main query
{query}

Respond with an appropriate answer only, do not explain yourself or output anything else."""

    messages: List[Dict] = [
        {'role': 'user', 'content': prompt}
    ]
    return messages
