import os
from dataclasses import dataclass, field
from typing import Optional, Union
from PIL import Image
import pandas as pd

from src.logger import get_logger
from src.modeling.llava.eval import (
    Evaluator,
    LLMValidator,
    LLMValidatorConfig,
    MCQuestion,
    RegexValidator,
)
from src.modeling.multimodal import ImageTextModelConfig


logger = get_logger(__name__)


def question_en(row: pd.Series) -> str:
    text = ""
    if pd.notna(row["hint"]):
        text += f"Hint: {row['hint']}\n"

    text += f"{row['question']}\n\n"
    text += "Options:\n"
    for option in ["A", "B", "C", "D"]:
        if pd.notna(row[option]):
            text += f"{option}. {row[option]}\n"

    text += "\n"
    text += "Answer using only one of the letters A, B, C, or D.\n"
    text += "Your answer: "

    return text

def question_pl(row: pd.Series) -> str:
    text = ""
    if pd.notna(row["hint"]):
        text += f"Podpowiedź: {row['hint']}\n"

    text += f"{row['question']}\n\n"
    text += "Opcje:\n"
    for option in ["A", "B", "C", "D"]:
        if pd.notna(row[option]):
            text += f"{option}. {row[option]}\n"

    text += "\n"
    text += "Odpowiedz używając tylko jednej z liter A, B, C lub D.\n"
    text += "Twoja odpowiedź: "
    return text


def paligemma_question_en(row: pd.Series) -> str:
    text = "answer en\n"
    if pd.notna(row["hint"]):
        text += f"Hint: {row['hint']}\n"

    text += f"{row['question']}\n\n"
    text += "Options:\n"
    for option in ["A", "B", "C", "D"]:
        if pd.notna(row[option]):
            text += f"{option}. {row[option]}\n"

    text += "\n"
    text += "Answer using only one of the letters A, B, C, or D.\n"
    text += "Your answer: "

    return text


def paligemma_question_pl(row: pd.Series) -> str:
    text = "answer pl\n"
    if pd.notna(row["hint"]):
        text += f"Podpowiedź: {row['hint']}\n"

    text += f"{row['question']}\n\n"
    text += "Opcje:\n"
    for option in ["A", "B", "C", "D"]:
        if pd.notna(row[option]):
            text += f"{option}. {row[option]}\n"

    text += "\n"
    text += "Odpowiedz używając tylko jednej z liter A, B, C lub D.\n"
    text += "Twoja odpowiedź: "
    return text


def evaluate(
    model_name_or_path: str,
    mmbench_df: pd.DataFrame,
    validator: Union[RegexValidator, LLMValidator],
    lang: Optional[str] = None,
    images_dir: Optional[str] = None,
) -> pd.DataFrame:
    q_en = paligemma_question_en if "paligemma" in model_name_or_path.lower() else question_en
    q_pl = paligemma_question_pl if "paligemma" in model_name_or_path.lower() else question_pl

    questions = [
        MCQuestion(
            index=row["index"],
            category=row["category"],
            question=q_en(row) if lang == "en" else q_pl(row),
            image=Image.open(os.path.join(images_dir, row["image"])),
            correct=row["answer"],
        )
        for _, row in mmbench_df.iterrows()
    ]

    logger.info("Using model: %s", model_name_or_path)
    config = ImageTextModelConfig(model_name_or_path=model_name_or_path)
    evaluator = Evaluator(config, validator)

    logger.info("Evaluating %d questions on language %s", len(questions), lang)
    prompt = lambda x: x.question
    results = evaluator.multichoice(questions, prompt)

    logger.info("Evaluation completed")
    return results


@dataclass
class EvaluationArgs:
    model_name_or_path: str = field(metadata={"help": "Path to the model."})
    input_file: str = field(
        metadata={"help": "Path to the input file containing MMBench data (legacy format)."}
    )
    output_file: str = field(metadata={"help": "Path to save the evaluation results."})
    images_dir: Optional[str] = field(
        default=None,
        metadata={
            "help": "Directory containing images if the input file contains image filenames instead of base64 strings."
        },
    )
    llm_extractor_config: Optional[str] = field(
        default=None,
        metadata={"help": "Path to the LLM extractor configuration YAML file."},
    )
    lang: Optional[str] = field(default=None, metadata={"help": "Language of the questions."})


if __name__ == "__main__":
    from transformers import HfArgumentParser

    parser = HfArgumentParser(EvaluationArgs)
    (args,) = parser.parse_args_into_dataclasses()

    logger.info("Loading MMBench data from %s", args.input_file)
    mmbench_df = pd.read_json(args.input_file, lines=args.input_file.endswith(".jsonl"))

    validator = None
    if args.llm_extractor_config:
        logger.info("Loading LLM extractor configuration from %s", args.validator_config)
        config = LLMValidatorConfig.from_yaml(args.llm_extractor_config)
        validator = LLMValidator(config)
        logger.info("LLM-extraction based validator initialized")
    else:
        validator = RegexValidator()
        logger.info("Regex-based validator initialized")

    results = evaluate(
        args.model_name_or_path,
        mmbench_df,
        validator,
        lang=args.lang,
        images_dir=args.images_dir,
    )

    accuracy = results["is_correct"].mean()
    logger.info("Overall accuracy: %.2f%%", accuracy * 100)

    unmatched_rate = results["answer_matched"].isna().mean()
    logger.info("Unmatched answers rate: %.2f%%", unmatched_rate * 100)

    logger.info("Saving results to %s", args.output_file)
    os.makedirs(os.path.dirname(args.output_file), exist_ok=True)
    use_lines = args.output_file.endswith(".jsonl")
    results.to_json(
        args.output_file,
        orient="records",
        lines=use_lines,
        indent=2 if not use_lines else None,
        force_ascii=False,
    )

    logger.info("Evaluation completed successfully.")
