from typing import TypedDict, List, Dict, Any

from repenseai.genai.benchmark.core.base_provider import BaseagentProvider
from repenseai.genai.benchmark.core.base_test import BaseTest


class Question(TypedDict):
    text: str
    options: List[str]
    correct_answer: str


class TestValidationError(Exception):
    pass


class OptionQuestionTest(BaseTest):
    def __init__(
        self,
        name: str,
        questions: List[Question],
        description: str = None,
        prompt_template: str = None,
    ):
        super().__init__(name, description)
        self.questions = questions

        DEFAULT_PROMPT_TEMPLATE = (
            "Question {question_number}:\n"
            "{question_text}\n\n"
            "Options:\n"
            "{options_text}\n\n"
            "Please respond with only the letter of the correct option (A, B, C, etc.).\n"
            "Do not explain your answer or add any additional text."
        )

        self.prompt_template = prompt_template or DEFAULT_PROMPT_TEMPLATE

    def validate(self) -> bool:
        try:
            self._validate_questions()
            return True
        except TestValidationError as e:
            print(e)
            return False

    def _validate_questions(self) -> None:
        for i, question in enumerate(self.questions, 1):
            if not all(
                key in question for key in ["text", "options", "correct_answer"]
            ):
                raise TestValidationError(f"Question {i} is missing required fields")
            if (
                not isinstance(question["options"], list)
                or len(question["options"]) < 2
            ):
                raise TestValidationError(f"Question {i} must have at least 2 options")
            if not any(
                question["correct_answer"] in option for option in question["options"]
            ):
                raise TestValidationError(
                    f"Question {i}'s correct answer must be in options"
                )

    async def run(self, agent_provider: BaseagentProvider) -> Dict[str, Any]:
        results = []

        for i, question in enumerate(self.questions, 1):
            prompt = self._create_prompt(question, i)
            response = await agent_provider.generate(prompt)

            results.append(
                {
                    "question_number": i,
                    "question": question["text"],
                    "options": question["options"],
                    "correct_answer": question["correct_answer"],
                    "agent_answer": response,
                }
            )

        return {"test_type": "options_exam", "results": results}

    def _create_prompt(self, question: Question, question_number: int) -> str:
        options_text = "\n".join(
            f"{chr(65 + i)}. {option}" for i, option in enumerate(question["options"])
        )

        return self.prompt_template.format(
            question_number=question_number,
            question_text=question["text"],
            options_text=options_text,
        )
