import logging
from typing import List

from autorag.nodes.promptmaker.base import prompt_maker_node

logger = logging.getLogger("AutoRAG")


@prompt_maker_node
def long_context_reorder(prompt: str, queries: List[str], retrieved_contents: List[List[str]],
                         retrieve_scores: List[List[float]]) -> List[str]:
    """
    Models struggle to access significant details found
    in the center of extended contexts. A study
    (https://arxiv.org/abs/2307.03172) observed that the best
    performance typically arises when crucial data is positioned
    at the start or conclusion of the input context. Additionally,
    as the input context lengthens, performance drops notably, even
    in models designed for long contexts.".

    .. Code:: yaml
    nodes:
    - node_type: prompt_maker
      modules:
      - module_type: long_context_reorder
        prompt: [Answer this question: {query} \n\n {retrieved_contents},
        Read the passages carefully and answer this question: {query} \n\n Passages: {retrieved_contents}]

    :param prompt: A prompt string.
    :param queries: List of query strings.
    :param retrieved_contents: List of retrieved contents.
    :param retrieve_scores: List of retrieve scores.
    :return: Prompts that made by long context reorder.
    """

    def long_context_reorder_row(_prompt: str, _query: str, _retrieved_contents: List[str],
                                 _retrieve_scores: List[float]) -> str:
        if not len(_retrieved_contents) == len(_retrieve_scores):
            logger.info("If you use a summarizer, the reorder will not proceed.")
            return _prompt.format(query=_query, retrieved_contents="\n\n".join(_retrieved_contents))
        content_scores = list(zip(_retrieved_contents, _retrieve_scores))
        sorted_content_scores = sorted(content_scores, key=lambda x: x[1], reverse=True)
        content_result, score_result = zip(*sorted_content_scores)
        _retrieved_contents.append(content_result[0])
        contents_str = "\n\n".join(_retrieved_contents)
        return _prompt.format(query=_query, retrieved_contents=contents_str)

    return list(map(lambda x: long_context_reorder_row(prompt, x[0], x[1], x[2]),
                    zip(queries, retrieved_contents, retrieve_scores)))
