"""IFEval benchmark loader for instruction following evaluation.

IFEval (Instruction Following Evaluation) tests whether models can follow
specific verifiable constraints in prompts like "respond in exactly 3 sentences"
or "include the word 'absolutely' at least twice."
"""

import re
from typing import Any, Iterator

from ember.api.data import HuggingFaceSource
from embergraphrouter.evals.benchmarks.base import (
    Benchmark,
    BenchmarkDomain,
    BenchmarkExample,
    EvaluationScore,
    MetricType,
)
from embergraphrouter.evals.benchmarks.registry import register_benchmark


def _count_sentences(text: str) -> int:
    """Count sentences using simple heuristics."""
    # Split on sentence-ending punctuation followed by space/end
    sentences = re.split(r"[.!?]+(?:\s+|$)", text.strip())
    return len([s for s in sentences if s.strip()])


def _count_words(text: str) -> int:
    """Count words in text."""
    return len(text.split())


def _count_paragraphs(text: str) -> int:
    """Count paragraphs (double-newline separated)."""
    paragraphs = text.strip().split("\n\n")
    return len([p for p in paragraphs if p.strip()])


def _check_word_presence(text: str, word: str, min_count: int = 1) -> bool:
    """Check if word appears at least min_count times."""
    pattern = rf"\b{re.escape(word)}\b"
    matches = re.findall(pattern, text, re.IGNORECASE)
    return len(matches) >= min_count


def _check_no_word(text: str, word: str) -> bool:
    """Check if word does NOT appear."""
    pattern = rf"\b{re.escape(word)}\b"
    return not re.search(pattern, text, re.IGNORECASE)


def _check_starts_with(text: str, prefix: str) -> bool:
    """Check if text starts with prefix."""
    return text.strip().lower().startswith(prefix.lower())


def _check_ends_with(text: str, suffix: str) -> bool:
    """Check if text ends with suffix."""
    return text.strip().lower().endswith(suffix.lower())


def _check_all_caps(text: str) -> bool:
    """Check if all letters are uppercase."""
    letters = [c for c in text if c.isalpha()]
    return all(c.isupper() for c in letters) if letters else False


def _check_no_caps(text: str) -> bool:
    """Check if no letters are uppercase."""
    letters = [c for c in text if c.isalpha()]
    return all(c.islower() for c in letters) if letters else False


def _evaluate_constraint(constraint: dict[str, Any], response: str) -> bool:
    """Evaluate a single constraint against response.

    Args:
        constraint: Constraint specification with type and parameters.
        response: Model's output.

    Returns:
        True if constraint is satisfied.
    """
    constraint_type = constraint.get("type", "")
    params = constraint.get("params", {})

    if constraint_type == "num_sentences":
        expected = params.get("count")
        relation = params.get("relation", "eq")
        actual = _count_sentences(response)
        if relation == "eq":
            return actual == expected
        elif relation == "ge":
            return actual >= expected
        elif relation == "le":
            return actual <= expected

    elif constraint_type == "num_words":
        expected = params.get("count")
        relation = params.get("relation", "eq")
        actual = _count_words(response)
        if relation == "eq":
            return actual == expected
        elif relation == "ge":
            return actual >= expected
        elif relation == "le":
            return actual <= expected

    elif constraint_type == "num_paragraphs":
        expected = params.get("count")
        relation = params.get("relation", "eq")
        actual = _count_paragraphs(response)
        if relation == "eq":
            return actual == expected
        elif relation == "ge":
            return actual >= expected
        elif relation == "le":
            return actual <= expected

    elif constraint_type == "include_word":
        word = params.get("word", "")
        min_count = params.get("min_count", 1)
        return _check_word_presence(response, word, min_count)

    elif constraint_type == "exclude_word":
        word = params.get("word", "")
        return _check_no_word(response, word)

    elif constraint_type == "starts_with":
        prefix = params.get("text", "")
        return _check_starts_with(response, prefix)

    elif constraint_type == "ends_with":
        suffix = params.get("text", "")
        return _check_ends_with(response, suffix)

    elif constraint_type == "all_caps":
        return _check_all_caps(response)

    elif constraint_type == "no_caps":
        return _check_no_caps(response)

    elif constraint_type == "format_json":
        try:
            import json

            json.loads(response.strip())
            return True
        except (json.JSONDecodeError, ValueError):
            return False

    elif constraint_type == "format_bullets":
        lines = [ln.strip() for ln in response.strip().split("\n") if ln.strip()]
        return all(ln.startswith(("- ", "* ", "• ")) for ln in lines if ln)

    # Unknown constraint type: fail safe
    return False


def _parse_ifeval_constraints(row: dict[str, Any]) -> list[dict[str, Any]]:
    """Parse IFEval constraint format from dataset row.

    IFEval uses instruction_id_list and kwargs to specify constraints.
    """
    constraints = []

    instruction_ids = row.get("instruction_id_list", [])
    kwargs_list = row.get("kwargs", [])

    for i, inst_id in enumerate(instruction_ids):
        kwargs = kwargs_list[i] if i < len(kwargs_list) else {}

        # Map IFEval instruction IDs to our constraint format
        constraint: dict[str, Any] = {"instruction_id": inst_id, "params": kwargs}

        if "num_sentences" in inst_id:
            constraint["type"] = "num_sentences"
            constraint["params"] = {
                "count": kwargs.get("num_sentences", 0),
                "relation": kwargs.get("relation", "eq"),
            }

        elif "num_words" in inst_id:
            constraint["type"] = "num_words"
            constraint["params"] = {
                "count": kwargs.get("num_words", 0),
                "relation": kwargs.get("relation", "eq"),
            }

        elif "num_paragraphs" in inst_id:
            constraint["type"] = "num_paragraphs"
            constraint["params"] = {
                "count": kwargs.get("num_paragraphs", 0),
                "relation": kwargs.get("relation", "eq"),
            }

        elif "include_word" in inst_id or "keywords" in inst_id:
            constraint["type"] = "include_word"
            word = kwargs.get("keyword") or kwargs.get("keywords", [""])[0]
            constraint["params"] = {"word": word, "min_count": 1}

        elif "exclude_word" in inst_id or "forbidden_words" in inst_id:
            constraint["type"] = "exclude_word"
            word = kwargs.get("forbidden_word") or kwargs.get("forbidden_words", [""])[0]
            constraint["params"] = {"word": word}

        elif "start" in inst_id:
            constraint["type"] = "starts_with"
            constraint["params"] = {"text": kwargs.get("starter", "")}

        elif "end" in inst_id:
            constraint["type"] = "ends_with"
            constraint["params"] = {"text": kwargs.get("ender", "")}

        elif "all_uppercase" in inst_id or "all_caps" in inst_id:
            constraint["type"] = "all_caps"

        elif "lowercase" in inst_id:
            constraint["type"] = "no_caps"

        elif "json" in inst_id.lower():
            constraint["type"] = "format_json"

        elif "bullet" in inst_id.lower():
            constraint["type"] = "format_bullets"

        else:
            # Generic constraint - keep raw instruction ID
            constraint["type"] = "unknown"

        constraints.append(constraint)

    return constraints


@register_benchmark("ifeval")
class IFEvalBenchmark(Benchmark):
    """IFEval Instruction Following benchmark.

    Tests whether models can follow verifiable constraints in prompts.
    Each example has one or more constraints that can be checked programmatically.

    This benchmark is NOT in EB1 router training data.
    """

    name = "ifeval"
    version = "1.0"
    domain = BenchmarkDomain.INSTRUCTION
    metric_type = MetricType.CONSTRAINT
    description = "Instruction following with verifiable constraints"

    def __init__(self, max_examples: int | None = None):
        """Initialize IFEval benchmark.

        Args:
            max_examples: Maximum examples to load (None = all ~541).
        """
        self._max_examples = max_examples

    def load(self, subset: str | None = None) -> Iterator[BenchmarkExample]:
        """Load IFEval examples via HuggingFace.

        Args:
            subset: Unused for IFEval (single split).

        Yields:
            BenchmarkExample instances with constraint specifications.
        """
        # IFEval on HuggingFace: google/IFEval (or wis-k/instruction-following-eval)
        source = HuggingFaceSource(
            "google/IFEval",
            split="train",  # IFEval typically has one split
        )

        count = 0
        for batch in source.read_batches():
            for row in batch:
                if self._max_examples and count >= self._max_examples:
                    return

                prompt = str(row.get("prompt", ""))
                constraints = _parse_ifeval_constraints(row)

                # IFEval doesn't have "reference answers" - just constraints
                # The reference is the set of constraint specifications
                constraint_descriptions = [
                    c.get("instruction_id", c.get("type", "unknown"))
                    for c in constraints
                ]

                yield BenchmarkExample(
                    id=f"ifeval_{count}",
                    query=prompt,
                    reference="[constraint-based evaluation]",
                    constraints=constraint_descriptions,
                    metadata={
                        "domain": "instruction",
                        "source": "ifeval",
                        "n_constraints": len(constraints),
                        "constraint_specs": constraints,
                        "raw_instruction_ids": row.get("instruction_id_list", []),
                    },
                )
                count += 1

    def evaluate(
        self,
        prediction: str,
        example: BenchmarkExample,
    ) -> EvaluationScore:
        """Evaluate prediction against all constraints.

        Args:
            prediction: Model's output.
            example: Ground truth with constraint specifications.

        Returns:
            EvaluationScore with constraint satisfaction metrics.
        """
        constraint_specs = example.metadata.get("constraint_specs", [])
        if not constraint_specs:
            return EvaluationScore(correct=True, score=1.0)

        results = []
        for constraint in constraint_specs:
            satisfied = _evaluate_constraint(constraint, prediction)
            results.append({
                "constraint_id": constraint.get("instruction_id", constraint.get("type")),
                "type": constraint.get("type"),
                "satisfied": satisfied,
            })

        n_satisfied = sum(1 for r in results if r["satisfied"])
        n_total = len(results)
        all_satisfied = n_satisfied == n_total

        return EvaluationScore(
            correct=all_satisfied,
            score=n_satisfied / n_total if n_total > 0 else 0.0,
            details={
                "metric": "constraint_satisfaction",
                "n_satisfied": n_satisfied,
                "n_total": n_total,
                "constraint_results": results,
            },
        )
