import pandas as pd
from dataclasses import dataclass
from typing import Callable, Optional, Union
from PIL import Image
import re

from src.modeling.multimodal import ImageTextModelConfig, AutoImageTextGenerator
from src.logger import get_logger
from src.generator import LLMGenerator, ModelConfig

logger = get_logger(__name__)


@dataclass
class MCQuestion:
    index: str
    question: str
    image: Image.Image
    correct: str  # correct choice_id
    category: Optional[str] = None


@dataclass
class MCQuestionResult:
    index: str
    correct: str  # correct choice_id
    is_correct: bool
    answer_raw: str
    answer_matched: Optional[str] = None


@dataclass
class ValidationResult:
    is_correct: bool
    matched: Optional[str] = None
    output: Optional[str] = None


class RegexValidator:
    def validate_multichoice(
        self, questions: list[MCQuestion], answers: list[str], lang: str = "en"
    ) -> list[ValidationResult]:
        results = []

        for question, answer in zip(questions, answers):
            match = re.match(r"^\s*([A-D])\b", answer, re.IGNORECASE)
            if match:
                answer_choice = match.group(1).upper()
                is_correct = answer_choice == question.correct
                results.append(ValidationResult(is_correct=is_correct, matched=answer_choice))
            else:
                results.append(ValidationResult(is_correct=False))

        return results


@dataclass
class LLMValidatorConfig(ModelConfig):
    match_pattern: str = r'{"matched_option"\s*:\s*"([A-D]|NO_MATCH)"}'


class LLMValidator:
    def __init__(self, config: LLMValidatorConfig):
        self.config = config

    def validate_multichoice(
        self, questions: list[MCQuestion], answers: list[str], lang: str = "en"
    ) -> list[ValidationResult]:
        logger.info("Using LLMValidator with model: %s", self.config.model_id)

        samples = list(zip(questions, answers))
        prompt = lambda sample: self.config.prompt_template.format(
            question=sample[0].question,
            answer=sample[1],
        ).strip()

        logger.info("Validating %d answers", len(answers))
        with LLMGenerator(self.config) as generator:
            outputs = generator.generate(samples, prompt)

        results = []
        for question, output in zip(questions, outputs):
            match = re.search(self.config.match_pattern, output, re.IGNORECASE)
            if match:
                answer_choice = match.group(1).upper()
                results.append(
                    ValidationResult(
                        is_correct=answer_choice == question.correct,
                        matched=answer_choice,
                        output=output,
                    )
                )
            else:
                results.append(ValidationResult(is_correct=False, output=output))

        return results


class Evaluator:
    def __init__(
        self,
        config: ImageTextModelConfig,
        validator: Union[RegexValidator, LLMValidator] = RegexValidator(),
    ):
        self.model_config = config
        self.validator = validator

    def multichoice(
        self, questions: list[MCQuestion], prompt: Callable[[MCQuestion], str]
    ) -> pd.DataFrame:
        prompts = [prompt(question) for question in questions]
        images = [question.image for question in questions]

        logger.info("Generating answers for %d questions", len(questions))
        with AutoImageTextGenerator.from_config(self.model_config) as generator:
            answers = generator.generate(prompts, images)

        validations = self.validator.validate_multichoice(questions, answers)
        return pd.DataFrame(
            {
                "index": [q.index for q in questions],
                "category": [q.category for q in questions],
                "question": [q.question for q in questions],
                "answer_raw": [a for a in answers],
                "answer_matched": [v.matched for v in validations],
                "answer_extractor_output": [v.output for v in validations],
                "correct": [q.correct for q in questions],
                "is_correct": [v.is_correct for v in validations],
            }
        )
