from ..models.base.model import BaseModel
from .client import LLMClient, get_default_client
from .errors import LLMGenerationError

GENERATE_REQUIREMENTS_PROMPT = """
You are auditing AI models. Your task is to generate a set of requirements on which the AI model will be tested.

A requirement is a short description of a behaviour that the model must satisfy. You will create requirements focus the following category of issues:

{issue_description}

This is the model you will be testing:

Model name: {model_name}
Model description: {model_description}

You must generate up to {count} requirements that comprehensively evaluate the behavior of this model in various scenarios related to the issue category above.
These assertions should cover both typical and edge cases relevant to the issue category.
Think step by step about the best requirements that cover all the scenarios relevant to the issue category. Then call the `generate_requirements` function with the generated requirements.
If you don’t find any relevant requirement, you can skip this test case by calling the `skip` function.
"""

GENERATE_REQUIREMENTS_FUNCTIONS = [
    {
        "name": "generate_requirements",
        "description": "Generates requirements for model testing",
        "parameters": {
            "type": "object",
            "properties": {
                "requirements": {
                    "type": "array",
                    "items": {"type": "string", "description": "A requirement the model must satisfy"},
                }
            },
            "required": ["requirements"],
        },
    },
    {
        "name": "skip",
        "description": "Skips the generation when no relevant requirements can be generated",
        "parameters": {
            "type": "object",
            "properties": {},
            "required": [],
        },
    },
]


class TestcaseRequirementsGenerator:
    def __init__(self, issue_description: str, llm_model="gpt-4", llm_temperature=0.1, llm_client: LLMClient = None):
        self.issue_description = issue_description
        self.llm_model = llm_model
        self.llm_temperature = llm_temperature
        self.llm_client = llm_client or get_default_client()

    def _make_generate_requirements_prompt(self, model: BaseModel, num_requirements: int):
        return GENERATE_REQUIREMENTS_PROMPT.format(
            issue_description=self.issue_description,
            model_name=model.meta.name,
            model_description=model.meta.description,
            count=num_requirements,
        )

    def _make_generate_requirements_functions(self):
        return GENERATE_REQUIREMENTS_FUNCTIONS

    def generate_requirements(self, model: BaseModel, max_requirements: int = 5):
        """Generates a set of requirements for a given a model."""
        prompt = self._make_generate_requirements_prompt(model, max_requirements)
        functions = self._make_generate_requirements_functions()
        out = self.llm_client.complete(
            messages=[{"role": "system", "content": prompt}],
            functions=functions,
            function_call={"name": "generate_requirements"},
            model=self.llm_model,
            temperature=self.llm_temperature,
        )

        if out.function_call is None or "requirements" not in out.function_call.args:
            raise LLMGenerationError("Could not parse test case requirements.")

        return out.function_call.args["requirements"]
