from __future__ import annotations

from datasets import Dataset
from lora_data import GOALS, LESSONS, split_examples
from transformers import PreTrainedTokenizerFast

CONTEXT_LENGTH = 128


def corrupt_response(response: str, index: int) -> tuple[str, str]:
    if index % 2 == 0:
        present = next(goal for goal in GOALS if goal in response)
        replacement = GOALS[(GOALS.index(present) + 2) % len(GOALS)]
        return response.replace(present, replacement), "wrong_goal"
    present = next(lesson for lesson in LESSONS if lesson in response)
    replacement = LESSONS[(LESSONS.index(present) + 2) % len(LESSONS)]
    return response.replace(present, replacement), "wrong_lesson"


def build_pairs() -> tuple[list[dict[str, str]], list[dict[str, str]]]:
    train_examples, eval_examples = split_examples()

    def convert(examples: list[dict[str, str]], offset: int) -> list[dict[str, str]]:
        pairs = []
        for index, example in enumerate(examples):
            rejected, corruption = corrupt_response(example["response"], index + offset)
            pairs.append(
                {
                    "prompt": example["prompt"],
                    "chosen": example["response"],
                    "rejected": rejected,
                    "corruption": corruption,
                }
            )
        return pairs

    return convert(train_examples, 0), convert(eval_examples, len(train_examples))


def encode_pairs(
    pairs: list[dict[str, str]],
    tokenizer: PreTrainedTokenizerFast,
) -> Dataset:
    rows: dict[str, list] = {
        "chosen_input_ids": [],
        "chosen_attention_mask": [],
        "rejected_input_ids": [],
        "rejected_attention_mask": [],
        "corruption": [],
    }
    for pair in pairs:
        required_goal = next(goal for goal in GOALS if goal in pair["prompt"])
        required_lesson = next(lesson for lesson in LESSONS if lesson in pair["prompt"])

        def compact_candidate(
            response: str,
            goal: str = required_goal,
            lesson: str = required_lesson,
        ) -> str:
            candidate_goal = next(goal for goal in GOALS if goal in response)
            candidate_lesson = next(lesson for lesson in LESSONS if lesson in response)
            return (
                f"<bos>Required goal: {goal}. "
                f"Required lesson: {lesson}. "
                f"Candidate goal: {candidate_goal}. "
                f"Candidate lesson: {candidate_lesson}.<eos>"
            )

        chosen = tokenizer(
            compact_candidate(pair["chosen"]),
            max_length=CONTEXT_LENGTH,
            truncation=True,
            padding="max_length",
            add_special_tokens=False,
        )
        rejected = tokenizer(
            compact_candidate(pair["rejected"]),
            max_length=CONTEXT_LENGTH,
            truncation=True,
            padding="max_length",
            add_special_tokens=False,
        )
        rows["chosen_input_ids"].append(chosen["input_ids"])
        rows["chosen_attention_mask"].append(chosen["attention_mask"])
        rows["rejected_input_ids"].append(rejected["input_ids"])
        rows["rejected_attention_mask"].append(rejected["attention_mask"])
        rows["corruption"].append(pair["corruption"])
    return Dataset.from_dict(rows)