from dataclasses import dataclass, field
import os
import json
import pandas as pd
from typing import Optional

from src.generator import GenerationConfig, LLMGenerator, ModelConfig
from src.logger import get_logger

logger = get_logger(__name__)


def build_question(row: pd.Series) -> str:
    prompt = row["question"]
    if row["hint"] and not pd.isna(row["hint"]):
        prompt += "\n"
        prompt += row["hint"]

    prompt += "\n"

    kyes = ["A", "B", "C", "D"]
    choices = {key: row[key] for key in kyes if row[key] and not pd.isna(row[key])}
    for key, choice in choices.items():
        prompt += f"{key}. {choice}\n"

    return prompt.strip()


def judge(
    source_target_pairs: list[tuple[str, str]],
    judge_config: Optional[ModelConfig],
    generation_config: GenerationConfig,
) -> list[Optional[dict]]:
    if judge_config is None:
        return [None] * len(source_target_pairs)

    logger.info("Starting judging process...")
    logger.info(f"Using model configuration: {judge_config}")

    prompt = lambda sample: judge_config.prompt_template.format(
        src_lang="English",
        target_lang="Polish",
        src_text=sample[0],
        target_text=sample[1],
    )

    with LLMGenerator(judge_config, generation_config) as generator:
        judgments = generator.generate(source_target_pairs, prompt=prompt)

    parsed_judgments = []
    for judgment in judgments:
        judgement = judgment.strip()
        try:
            parsed_judgement = json.loads(judgement)
            parsed_judgments.append(parsed_judgement)
        except json.JSONDecodeError:
            logger.warning(f"Failed to parse judgment: {judgment}")
            parsed_judgments.append(None)

    total, parsed = len(judgments), sum(1 for j in parsed_judgments if j is not None)
    logger.info(f"Judging process completed. Parsed {parsed}/{total} judgements.")
    return parsed_judgments


def translate(
    mmbench_df: pd.DataFrame,
    translator_config: ModelConfig,
    generation_config: GenerationConfig,
    judge_config: Optional[ModelConfig] = None,
) -> pd.DataFrame:
    logger.info("Starting translation process...")
    logger.info(f"Using model configuration: {translator_config}")

    samples = []
    for i, row in mmbench_df.iterrows():
        question = build_question(row)
        samples.append(question)

    logger.info(f"Total samples to translate: {len(samples)}")
    prompt = lambda text: translator_config.prompt_template.format(text=text)

    with LLMGenerator(translator_config, generation_config) as generator:
        translations = generator.generate(samples, prompt=prompt)

    translated_df = mmbench_df.copy()
    translated_df["en"] = samples
    translated_df["pl"] = translations
    translated_df["translation_quality"] = judge(
        list(zip(translated_df["en"], translated_df["pl"])),
        judge_config,
        generation_config,
    )

    logger.info("Translation process completed.")
    return translated_df


@dataclass
class TranslationArgs:
    input: str = field(metadata={"help": "Path to the input file"})
    output: str = field(metadata={"help": "Path to the output file"})
    translator_config: str = field(
        metadata={"help": "Path to the translation model configuration file"}
    )
    judge_config: Optional[str] = field(
        default=None, metadata={"help": "Path to the judging model configuration file"}
    )
    remove_images: bool = field(
        default=False,
        metadata={"help": "Whether to remove image data from the dataset"},
    )


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

    parser = HfArgumentParser([TranslationArgs, GenerationConfig])
    (
        args,
        generation_config,
    ) = parser.parse_args_into_dataclasses()

    data = pd.read_csv(args.input, sep="\t")
    # data = data.sample(n=10, random_state=42)

    # Setup configurations
    translator_config = ModelConfig.from_yaml(args.translator_config)

    judge_config = None
    if args.judge_config:
        judge_config = ModelConfig.from_yaml(args.judge_config)

    # Perform translation
    result = translate(data, translator_config, generation_config, judge_config)

    # Remove image data if specified
    if args.remove_images:
        logger.info("Removing image data from the dataset...")
        result = result.drop(columns=["image"])

    # Save results
    os.makedirs(os.path.dirname(args.output), exist_ok=True)
    use_lines = args.output.endswith(".jsonl")
    result.to_json(
        args.output,
        orient="records",
        lines=use_lines,
        indent=2 if not use_lines else None,
    )
    logger.info(f"Translation saved to {args.output}")
