"""QA and specialized format converters for FullBenchmarkDownloader."""

from typing import Any, Dict, List


class QAConvertersMixin:
    """Mixin providing QA and specialized format conversion methods."""

    def _convert_winogrande_format(self, sample: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Convert Winogrande format (sentence, option1, option2, answer)."""
        sentence = sample.get("sentence", "")
        option1 = sample.get("option1", "")
        option2 = sample.get("option2", "")
        answer = sample.get("answer", "")

        if not sentence or not option1 or not option2 or not answer:
            return []

        # Determine correct and incorrect answers
        if answer == "1":
            correct_answer = option1
            incorrect_answer = option2
        elif answer == "2":
            correct_answer = option2
            incorrect_answer = option1
        else:
            # If answer format is unexpected, default to option1 as correct
            correct_answer = option1
            incorrect_answer = option2

        # Create contrastive pair
        return [
            {
                "question": sentence,  # The sentence with blank to fill
                "good_response": correct_answer,
                "bad_response": incorrect_answer,
                "metadata": {
                    "option1": option1,
                    "option2": option2,
                    "answer": answer,
                    "benchmark_type": "winogrande",
                    "task_type": "coreference_resolution",
                },
            }
        ]

    def _convert_wikitext_format(self, sample: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Convert WikiText format (page)."""
        page = sample.get("page", "")

        if not page or len(page.strip()) < 50:  # Skip very short pages
            return []

        # For WikiText, we create language modeling pairs
        # Split the page into sentences and create good/corrupted pairs
        sentences = page.split(". ")
        if len(sentences) < 2:
            return []

        pairs = []
        for i, sentence in enumerate(sentences):
            if len(sentence.strip()) > 20:  # Only use substantial sentences
                # Create a corrupted version by replacing some words
                words = sentence.split()
                if len(words) > 3:
                    # Simple corruption: duplicate a word in the middle
                    mid_idx = len(words) // 2
                    corrupted_words = words.copy()
                    corrupted_words.insert(mid_idx, words[mid_idx])
                    corrupted_sentence = " ".join(corrupted_words)

                    pairs.append(
                        {
                            "question": "Complete the text naturally:",
                            "good_response": sentence.strip(),
                            "bad_response": corrupted_sentence,
                            "metadata": {
                                "benchmark_type": "wikitext",
                                "task_type": "language_modeling",
                                "sentence_index": i,
                            },
                        }
                    )

                    # Limit to 3 pairs per page to avoid too many
                    if len(pairs) >= 3:
                        break

        return pairs

    def _convert_webqs_format(self, sample: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Convert WebQS format (question, answers list)."""
        question = sample.get("question", "")
        answers = sample.get("answers", [])

        if not question or not answers:
            return []

        # Take the first answer as the correct one
        correct_answer = answers[0] if answers else ""

        if not correct_answer:
            return []

        # Generate incorrect answers (simple approach)
        incorrect_answers = []

        # Strategy 1: Use other answers from the same dataset if available
        if len(answers) > 1:
            incorrect_answers.extend(answers[1:3])  # Take up to 2 more answers as distractors

        # Strategy 2: Generate simple incorrect answers
        if len(incorrect_answers) < 2:
            # Simple factual distractors
            incorrect_answers.append("Unknown")
            incorrect_answers.append("No information available")

        # Create contrastive pairs
        pairs = []
        for incorrect in incorrect_answers[:2]:  # Limit to 2 pairs
            pairs.append(
                {
                    "question": question,
                    "good_response": correct_answer,
                    "bad_response": incorrect,
                    "metadata": {"benchmark_type": "webqs", "task_type": "factual_qa", "url": sample.get("url", "")},
                }
            )

        return pairs

    def _convert_naturalqs_format(self, sample: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Convert NaturalQS format (question, answer as list)."""
        question = sample.get("question", "")
        answer_list = sample.get("answer", [])

        if not question or not answer_list:
            return []

        # Take the first answer as the correct one (shortest/most direct)
        correct_answer = answer_list[0] if answer_list else ""

        if not correct_answer:
            return []

        # Generate incorrect answers
        incorrect_answers = []

        # Strategy 1: Use other answers from the list as distractors if available
        if len(answer_list) > 1:
            incorrect_answers.extend(answer_list[1:3])  # Take up to 2 more answers

        # Strategy 2: Generate generic incorrect answers
        if len(incorrect_answers) < 2:
            incorrect_answers.append("I don't know the answer to this question.")
            incorrect_answers.append("This information is not available.")

        # Create contrastive pairs
        pairs = []
        for incorrect in incorrect_answers[:2]:  # Limit to 2 pairs
            pairs.append(
                {
                    "context": question,
                    "good_response": correct_answer,
                    "bad_response": incorrect,
                    "metadata": {
                        "benchmark_type": "naturalqs",
                        "task_type": "factual_qa",
                        "total_answers": len(answer_list),
                    },
                }
            )

        return pairs

    def _convert_triviaqa_format(self, sample: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Convert TriviaQA format (question, answer as dict with aliases)."""
        question = sample.get("question", "")
        answer_dict = sample.get("answer", {})

        if not question or not answer_dict:
            return []

        # Extract the correct answer from aliases
        aliases = answer_dict.get("aliases", [])
        if not aliases:
            # Use other fields
            correct_answer = (
                answer_dict.get("value", "") or answer_dict.get("normalized_value", "") or str(answer_dict)
            )[:100]  # Truncate if too long
        else:
            correct_answer = aliases[0]  # Use first alias as primary answer

        if not correct_answer:
            return []

        # Generate incorrect answers
        incorrect_answers = []

        # Strategy 1: Use other aliases as distractors if available
        if len(aliases) > 1:
            incorrect_answers.extend(aliases[1:3])  # Take up to 2 more aliases

        # Strategy 2: Generate generic incorrect answers for trivia
        if len(incorrect_answers) < 2:
            incorrect_answers.append("Unknown")
            incorrect_answers.append("I don't know")

        # Create contrastive pairs
        pairs = []
        for incorrect in incorrect_answers[:2]:  # Limit to 2 pairs
            pairs.append(
                {
                    "context": question,
                    "good_response": correct_answer,
                    "bad_response": incorrect,
                    "metadata": {
                        "benchmark_type": "triviaqa",
                        "task_type": "trivia_qa",
                        "total_aliases": len(aliases),
                        "entity_name": answer_dict.get("matched_wiki_entity_name", ""),
                    },
                }
            )

        return pairs
