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

import inspect
from typing import Any

from loguru import logger

from app.data_structures import MessageThread
from app.model import common
from app.post_process import ExtractStatus, is_valid_json

SYSTEM_PROMPT = """You are an intelligent software developer working on a code generation task.
The task contains a description marked between <issue> and </issue>.
Your task is to extract the code generation location in the text into a json file like this:
{
    "Locations": [{"file": "path/to/file", "class": "class_name", "method": "func_name", "start_line": start_line, "end_line": end_line}, ... ]
}
Leave "class" / "start_line" / "end_line" empty if you didn't detect them.
For example:
<issue>
 The function you need to generate is in `projectA/path/file.py` at Class1.func_1.                                                                     
 The function signature and docstring of the code you need to generate are:
 ```                                                                          
 def func_1(                                                                   
         param1,
         param2,
...
```
</issue>
Your return should be:
{
    "Locations": [{"file": "path/to/file.py", "class": "Continuous", "method": "label", "start_line": None, "end_line": None}]
}

"""

def run_with_retries(text: str, retries=5) -> tuple[str | None, list[MessageThread]]:
    msg_threads = []
    for idx in range(1, retries + 1):
        logger.debug(
            "Trying to extract code gen location into json. Try {} of {}.",
            idx,
            retries,
        )

        res_text, new_thread = run(text)
        msg_threads.append(new_thread)

        extract_status, data = is_valid_json(res_text)

        if extract_status != ExtractStatus.IS_VALID_JSON:
            logger.debug("Invalid json. Will retry.")
            continue

        valid, diagnosis = is_valid_response(data) # diagnosis is the problem description of is_valid_response func.
        if not valid:
            logger.debug(f"{diagnosis}. Will retry.")
            continue
        logger.debug("Extracted a valid json.")
        return res_text, msg_threads
    return None, msg_threads


def run(text: str) -> tuple[str, MessageThread]:
    """
    Run the agent to extract issue to json format.
    """
    msg_thread = MessageThread()
    msg_thread.add_system(SYSTEM_PROMPT)
    msg_thread.add_user(text)
    res_text, *_ = common.SELECTED_MODEL.call(
        msg_thread.to_msg(), response_format="json_object"
    )

    msg_thread.add_model(res_text, [])  # no tools
    return res_text, msg_thread

def is_valid_response(data: Any) -> tuple[bool, str]:
    if not isinstance(data, dict):
        return False, "Json is not a dict"

    if data.get("Locations"):
        Locations = data.get("Locations")
        if not isinstance(Locations, list) or not Locations:
            return False, "Locations are empty"
        if len(Locations) == 0:
            return False, "Locations can't be a empty list"
        
        for loc in data["Locations"]:
            if not isinstance(loc, dict):
                return False, "Every location must be a dict"
            if loc.get("file") and (loc.get("class") or loc.get("method")):
                continue
            else:
                return (
                    False,
                    "Locations not detailed enough. Each locations must at least contain a file path and one of class and method.",
                )
    else:
            return False, "Json should only contain one key-value pair. Key is 'Locations' and the value should be a list"
    return True, "OK"