from __future__ import annotations

from typing import Any
from wisent.core.cli_logger import setup_logger

from wisent.core.contrastive_pairs.core.pair import ContrastivePair
from wisent.core.contrastive_pairs.huggingface_pairs.atoms import HuggingFaceBenchmarkExtractor

__all__ = ["MBPPExtractor"]

log = setup_logger(__name__)

task_names = ("mbpp_plus",)

class MBPPExtractor(HuggingFaceBenchmarkExtractor):
    """
    Extractor for MBPP (Mostly Basic Python Problems) dataset.

    Schema (mbpp or google-research/mbpp):
        - text: str (problem description)
        - code: str (correct solution code)
        - test_list: list[str] (test cases)
    """


    evaluator_name = "coding"

    def extract_contrastive_pairs(
        self,
        limit: int | None = None,
    ) -> list[ContrastivePair]:
        """
        Build contrastive pairs from MBPP examples.

        Args:
            limit: Optional maximum number of pairs to produce.

        Returns:
            A list of ContrastivePair objects.
        """
        max_items = self._normalize_limit(limit)

        # Load dataset
        docs = self.load_dataset(
            dataset_name="mbpp",
            split="test",
            limit=max_items,
        )

        pairs: list[ContrastivePair] = []

        log.info(f"Extracting contrastive pairs from {len(docs)} MBPP examples")

        for doc in docs:
            pair = self._extract_pair_from_doc(doc)
            if pair is not None:
                pairs.append(pair)
                if max_items is not None and len(pairs) >= max_items:
                    break

        if not pairs:
            log.warning("No valid MBPP pairs extracted")

        return pairs

    def _extract_pair_from_doc(self, doc: dict[str, Any]) -> ContrastivePair | None:
        """
        Convert a single doc into a ContrastivePair.

        Returns None when required fields are missing or malformed.
        """
        try:
            problem_text = doc.get("text", "").strip()
            correct_code = doc.get("code", "").strip()
            test_list = doc.get("test_list", [])
            test_imports = doc.get("test_imports", [])

            if not problem_text or not correct_code:
                log.debug("Skipping: missing problem text or code")
                return None

            # Create incorrect code (add syntax error or logical error)
            incorrect_code = self._create_incorrect_code(correct_code)

            # Format the prompt
            formatted_prompt = f"{problem_text}\n\nWrite a Python function to solve this problem."

            # Build test code from test_list
            # The test_list contains assertion strings like "assert function_name(...) == expected"
            test_code = ""
            entry_point = None
            if test_list:
                # Extract function name from correct_code
                import re
                func_match = re.search(r'def\s+(\w+)\s*\(', correct_code)
                if func_match:
                    entry_point = func_match.group(1)
                
                # Add imports if provided
                if test_imports:
                    test_code += "\n".join(test_imports) + "\n\n"

                # Build test function that uses 'candidate' parameter
                test_code += "def check(candidate):\n"
                for test_case in test_list:
                    # Replace function name with 'candidate' in assertions
                    if entry_point:
                        modified_test = test_case.replace(f"{entry_point}(", "candidate(")
                    else:
                        modified_test = test_case
                    test_code += f"    {modified_test}\n"

            metadata = {
                "label": "mbpp",
                "source": "mbpp",
                "test_code": test_code if test_code else None,
                "entry_point": entry_point,
            }

            return self._build_pair(
                question=formatted_prompt,
                correct=correct_code,
                incorrect=incorrect_code,
                metadata=metadata,
            )

        except Exception as exc:
            log.error(f"Error extracting pair from doc: {exc}", exc_info=True)
            return None

    def _create_incorrect_code(self, correct: str) -> str:
        """Create an incorrect version of the code."""
        # Add a syntax error by removing closing parenthesis
        if "(" in correct and ")" in correct:
            # Find last closing paren and remove it
            idx = correct.rfind(")")
            if idx > 0:
                return correct[:idx] + correct[idx+1:]

        # Fallback: add comment that breaks the code
        return correct + "\n# Missing return statement"

