# Copyright (c) Alibaba, Inc. and its affiliates.

import os
import re
from typing import Any, Dict

from evalscope.api.benchmark import BenchmarkMeta, DefaultDataAdapter
from evalscope.api.dataset import Sample
from evalscope.api.evaluator import TaskState
from evalscope.api.registry import register_benchmark
from evalscope.constants import Tags
from evalscope.utils.logger import get_logger

logger = get_logger()

# BBH multiple choice subset list
MULTIPLE_CHOICE = 'multiple_choice'
MULTIPLE_CHOICE_LIST = [
    'temporal_sequences',
    'disambiguation_qa',
    'date_understanding',
    'tracking_shuffled_objects_three_objects',
    'penguins_in_a_table',
    'geometric_shapes',
    'snarks',
    'ruin_names',
    'tracking_shuffled_objects_seven_objects',
    'tracking_shuffled_objects_five_objects',
    'logical_deduction_three_objects',
    'hyperbaton',
    'logical_deduction_five_objects',
    'logical_deduction_seven_objects',
    'movie_recommendation',
    'salient_translation_error_detection',
    'reasoning_about_colored_objects',
]

# The free form subset list of BBH dataset
FREE_FORM = 'free_form'
FREE_FORM_LIST = [
    'multistep_arithmetic_two',
    'navigate',
    'dyck_languages',
    'word_sorting',
    'sports_understanding',
    'boolean_expressions',
    'object_counting',
    'formal_fallacies',
    'causal_judgement',
    'web_of_lies',
]

# BBH sub-task type
TASK_TYPE = 'task_type'
SUBSET_LIST = MULTIPLE_CHOICE_LIST + FREE_FORM_LIST

PROMPT_TEMPLATE = """
Q: {question}
A: Let's think step by step. Put your final answer in the format of "So the answer is [ANSWER]" (without quotes and markdown) where [ANSWER] is the answer to the problem.
""".lstrip()  # noqa: E501

FEWSHOT_TEMPLATE = """
{fewshot}

""".lstrip() + PROMPT_TEMPLATE


@register_benchmark(
    BenchmarkMeta(
        name='bbh',
        pretty_name='BBH',
        dataset_id='evalscope/bbh',
        tags=[Tags.REASONING],
        description="""
## Overview

BBH (BIG-Bench Hard) is a subset of 23 challenging tasks from the BIG-Bench benchmark that are specifically selected because language models initially struggled with them. These tasks require complex reasoning abilities that benefit from Chain-of-Thought (CoT) prompting.

## Task Description

- **Task Type**: Mixed (Multiple-Choice and Free-Form)
- **Input**: Task-specific questions requiring reasoning
- **Output**: Answers in specified format
- **Subsets**: 27 reasoning tasks divided into multiple-choice (17) and free-form (10)

## Key Features

- 27 challenging reasoning tasks from BIG-Bench
- Multiple-choice tasks: temporal sequences, disambiguation, logical deduction, etc.
- Free-form tasks: arithmetic, navigation, boolean expressions, etc.
- Each task comes with curated Chain-of-Thought examples
- Designed to test advanced reasoning capabilities

## Evaluation Notes

- Default configuration uses **3-shot** with CoT prompting (recommended)
- CoT prompts are pre-defined for each subset in `cot_prompts/` directory
- Answers should follow the format: "So the answer is [ANSWER]"
- Setting `few_shot_num=0` disables few-shot examples
- Multiple-choice answers are normalized to single letters (A, B, C, etc.)
""",
        subset_list=SUBSET_LIST,
        few_shot_num=3,
        train_split=None,
        eval_split='test',
        metric_list=['acc'],
        prompt_template=PROMPT_TEMPLATE,
        few_shot_prompt_template=FEWSHOT_TEMPLATE,
    )
)
class BBHAdapter(DefaultDataAdapter):
    """
    Adapter for BBH free-form and multiple-choices sub-tasks.
    """

    def __init__(self, **kwargs):
        few_shot_num = kwargs.get('few_shot_num', 3)

        if few_shot_num != 3 and few_shot_num != 0:
            logger.error(
                f'BBH uses 3-shot examples with CoT or 0-shot by system, but got {few_shot_num}. '
                f'Use 3-shot by default.'
            )
            kwargs['few_shot_num'] = 3

        super().__init__(**kwargs)

    def record_to_sample(self, record: Dict[str, Any]) -> Sample:
        input = record['input']
        target = record['target'].replace('(', '').replace(')', '').strip()  # Clean up the target answer

        # Determine task type based on subset name
        task_type = None
        subset_name = self.current_subset_name
        if subset_name in MULTIPLE_CHOICE_LIST:
            task_type = MULTIPLE_CHOICE
        elif subset_name in FREE_FORM_LIST:
            task_type = FREE_FORM

        metadata = {TASK_TYPE: task_type}

        return Sample(input=input, target=target, metadata=metadata, subset_key=subset_name)

    def format_fewshot_template(self, fewshot: str, sample: Sample) -> str:
        # Load CoT prompts from file for BBH
        subset_name = sample.subset_key
        if subset_name:
            cot_file_path = os.path.join(os.path.dirname(__file__), 'cot_prompts', f'{subset_name}.txt')
            if os.path.exists(cot_file_path):
                with open(cot_file_path, 'r', encoding='utf-8') as f:
                    fewshot = f.read().strip()
        return self.few_shot_prompt_template.format(
            fewshot=fewshot,
            question=sample.input,
        )

    def extract_answer(self, prediction: str, task_state: TaskState):
        task_type = task_state.metadata.get(TASK_TYPE)

        if task_type == MULTIPLE_CHOICE:
            return self._extract_mc_answer(prediction)
        elif task_type == FREE_FORM:
            return self._extract_ff_answer(prediction)
        else:
            return prediction.strip()

    @classmethod
    def _extract_mc_answer(cls, ans: str) -> str:
        """
        Extract normalized answer for BBH multiple-choice tasks.
        Handles formats like:
        - "answer is (A)"
        - "The answer is A."
        - Extra text after answer.
        Always uses the *last* occurrence of "answer is".
        """
        ans = ans.strip()

        parts = ans.split('So the answer is ')
        if len(parts) > 1:
            ans = parts[-1].strip()
        ans = ans.split('\n')[0].strip()

        # Remove trailing period
        if ans.endswith('.'):
            ans = ans[:-1].strip()

        # Capture uppercase letter inside parentheses (A) (B) ...
        match = re.search(r'\(([A-Z])\)', ans)
        if match:
            return match.group(1)

        # Capture single uppercase letter
        match = re.search(r'\b([A-Z])\b', ans)
        if match:
            return match.group(1)

        return ans

    @classmethod
    def _extract_ff_answer(cls, ans: str):
        """
        Extract the normalized answer for BBH free-form tasks.
        Handles patterns like:
        - "answer is XXX."
        - "The answer is **valid**."
        - Extra trailing dots / line breaks.
        - Bold-marked answers (**xxx**).
        Always uses the *last* occurrence of "answer is".
        """
        ans = ans.strip()

        parts = ans.split('So the answer is ')
        if len(parts) > 1:
            ans = parts[-1].strip()
        ans = ans.split('\n')[0].strip()

        # Remove trailing period
        if ans.endswith('.'):
            ans = ans[:-1].strip()

        # If answer is in bold (**xxx**), prefer the content inside
        match = re.search(r'\*\*(.*?)\*\*', ans)
        if match:
            ans = match.group(1).strip()

        return ans
