# align with official CL-bench eval.py(https://github.com/Tencent-Hunyuan/CL-bench/blob/main/eval.py)
import json
from typing import Any, Dict, List

from evalscope.api.benchmark import BenchmarkMeta, DefaultDataAdapter
from evalscope.api.dataset import Sample
from evalscope.api.evaluator import TaskState
from evalscope.api.messages import dict_to_chat_message
from evalscope.api.messages.chat_message import ChatMessageUser
from evalscope.api.metric import Score
from evalscope.api.registry import register_benchmark
from evalscope.constants import Tags
from evalscope.utils.logger import get_logger

logger = get_logger()

CL_BENCH_DESCRIPTION = """
## Overview

CL-bench represents a step towards building LMs with this fundamental capability (Context Learning), making them more intelligent and advancing their deployment in real-world scenarios. This benchmark is specifically designed to evaluate a model's ability to learn specific, often novel or long-tail knowledge directly from the context and apply it to solve problems, simulating real-world learning processes.

**Resources:**
[Homepage](https://github.com/Tencent-Hunyuan/CL-bench) | [Dataset](https://huggingface.co/datasets/tencent/CL-bench)

## Task Description

- **Task Type**: In-Context Learning & Reasoning (Context-dependent QA)
- **Input**: A context containing new rules, fictional information, or specific logic, followed by a query
- **Output**: A solution derived strictly from the provided context (not pre-trained knowledge)
- **Difficulty**: Varied, requiring understanding of novel concepts defined in the prompt

## Key Features

- **Contamination-Free**: Uses synthetic or highly specific data (e.g., fictional laws, new programming syntax) to ensure models cannot rely on memorized training data
- **Real-world Simulation**: Mimics scenarios where humans learn new tasks by reading instructions or documentation
- **Diverse Domains**: Covers logic reasoning, rule following, language understanding, and puzzle solving within a given context
- **Evaluation of Adaptability**: Measures the "learning" capability rather than just "retrieval" or "memory"

## Evaluation Notes

- Focuses on the model's ability to follow strict instructions provided in the context
- Evaluation typically checks if the reasoning process utilizes the unique information given in the prompt
- Answers are often evaluated against specific ground-truth rules defined in the context
- Crucial for assessing how well models can adapt to private data or dynamic environments without fine-tuning
"""


@register_benchmark(
    BenchmarkMeta(
        name='cl_bench',
        pretty_name='CL-bench',
        tags=[Tags.INSTRUCTION_FOLLOWING, Tags.REASONING],
        description=CL_BENCH_DESCRIPTION,
        dataset_id='tencent-community/CL-bench',
        metric_list=['acc'],
        few_shot_num=0,
        train_split=None,
        eval_split='train',
        prompt_template='',
    )
)
class CLBenchAdapter(DefaultDataAdapter):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._use_llm_judge = True

    def load(self):
        last_error = None
        original_split = self.eval_split
        candidate_splits = [original_split]  # only eval on the original split, but or subsequent adaptation, if needed
        seen = set()
        for split in candidate_splits:
            if not split or split in seen:
                continue
            seen.add(split)
            try:
                self.eval_split = split
                return super().load()
            except Exception as exc:
                last_error = exc
                logger.warning(f'Failed to load CL-bench split "{split}": {exc}')
        if original_split:
            self.eval_split = original_split
        if last_error:
            raise last_error
        return super().load()

    def record_to_sample(self, record: Dict[str, Any]) -> Sample:
        messages = [dict_to_chat_message(msg) for msg in record.get('messages', [])]
        rubrics = record.get('rubrics', [])
        metadata = record.get('metadata') or record.get('meta_data') or {}
        return Sample(input=messages, target=rubrics, metadata=metadata)

    def llm_match_score(
        self,
        original_prediction: str,
        filtered_prediction: str,
        reference: List[str],
        task_state: TaskState,
    ) -> Score:
        from .utils import build_rubrics_text, extract_json_block

        score = Score(
            extracted_prediction=filtered_prediction,
            prediction=original_prediction,
        )

        if not filtered_prediction or not filtered_prediction.strip():
            score.value = {'acc': 0.0}
            score.explanation = 'Empty model output; scored as 0.'
            score.main_score_name = 'acc'
            return score

        rubrics_text = build_rubrics_text(reference)
        grading_prompt = (
            "Starting now, you are a rigorous instruction-following grading teacher. Your task is to accurately grade and score student answers based on the 【Rubrics】.\n\n"
            "Grading Criteria\n"
            "This is a strict, all-or-nothing grading system. The final score is binary.\n"
            "To receive a score of 1, the student's answer must perfectly satisfy every single requirement listed in the 【Rubrics】.\n"
            "If even one requirement is not fully met, the final score will be 0.\n"
            "Grading Process\n"
            "Please strictly follow the steps below for analysis—no steps may be skipped:\n"
            "Step 1: Analyze the Standard Answer\n"
            "List all explicit requirements in the 【Rubrics】 item by item (including format, content, quantity, order, etc.).\n"
            "Identify implicit requirements in the 【Rubrics】 (e.g., language style, logical structure).\n"
            "Define specific evaluation criteria for each requirement (e.g., \"must include X,\" \"must not exceed Y\").\n"
            "Step 2: Check Each Requirement Against the Student's Answer\n"
            "For every requirement in the 【Rubrics】, verify one by one whether the student's answer fully satisfies it.\n"
            "Step 3: Self-Reflection\n"
            "Before giving the final score, you must conduct the following checks:\n"
            "  Completeness Check: Whether all requirements in the standard answer have been reviewed with no omissions.\n"
            "  Strictness Check: Whether the evaluation strictly adheres to the \"fully satisfied\" standard without relaxing requirements due to subjective judgment.\n"
            "  Consistency Check: Whether the grading rationale aligns logically with the final score.\n"
            "  Objectivity Check: Whether judgments are based on objective facts rather than subjective speculation.\n"
            "Output Format Requirements\n"
            "【Grading Rationale】: xxx\n"
            "【List of Requirement Satisfaction Status】: [x₁, x₂, …, xᵢ, …, xₙ] (where n is the total number of requirements in the 【Rubrics】, and xᵢ indicates whether the student's answer meets the i-th requirement, with values \"yes\"/\"no\")\n"
            "【Overall Score】: x points (x is an integer, either 0 or 1.)\n\n"
            "Content to Be Graded\n"
            f"【Rubrics】:\n{rubrics_text}\n"
            f"【Student Response】:\n{filtered_prediction}\n"
            "\nPlease strictly output ONLY the following JSON format (do not output any other content):\n"
            "{\n"
            '  "Grading Rationale": "Your detailed grading rationale",\n'
            '  "List of Requirement Satisfaction Status": ["yes", "no", ...],\n'
            '  "Overall Score": 0 or 1\n'
            "}\n"
        )
        messages = [ChatMessageUser(content=grading_prompt)]

        judge_response = self.llm_judge.judge(messages=messages)

        parsed = None
        json_block = extract_json_block(judge_response)
        if json_block:
            try:
                parsed = json.loads(json_block)
            except json.JSONDecodeError as exc:
                logger.warning(f'Failed to parse CL-bench judge response: {exc}')

        if not isinstance(parsed, dict):
            score.value = {'acc': 0.0}
            score.explanation = f'LLM judge parsing failed: {judge_response}'
            score.metadata = {
                'source': 'llm_judge',
                'judge_strategy': self.judge_strategy,
                'model': self.llm_judge.model_id,
                'judge_response': judge_response,
            }
            score.main_score_name = 'acc'
            return score

        overall_score = parsed.get('Overall Score', 0)
        try:
            overall_score = int(overall_score)
        except (TypeError, ValueError):
            overall_score = 0
        overall_score = 1 if overall_score == 1 else 0

        score.value = {'acc': float(overall_score)}
        score.explanation = parsed.get('Grading Rationale', '')
        score.metadata = {
            'source': 'llm_judge',
            'judge_strategy': self.judge_strategy,
            'model': self.llm_judge.model_id,
            'requirement_status': parsed.get('List of Requirement Satisfaction Status', []),
            'judge_response': parsed,
        }
        score.main_score_name = 'acc'
        return score
