"""
This agent selects the search APIs to use, and returns the selected APIs in its response in
non-json format.
"""

import re
from collections.abc import Generator

from loguru import logger

from app.config import config
from app.data_structures import MessageThread
from app.log import print_acr, print_retrieval
from app.model import common, ollama

SELECT_PROMPT = (
    "Based on the previous analysis of test feedback and possible missing information, select necessary APIs to get more context. It's better to provide the APIs you need to call and their arguments in your response."
)


QUESTION_PROMPT = """To help diagnose and refine the generated code with feedback, please carefully answer the following questions. Be objective and analytical.

Please answer the following two questions:

1. **Do you have enough context to confidently fix the code so it passes the test?**  
   - If yes, describe your fix strategy in detail.  
   - If not, describe what information is needed to improve our understanding of the error feedback.

2. **If not enough context is available, which specific API calls (e.g., `search_relevant_method`, `search_target_usage_example`, `pdb_debug_trace`) should be invoked next to gather the missing information?**  
   - If the context is already sufficient, leave this blank.

Instructions:
- If the test failure's root cause is uncertain, it is better to request more information than to guess.
- Remember: you have access to multiple code search APIs. Use them when needed.
"""

SPECIFY_API_PRMPT = """Based on above analysis, What API calls are needed?"""



def generator(
    prefix_thread: MessageThread
) -> Generator[tuple[str, MessageThread], tuple[str, bool] | None, None]:
    
    msg_thread = prefix_thread.copy()
    msg_thread.add_user(SELECT_PROMPT)

    while True:
        # first call is to select some APIs to call
        logger.debug("<Agent Search in Debugger> Selecting APIs to call.")
        res_text, *_ = common.SELECTED_MODEL.call(msg_thread.to_msg())
        msg_thread.add_model(res_text)
        print_retrieval(res_text, "API Selection in Debugger")

        # the search result should be sent here by our backend AST search tool
        generator_input = yield res_text, msg_thread
        assert generator_input is not None
        search_result, re_search = generator_input

        if re_search:
            # the search APIs selected have some issue
            logger.debug(
                "<Agent Search in Debugger> Downstream could not consume our last response. Will retry."
            )
            msg_thread.add_user(search_result)
            continue

        # the search APIs selected are ok and the results are back
        # second call is to analyze the search results
        logger.debug("<Agent Search in Debugger> Analyzing search results.")
        msg_thread.add_user(search_result)
        msg_thread.add_user(QUESTION_PROMPT)
        print_acr(QUESTION_PROMPT, "Agent Debugger Question")

        res_text, *_ = common.SELECTED_MODEL.call(msg_thread.to_msg())
        msg_thread.add_model(res_text)
        print_retrieval(res_text, "Agent Debugger Analysis")
        msg_thread.add_user(SPECIFY_API_PRMPT)

    