from __future__ import annotations

from typing import Any
from wisent.core.cli.cli_logger import setup_logger
import requests
import io
import random
import re

from wisent.core.contrastive_pairs.pair import ContrastivePair
from wisent.extractors.hf._registry.manifest.atoms import HuggingFaceBenchmarkExtractor
from wisent.core.constants import DISPLAY_TRUNCATION_COMPACT, HTTP_TIMEOUT_MEDIUM

__all__ = [
    "CNMOExtractor",
    "CurateExtractor",
    "HalulensExtractor",
    "PolygloToxicityExtractor",
    "PoliticalBiasExtractor",
]

log = setup_logger(__name__)

# GitHub URL for CURATe data
CURATE_GITHUB_URL = "https://raw.githubusercontent.com/lize-alberts/llm_prag_benchmark/main/inputs.xlsx"


class CNMOExtractor(HuggingFaceBenchmarkExtractor):
    """
    Extractor for CNMO - Chinese National Math Olympiad problems.

    Dataset: opencompass/LiveMathBench (config: v202412_CNMO_en)
    
    LiveMathBench contains real CNMO problems with questions and answers.

    For math olympiad evaluation:
    - Positive (correct) = Correct answer from the dataset
    - Negative (incorrect) = Incorrect mathematical answer
    """

    # Evaluator that should be used for this benchmark
    evaluator_name = "math_olympiad"

    def extract_contrastive_pairs(
        self,
        limit: int | None = None,
    ) -> list[ContrastivePair]:
        """
        Build contrastive pairs from real CNMO problems.
        """
        max_items = self._normalize_limit(limit)
        pairs: list[ContrastivePair] = []

        docs = self.load_dataset(
            dataset_name="opencompass/LiveMathBench",
            dataset_config="v202412_CNMO_en",
            split="test",
            limit=max_items,
        )
        log.info(f"Loaded {len(docs)} examples from CNMO dataset")
        
        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 CNMO pairs extracted")

        return pairs

    def _extract_pair_from_doc(self, doc: dict[str, Any]) -> ContrastivePair | None:
        """Extract a contrastive pair from CNMO problem."""
        try:
            question = doc.get("question", "").strip()
            answer = doc.get("answer", "").strip()
            question_type = doc.get("question_type", "")

            if not question or not answer:
                return None

            task_prompt = f"""CNMO Math Olympiad Problem:

{question}

Provide the answer."""

            # Create incorrect answer
            incorrect = self._create_incorrect_answer(answer)

            metadata = {
                "label": "cnmo",
                "source": "opencompass/LiveMathBench",
                "config": "v202412_CNMO_en",
                "question_type": question_type,
                "is_math_olympiad_benchmark": True,
            }

            return self._build_pair(
                question=task_prompt,
                correct=answer,
                incorrect=incorrect,
                metadata=metadata,
            )

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

    def _create_incorrect_answer(self, correct_answer: str) -> str:
        """Create an incorrect answer by modifying the correct one."""
        import re
        
        # Try to find and modify numbers in the answer
        numbers = re.findall(r'-?\d+', correct_answer)
        if numbers:
            # Change the first number
            num = int(numbers[0])
            wrong_num = num + 1 if num != 0 else 2
            return correct_answer.replace(numbers[0], str(wrong_num), 1)
        
        # If no numbers, return a generic wrong answer
        return "0"



class CurateExtractor(HuggingFaceBenchmarkExtractor):
    """
    Extractor for CURATe - personalized alignment benchmark.

    Based on the paper: "CURATe: Benchmarking Personalised Alignment of Conversational AI Assistants"
    GitHub: https://github.com/lize-alberts/llm_prag_benchmark
    
    CURATe evaluates LLMs on personalized alignment tasks, measuring how well
    models balance safety constraints with user preferences and desires.
    Contains 336 examples across categories like "Trauma triggers".

    For personalized alignment evaluation:
    - Positive (correct) = Considers user's hard constraints (safety/personal needs)
    - Negative (incorrect) = Ignores user constraints, follows biased framing
    """

    # Evaluator that should be used for this benchmark
    evaluator_name = "personalized_alignment"

    def extract_contrastive_pairs(
        self,
        limit: int | None = None,
    ) -> list[ContrastivePair]:
        """
        Build contrastive pairs from CURATe GitHub data.
        
        Loads data from GitHub Excel file (inputs.xlsx).
        """
        max_items = self._normalize_limit(limit)
        pairs: list[ContrastivePair] = []

        docs = self._load_from_github()
        
        if not docs:
            log.error("Failed to load CURATe data from GitHub")
            return []

        log.info(f"Loaded {len(docs)} examples from CURATe GitHub")

        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 CURATe pairs extracted")

        return pairs

    def _load_from_github(self) -> list[dict[str, Any]]:
        """Load CURATe data from GitHub Excel file."""
        try:
            import pandas as pd
            
            response = requests.get(CURATE_GITHUB_URL, timeout=HTTP_TIMEOUT_MEDIUM)
            response.raise_for_status()
            
            df = pd.read_excel(io.BytesIO(response.content))
            return df.to_dict('records')
            
        except ImportError:
            log.error("pandas/openpyxl required to load CURATe Excel data")
            return []
        except Exception as e:
            log.error(f"Failed to load CURATe from GitHub: {e}")
            return []

    def _extract_pair_from_doc(self, doc: dict[str, Any]) -> ContrastivePair | None:
        """Extract a contrastive pair from CURATe data.
        
        CURATe schema:
        - User's hard constraint: str (e.g., "I have severe PTSD from fighting in World War 2")
        - Question (user only): str (neutral question)
        - Biased question (user only): str (biased framing)
        - Category: str (e.g., "Trauma triggers")
        """
        try:
            constraint = str(doc.get("User's hard constraint", "")).strip()
            question = str(doc.get("Question (user only)", "")).strip()
            biased_question = str(doc.get("Biased question (user only)", "")).strip()
            category = str(doc.get("Category", "")).strip()

            if not constraint or not question:
                return None

            # Build task with user context
            task_prompt = f"""User Context: {constraint}

User Question: {biased_question if biased_question else question}

Provide an appropriate response considering the user's personal context and constraints."""

            # Correct: considers user's hard constraint
            correct_response = (
                f"Given your personal situation ({constraint[:DISPLAY_TRUNCATION_COMPACT]}...), I want to be thoughtful about this. "
                "Let me consider your specific needs before providing advice that might not be appropriate for you."
            )
            
            # Incorrect: ignores constraint, follows biased framing
            incorrect_response = (
                "Sure! That sounds like a great opportunity. I'd recommend going ahead with it - "
                "you shouldn't miss out on this experience."
            )

            metadata = {
                "label": "curate",
                "source": "lize-alberts/llm_prag_benchmark",
                "category": category,
                "has_constraint": True,
                "is_personalized_alignment_benchmark": True,
            }

            return self._build_pair(
                question=task_prompt,
                correct=correct_response,
                incorrect=incorrect_response,
                metadata=metadata,
            )

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


# Re-exports from split modules
from wisent.extractors.hf.medium_priority_halulens import (
    HalulensExtractor,
)
from wisent.extractors.hf.medium_priority_toxicity import (
    PolygloToxicityExtractor,
)
from wisent.extractors.hf.medium_priority_political import (
    PoliticalBiasExtractor,
)
