# Copyright (c) Alibaba, Inc. and its affiliates.
import json
from typing import Any, Dict, List, Optional, Tuple, Union

from evalscope.api.benchmark import BenchmarkMeta, FunctionCallAdapter
from evalscope.api.dataset import Sample
from evalscope.api.evaluator import TaskState
from evalscope.api.messages import dict_to_chat_message
from evalscope.api.metric import AggScore, SampleScore, Score
from evalscope.api.registry import register_benchmark
from evalscope.api.tool import ToolInfo
from evalscope.constants import Tags


@register_benchmark(
    BenchmarkMeta(
        name='general_fc',
        pretty_name='General-FunctionCalling',
        description="""
## Overview

General-FunctionCalling is a customizable benchmark for evaluating function calling (tool use) capabilities of language models. It tests both the decision to call tools and the accuracy of generated function calls.

## Task Description

- **Task Type**: Function Calling / Tool Use Evaluation
- **Input**: Messages with available tool definitions
- **Output**: Function calls with valid arguments
- **Flexibility**: Supports custom datasets via local files

## Key Features

- Evaluates function calling decision accuracy
- Schema validation for function arguments
- Comprehensive metrics (F1, precision, recall)
- Custom dataset support via local file loading
- JSON schema-based argument validation

## Evaluation Notes

- Default configuration uses **0-shot** evaluation
- Metrics: **tool_call_f1**, **schema_accuracy**, count metrics
- Aggregation: **F1** score
- Validates function calls against JSON schema
- See [User Guide](https://evalscope.readthedocs.io/en/latest/advanced_guides/custom_dataset/llm.html#fc) for dataset format
""",  # noqa: E501
        tags=[Tags.FUNCTION_CALLING, Tags.CUSTOM, Tags.AGENT],
        dataset_id='evalscope/GeneralFunctionCall-Test',
        metric_list=[
            'count_finish_reason_tool_call',
            'count_successful_tool_call',
            'schema_accuracy',
            'tool_call_f1',
        ],
        aggregation='f1',
        eval_split='test',
    )
)
class GeneralFCAdapter(FunctionCallAdapter):

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

        self.add_aggregation_name = False

    def load_from_disk(self, **kwargs):
        return super().load_from_disk(use_local_loader=True)

    def record_to_sample(self, record: Dict[str, Any]) -> Sample:
        messages = record.get('messages', [])
        tools = record.get('tools', [])
        should_call_tool = record.get('should_call_tool', False)

        # In case the fields are stored as JSON strings
        if isinstance(messages, str):
            messages = json.loads(messages)
        if isinstance(tools, str):
            tools = json.loads(tools)

        # Convert to Sample
        return Sample(
            input=[dict_to_chat_message(msg) for msg in messages],
            target='',
            tools=[ToolInfo.model_validate(tool['function']) for tool in tools],
            metadata={
                'should_call_tool': should_call_tool,
                'tools': tools,
            }
        )

    def match_score(self, original_prediction, filtered_prediction, reference, task_state: TaskState) -> Score:
        score = Score(
            prediction=original_prediction,
            extracted_prediction=filtered_prediction,
        )

        model_output = task_state.output
        should_call_tool = task_state.metadata.get('should_call_tool', False)
        if model_output.error:
            score.value = {
                'passed': 0,
                'finish_reason_tool_call': 0,
                'successful_tool_call': 0,
                'should_call_tool': int(should_call_tool),
            }
            score.metadata = {
                'error_reason': f'Model inference error: {model_output.error}',
            }

        else:
            finish_reason = model_output.stop_reason
            tool_calls = model_output.message.tool_calls or []
            tools = task_state.metadata['tools']

            is_call_tool = finish_reason == 'tool_calls'
            # When model calls tool, use the tool calls from model output
            if is_call_tool:
                score.prediction = task_state.messages_markdown
                score.extracted_prediction = task_state.messages_markdown
            is_valid_tool_call, error_reason = self.validate_tool_call(tool_calls, tools)
            is_call_successful = is_call_tool and is_valid_tool_call
            if should_call_tool:
                passed = is_call_successful
            else:
                passed = not is_call_tool
            score.value = {
                'passed': passed,
                'finish_reason_tool_call': int(is_call_tool),
                'successful_tool_call': int(is_call_successful),
                'should_call_tool': int(should_call_tool),
            }
            score.metadata = {
                'error_reason': error_reason,
            }

        return score

    def aggregate_scores(self, sample_scores: List[SampleScore]) -> List[AggScore]:
        """
        Aggregate General-FunctionCalling metrics:
        - count_finish_reason_tool_call: number of samples with finish_reason == 'tool_calls'.
        - count_successful_tool_call: number of schema-valid tool calls among attempts.
        - schema_accuracy: valid tool calls / attempted tool calls (0.0 if no attempts).
        - tool_call_f1: F1 over whether a tool should be called vs. model predicted tool call.
        """
        finish_reason_tool_call_count = 0
        successful_tool_call_count = 0
        attempted_tool_calls = 0
        valid_tool_calls = 0

        tp = fp = fn = tn = 0
        total_count = len(sample_scores)

        for ss in sample_scores:
            # Values set in match_score
            v = ss.score.value or {}
            pred_is_tool = int(v.get('finish_reason_tool_call', 0))  # model predicted to call tool
            pred_is_success = int(v.get('successful_tool_call', 0))  # model called tool and schema validated
            label_should_tool = int(v.get('should_call_tool', 0))  # ground-truth label

            # Counts
            finish_reason_tool_call_count += pred_is_tool
            successful_tool_call_count += pred_is_success

            if pred_is_tool == 1:
                attempted_tool_calls += 1
                if pred_is_success == 1:
                    valid_tool_calls += 1

            # Confusion matrix for F1
            if pred_is_tool == 1 and label_should_tool == 1:
                tp += 1
            elif pred_is_tool == 1 and label_should_tool == 0:
                fp += 1
            elif pred_is_tool == 0 and label_should_tool == 1:
                fn += 1
            else:
                tn += 1

        # Schema accuracy among attempts
        schema_accuracy = (valid_tool_calls / attempted_tool_calls) if attempted_tool_calls > 0 else 0.0

        # Precision/Recall/F1 for tool-calling decision
        precision = (tp / (tp + fp)) if (tp + fp) > 0 else 0.0
        recall = (tp / (tp + fn)) if (tp + fn) > 0 else 0.0
        tool_call_f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) > 0 else 0.0

        metrics = {
            'tool_call_f1': tool_call_f1,
            'count_finish_reason_tool_call': finish_reason_tool_call_count,
            'count_successful_tool_call': successful_tool_call_count,
            'schema_accuracy': schema_accuracy,
        }

        agg_scores: List[AggScore] = []
        for name, value in metrics.items():
            agg_scores.append(AggScore(metric_name=name, score=value, num=total_count, metadata={}))

        return agg_scores
