"""
LM-Evaluation-Harness task wrapper for task-agnostic architecture.
"""

from typing import Any, Dict, List, Optional

from ..benchmark_extractors import BenchmarkExtractor
from ..benchmark_extractors import get_extractor as get_benchmark_extractor
from ..contrastive_pairs.lm_eval_pairs.lm_extractor_registry import get_extractor as get_lm_extractor
from ..task_interface import TaskInterface
from ..utils.dataset_splits import get_test_docs, get_all_docs_from_task, create_deterministic_split


def get_extractor(task_name: str) -> BenchmarkExtractor:
    """Get extractor, trying lm_extractor_registry first, then benchmark_extractors."""
    try:
        return get_lm_extractor(task_name)
    except Exception:
        return get_benchmark_extractor(task_name)


class LMEvalTask(TaskInterface):
    """Wrapper for lm-evaluation-harness tasks."""

    def __init__(self, task_name: str, description: str, categories: List[str]):
        self.task_name = task_name
        self._description = description
        self._categories = categories
        self._extractor = get_extractor(task_name)

    def load_data(self, limit: Optional[int] = None) -> List[Dict[str, Any]]:
        """
        Load TEST data from lm-eval using our unified split strategy.

        All available splits are combined and split 80/20 into train/test.
        This method returns the TEST portion for evaluation.
        Training portion is used for contrastive pair generation (see atoms.py).
        """
        try:
            from lm_eval.tasks import get_task_dict

            # Get task directly from lm-eval
            task_dict = get_task_dict([self.task_name])
            if self.task_name not in task_dict:
                print(f"Warning: Task '{self.task_name}' not found in lm-eval")
                return []

            task = task_dict[self.task_name]

            # Get ALL docs from all splits
            all_docs, split_counts = get_all_docs_from_task(task)

            if not all_docs:
                print(f"Warning: No documents found for task '{self.task_name}'")
                return []

            # Apply our 80/20 split and get TEST docs only
            _, test_docs = create_deterministic_split(all_docs, self.task_name)

            # Ensure docs are in dictionary format
            processed_docs = []
            for doc in test_docs:
                if isinstance(doc, dict):
                    processed_docs.append(doc)
                elif isinstance(doc, str):
                    processed_docs.append({"text": doc})
                else:
                    try:
                        processed_docs.append(dict(doc))
                    except:
                        processed_docs.append({"data": str(doc)})

            docs = processed_docs

            # Apply limit if specified
            if limit and len(docs) > limit:
                docs = docs[:limit]

            total = len(all_docs)
            test_count = len(test_docs)
            returned = len(docs)
            print(f"Loaded {returned} test docs from {self.task_name} "
                  f"(total: {total}, test split: {test_count}, original splits: {split_counts})")

            return docs

        except Exception as e:
            print(f"Warning: Could not load lm-eval task '{self.task_name}': {e}")
            return []

    def get_extractor(self) -> BenchmarkExtractor:
        """Get the benchmark extractor for this task."""
        return self._extractor

    def get_name(self) -> str:
        """Get the task name."""
        return self.task_name

    def get_description(self) -> str:
        """Get the task description."""
        return self._description

    def get_categories(self) -> List[str]:
        """Get the task categories."""
        return self._categories


class MBPPTask(LMEvalTask):
    """MBPP task implementation."""

    def __init__(self):
        super().__init__(
            task_name="mbpp",
            description="MBPP: Mostly Basic Python Problems coding benchmark",
            categories=["coding", "reasoning", "python"],
        )


class HumanEvalTask(LMEvalTask):
    """HumanEval task implementation."""

    def __init__(self):
        super().__init__(
            task_name="humaneval",
            description="HumanEval: Human Evaluation of Python coding problems",
            categories=["coding", "reasoning", "python"],
        )


class MBPPPlusTask(LMEvalTask):
    """MBPP Plus task implementation."""

    def __init__(self):
        super().__init__(
            task_name="mbpp_plus",
            description="MBPP Plus: Extended version of MBPP with additional test cases",
            categories=["coding", "reasoning", "python"],
        )


class GSM8KTask(LMEvalTask):
    """GSM8K task implementation."""

    def __init__(self):
        super().__init__(
            task_name="gsm8k",
            description="GSM8K: Grade School Math 8K problems",
            categories=["mathematics", "reasoning", "arithmetic"],
        )

class ArithmeticBaseTask(LMEvalTask):
    """Base class for arithmetic tasks. Uses unified split strategy."""
    pass  # No longer needs special handling - unified split combines all available docs


class Arithmetic1dcTask(ArithmeticBaseTask):
    """Arithmetic 1dc task implementation"""

    def __init__(self):
        super().__init__(
            task_name="arithmetic_1dc",
            description="Arithmetic 1dc: 1 digit addition arithmetic problems",
            categories=["mathematics", "arithmetic"]
        )


class Arithmetic2daTask(ArithmeticBaseTask):
    """Arithmetic 2da task implementation"""

    def __init__(self):
        super().__init__(
            task_name="arithmetic_2da",
            description="Arithmetic 2da: 2 digit addition arithmetic problems",
            categories=["mathematics", "arithmetic"]
        )

class Arithmetic2dmTask(ArithmeticBaseTask):
    """Arithmetic 2dm task implementation"""

    def __init__(self):
        super().__init__(
            task_name="arithmetic_2dm",
            description="Arithmetic 2dm: 2 digit multiplication arithmetic problems",
            categories=["mathematics", "arithmetic"]
        )

class Arithmetic2dsTask(ArithmeticBaseTask):
    """Arithmetic 2ds task implementation"""

    def __init__(self):
        super().__init__(
            task_name="arithmetic_2ds",
            description="Arithmetic 2ds: 2 digit subtraction arithmetic problems",
            categories=["mathematics", "arithmetic"]
        )

class Arithmetic3daTask(ArithmeticBaseTask):
    """Arithmetic 3da task implementation"""

    def __init__(self):
        super().__init__(
            task_name="arithmetic_3da",
            description="Arithmetic 3da: 3 digit addition arithmetic problems",
            categories=["mathematics", "arithmetic"]
        )

class Arithmetic3dsTask(ArithmeticBaseTask):
    """Arithmetic 3ds task implementation"""

    def __init__(self):
        super().__init__(
            task_name="arithmetic_3ds",
            description="Arithmetic 3ds: 3 digit subtraction arithmetic problems",
            categories=["mathematics", "arithmetic"]
        )

class Arithmetic4daTask(ArithmeticBaseTask):
    """Arithmetic 4da task implementation"""

    def __init__(self):
        super().__init__(
            task_name="arithmetic_4da",
            description="Arithmetic 4da: 4 digit addition arithmetic problems",
            categories=["mathematics", "arithmetic"]
        )

class Arithmetic4dsTask(ArithmeticBaseTask):
    """Arithmetic 4ds task implementation"""

    def __init__(self):
        super().__init__(
            task_name="arithmetic_4ds",
            description="Arithmetic 4ds: 4 digit subtraction arithmetic problems",
            categories=["mathematics", "arithmetic"]
        )

class Arithmetic5daTask(ArithmeticBaseTask):
    """Arithmetic 5da task implementation"""

    def __init__(self):
        super().__init__(
            task_name="arithmetic_5da",
            description="Arithmetic 5da: 5 digit addition arithmetic problems",
            categories=["mathematics", "arithmetic"]
        )

class Arithmetic5dsTask(ArithmeticBaseTask):
    """Arithmetic 5ds task implementation"""

    def __init__(self):
        super().__init__(
            task_name="arithmetic_5ds",
            description="Arithmetic 5ds: 5 digit subtraction arithmetic problems",
            categories=["mathematics", "arithmetic"]
        )

class MULTIRCTask(LMEvalTask):
    """MULTIRC task implementation. Uses unified split strategy."""

    def __init__(self):
        super().__init__(
            task_name="multirc",
            description="MultiRC: Multi-Sentence Reading Comprehension",
            categories=["reasoning", "long context", "general knowledge"]
        )
    # No longer needs special handling - unified split combines all available docs


class TruthfulQATask(LMEvalTask):
    """TruthfulQA task implementation. Uses unified split strategy."""

    def __init__(self):
        super().__init__(
            task_name="truthfulqa_mc1",
            description="TruthfulQA: Truthfulness evaluation benchmark",
            categories=["hallucination", "general-knowledge", "reasoning"],
        )
    # No longer needs special handling - unified split combines all available docs


class MMLUTask(LMEvalTask):
    """MMLU task implementation."""

    def __init__(self):
        super().__init__(
            task_name="mmlu",
            description="MMLU: Massive Multitask Language Understanding",
            categories=["general-knowledge", "science", "reasoning"],
        )

class QA4MRETask(LMEvalTask):
    """QA4MRE task implementation"""

    def __init__(self):
        super().__init__(
            task_name="qa4mre_2013",
            description="QA4MRE: Question Answering for Machine Reading Evaluation",
            categories=["multiple-choice", "long context"],
        )


# === CODING TASKS ===


class InstructHumanEvalTask(LMEvalTask):
    """InstructHumanEval task implementation."""

    def __init__(self):
        super().__init__(
            task_name="instructhumaneval",
            description="InstructHumanEval: Instruction-following HumanEval benchmark",
            categories=["coding", "reasoning", "python", "instruction-following"],
        )


class HumanEvalPlusTask(LMEvalTask):
    """HumanEval Plus task implementation."""

    def __init__(self):
        super().__init__(
            task_name="humaneval_plus",
            description="HumanEval Plus: Extended HumanEval with more tests",
            categories=["coding", "reasoning", "python"],
        )


class ConalaTask(LMEvalTask):
    """Conala task implementation."""

    def __init__(self):
        super().__init__(
            task_name="conala",
            description="Conala: Code generation from natural language",
            categories=["coding", "reasoning", "python", "nl2code"],
        )


class ConcodeTask(LMEvalTask):
    """Concode task implementation."""

    def __init__(self):
        super().__init__(
            task_name="concode",
            description="Concode: Code completion benchmark",
            categories=["coding", "reasoning", "completion"],
        )


class MercuryTask(LMEvalTask):
    """Mercury task implementation."""

    def __init__(self):
        super().__init__(
            task_name="mercury",
            description="Mercury: Code generation benchmark",
            categories=["coding", "reasoning"],
        )


class AppsTask(LMEvalTask):
    """APPS task implementation."""

    def __init__(self):
        super().__init__(
            task_name="apps",
            description="APPS: Automated Programming Problems Synthesis",
            categories=["coding", "reasoning", "python", "competitive"],
        )


class DS1000Task(LMEvalTask):
    """DS1000 task implementation."""

    def __init__(self):
        super().__init__(
            task_name="ds1000",
            description="DS1000: Data Science coding tasks",
            categories=["coding", "reasoning", "python", "data-science"],
        )


class MultiplePyTask(LMEvalTask):
    """Multiple-Py task implementation."""

    def __init__(self):
        super().__init__(
            task_name="multiple_py",
            description="Multiple-Py: Multi-language Python tasks",
            categories=["coding", "reasoning", "python", "multi-language"],
        )


class MultipleJsTask(LMEvalTask):
    """Multiple-JS task implementation."""

    def __init__(self):
        super().__init__(
            task_name="multiple_js",
            description="Multiple-JS: Multi-language JavaScript tasks",
            categories=["coding", "reasoning", "javascript", "multi-language"],
        )


class MultipleJavaTask(LMEvalTask):
    """Multiple-Java task implementation."""

    def __init__(self):
        super().__init__(
            task_name="multiple_java",
            description="Multiple-Java: Multi-language Java tasks",
            categories=["coding", "reasoning", "java", "multi-language"],
        )


class MultipleCppTask(LMEvalTask):
    """Multiple-Cpp task implementation."""

    def __init__(self):
        super().__init__(
            task_name="multiple_cpp",
            description="Multiple-Cpp: Multi-language C++ tasks",
            categories=["coding", "reasoning", "cpp", "multi-language"],
        )


class MultipleRsTask(LMEvalTask):
    """Multiple-Rs task implementation."""

    def __init__(self):
        super().__init__(
            task_name="multiple_rs",
            description="Multiple-Rs: Multi-language Rust tasks",
            categories=["coding", "reasoning", "rust", "multi-language"],
        )


class MultipleGoTask(LMEvalTask):
    """Multiple-Go task implementation."""

    def __init__(self):
        super().__init__(
            task_name="multiple_go",
            description="Multiple-Go: Multi-language Go tasks",
            categories=["coding", "reasoning", "go", "multi-language"],
        )


class CodexglueCodeToTextPythonTask(LMEvalTask):
    """CodexGlue Code-to-Text Python task implementation."""

    def __init__(self):
        super().__init__(
            task_name="codexglue_code_to_text_python",
            description="CodexGlue Code-to-Text Python: Python code summarization",
            categories=["coding", "reasoning", "python", "code-to-text"],
        )


class CodexglueCodeToTextGoTask(LMEvalTask):
    """CodexGlue Code-to-Text Go task implementation."""

    def __init__(self):
        super().__init__(
            task_name="codexglue_code_to_text_go",
            description="CodexGlue Code-to-Text Go: Go code summarization",
            categories=["coding", "reasoning", "go", "code-to-text"],
        )


class CodexglueCodeToTextRubyTask(LMEvalTask):
    """CodexGlue Code-to-Text Ruby task implementation."""

    def __init__(self):
        super().__init__(
            task_name="codexglue_code_to_text_ruby",
            description="CodexGlue Code-to-Text Ruby: Ruby code summarization",
            categories=["coding", "reasoning", "ruby", "code-to-text"],
        )


class CodexglueCodeToTextJavaTask(LMEvalTask):
    """CodexGlue Code-to-Text Java task implementation."""

    def __init__(self):
        super().__init__(
            task_name="codexglue_code_to_text_java",
            description="CodexGlue Code-to-Text Java: Java code summarization",
            categories=["coding", "reasoning", "java", "code-to-text"],
        )


class CodexglueCodeToTextJavascriptTask(LMEvalTask):
    """CodexGlue Code-to-Text JavaScript task implementation."""

    def __init__(self):
        super().__init__(
            task_name="codexglue_code_to_text_javascript",
            description="CodexGlue Code-to-Text JavaScript: JavaScript code summarization",
            categories=["coding", "reasoning", "javascript", "code-to-text"],
        )


class CodexglueCodeToTextPhpTask(LMEvalTask):
    """CodexGlue Code-to-Text PHP task implementation."""

    def __init__(self):
        super().__init__(
            task_name="codexglue_code_to_text_php",
            description="CodexGlue Code-to-Text PHP: PHP code summarization",
            categories=["coding", "reasoning", "php", "code-to-text"],
        )


class RecodeTask(LMEvalTask):
    """Recode task implementation."""

    def __init__(self):
        super().__init__(
            task_name="recode",
            description="Recode: Perturbed HumanEval natural generation",
            categories=["coding", "reasoning", "python", "perturbation"],
        )


class Squad2Task(LMEvalTask):
    """SQuAD2 task implementation. Uses unified split strategy."""

    def __init__(self):
        super().__init__(
            task_name="squadv2",
            description="SQuAD2: Stanford Question Answering Dataset 2.0",
            categories=["reading-comprehension", "qa", "natural-language"],
        )
    # No longer needs special handling - unified split combines all available docs


class ArcEasyTask(LMEvalTask):
    """ARC-Easy task implementation."""

    def __init__(self):
        super().__init__(
            task_name="arc_easy",
            description="ARC-Easy: AI2 Reasoning Challenge (Easy Set)",
            categories=["reasoning", "science", "multiple-choice"],
        )


class ArcChallengeTask(LMEvalTask):
    """ARC-Challenge task implementation."""

    def __init__(self):
        super().__init__(
            task_name="arc_challenge",
            description="ARC-Challenge: AI2 Reasoning Challenge (Challenge Set)",
            categories=["reasoning", "science", "multiple-choice"],
        )


class HellaswagTask(LMEvalTask):
    """HellaSwag task implementation."""

    def __init__(self):
        super().__init__(
            task_name="hellaswag",
            description="HellaSwag: Commonsense NLI for sentence completion",
            categories=["reasoning", "commonsense", "multiple-choice"],
        )


class WinograndeTask(LMEvalTask):
    """WinoGrande task implementation."""

    def __init__(self):
        super().__init__(
            task_name="winogrande",
            description="WinoGrande: Large-scale Winograd Schema Challenge",
            categories=["reasoning", "commonsense", "coreference"],
        )


class PiqaTask(LMEvalTask):
    """PIQA task implementation."""

    def __init__(self):
        super().__init__(
            task_name="piqa",
            description="PIQA: Physical Interaction Question Answering",
            categories=["reasoning", "commonsense", "physical"],
        )


class BoolqTask(LMEvalTask):
    """BoolQ task implementation."""

    def __init__(self):
        super().__init__(
            task_name="boolq",
            description="BoolQ: Boolean Questions reading comprehension",
            categories=["reading-comprehension", "qa", "boolean"],
        )


class OpenbookqaTask(LMEvalTask):
    """OpenBookQA task implementation."""

    def __init__(self):
        super().__init__(
            task_name="openbookqa",
            description="OpenBookQA: Open-book science question answering",
            categories=["reasoning", "science", "multiple-choice"],
        )
