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

from typing import Any, Dict, List, Tuple

from evalscope.api.benchmark import BenchmarkMeta, DefaultDataAdapter
from evalscope.api.dataset import Sample
from evalscope.api.evaluator import TaskState
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
from .utils import build_full_code, normalize_languages

logger = get_logger()


@register_benchmark(
    BenchmarkMeta(
        name='multiple_mbpp',
        pretty_name='MultiPL-E MBPP',
        tags=[Tags.CODING],
        description="""
## Overview

MultiPL-E MBPP is a multilingual code generation benchmark derived from MBPP (Mostly Basic Python Programming). It extends the original MBPP to 18 programming languages, enabling cross-lingual evaluation of code generation capabilities.

## Task Description

- **Task Type**: Multilingual Code Generation
- **Input**: Programming problem prompt with docstring
- **Output**: Complete code solution that passes test cases
- **Languages**: 18 languages (C++, TypeScript, Shell, C#, Go, Java, Lua, JavaScript, PHP, Perl, Racket, R, Rust, Scala, Swift, Ruby, D, Julia)

## Key Features

- Multilingual evaluation across 18 programming languages
- Execution-based evaluation with test cases
- Supports pass@k metric for code generation
- Docker sandbox environment for safe code execution
- Derived from MBPP with consistent problem difficulty

## Evaluation Notes

- **Sandbox Required**: Requires sandbox environment for safe code execution
- Default evaluation uses **test** split
- Primary metric: **Accuracy** with **pass@k** aggregation
- Timeout: 30 seconds per test case
- See [sandbox documentation](https://evalscope.readthedocs.io/en/latest/user_guides/sandbox.html) for setup
""",  # noqa: E501
        dataset_id='evalscope/MultiPL-E',
        subset_list=[
            'mbpp-cpp',
            'mbpp-ts',
            'mbpp-sh',
            'mbpp-cs',
            'mbpp-go',
            'mbpp-java',
            'mbpp-lua',
            'mbpp-js',
            'mbpp-php',
            'mbpp-pl',
            'mbpp-rkt',
            'mbpp-r',
            'mbpp-rs',
            'mbpp-scala',
            'mbpp-swift',
            'mbpp-rb',
            'mbpp-d',
            'mbpp-jl',
        ],
        metric_list=['acc'],
        aggregation='mean_and_pass_at_k',
        eval_split='test',
        prompt_template='{prompt}',
        review_timeout=30,
        sandbox_config={
            'image': 'volcengine/sandbox-fusion:server-20250609',
            'tools_config': {
                'shell_executor': {},
                'python_executor': {},
                'multi_code_executor': {}  # Multi-language code executor
            },
            'memory_limit': '2g',
            'cpu_limit': '2.0',
        },
    )
)
class MultiPLEMBPPAdapter(DefaultDataAdapter):
    """
    MultiPL-E MBPP adapter using the new data processing framework.
    Assumptions:
    - Each subset is a single language suite.
    - Records contain: 'prompt', 'tests', optional 'stop_tokens', 'language', and id: 'task_id' or 'name'.
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def record_to_sample(self, record: Dict[str, Any]) -> Sample:
        """Convert a data record to a Sample object."""
        return Sample(
            input=record['prompt'],
            target='',
            metadata={
                'tests': record['tests'],
                'stop_tokens': record.get('stop_tokens', []),
                'task_id': record.get('name', record.get('task_id')),
                'language': record.get('language'),
                'doctests': record.get('doctests', ''),
            }
        )

    def format_prompt_template(self, sample: Sample) -> str:
        """
        Freeform SFT prompt:
        - Fence the given prompt with the language derived from metadata.language ("mbpp-<lang>").
        - Add a short instruction requesting full code without a Main entrypoint.
        """
        extract_lang, _ = normalize_languages(sample.metadata.get('language'))
        instruction = (
            'Please complete the above code according to the requirements in the docstring. '
            'Write the complete code and wrap it in markdown fenced code. The code should not contain `Main` function.'
        )
        return f'```{extract_lang}\n{sample.input}\n```\n\n{instruction}'

    def extract_answer(self, prediction: str, task_state: TaskState) -> str:
        """SFT extraction: extract fenced code, stop at language end tokens, remove entrypoints, append tests."""
        return build_full_code(
            prediction=prediction,
            language=task_state.metadata.get('language'),
            stop_tokens=task_state.metadata.get('stop_tokens', []),
            tests=task_state.metadata.get('tests', ''),
        )

    def match_score(
        self, original_prediction: str, filtered_prediction: str, reference: str, task_state: TaskState
    ) -> Score:
        """Run code in sandbox and return pass/fail."""
        if not self.use_sandbox:
            raise RuntimeError(
                'MultiPL-E MBPP requires sandboxed code execution for safety. Enable use_sandbox in TaskConfig.'
            )

        score = Score(extracted_prediction=filtered_prediction, prediction=original_prediction)
        _, run_language = normalize_languages(task_state.metadata.get('language'))

        res = self.execute_code_in_sandbox(
            code=filtered_prediction,
            timeout=self.review_timeout,
            language=run_language,
        )
        passed = res.get('status') == 'success'
        score.value = {'acc': passed}
        score.metadata = {
            'task_id': task_state.metadata.get('task_id'),
            'timeout': self.review_timeout,
            'execution_result': res,
            'run_language': run_language,
        }
        score.main_score_name = 'acc'
        return score
