# !/usr/bin/env python3
# -*- coding:utf-8 -*-

# @Time    : 2024/7/1 15:50
# @Author  : wangchongshi
# @Email   : wangchongshi.wcs@antgroup.com
# @FileName: data_agent.py
import datetime

from agentuniverse.agent.agent import Agent
from agentuniverse.agent.agent_manager import AgentManager
from agentuniverse.agent.input_object import InputObject
from agentuniverse.agent.output_object import OutputObject
from agentuniverse.base.util.logging.logging_util import LOGGER


class DataAgent(Agent):
    """Data Agent class."""

    def input_keys(self) -> list[str]:
        """Return the input keys of the Agent."""
        return ['queryset_path']

    def output_keys(self) -> list[str]:
        """Return the output keys of the Agent."""
        return []

    def parse_input(self, input_object: InputObject, agent_input: dict) -> dict:
        """Agent parameter parsing.

        Args:
            input_object (InputObject): input parameters passed by the user.
            agent_input (dict): agent input preparsed by the agent.
        Returns:
            dict: agent input parsed from `input_object` by the user.
        """
        return {}

    def parse_result(self, planner_result: dict) -> dict:
        """Planner result parser.

        Args:
            planner_result(dict): Planner result
        Returns:
            dict: Agent result object.
        """
        return planner_result

    def execute(self, input_object: InputObject, agent_input: dict):
        """Execute agent instance.

        Args:
            input_object (InputObject): input parameters passed by the user.
            agent_input (dict): agent input parsed from `input_object` by the user.
        """
        date = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
        input_object.add_data('date', date)

        # step1: build q&a dataset from the candidate agent which needs to be evaluated.
        query_answer_list = self.dataset_build_agent(input_object)
        input_object.add_data('query_answer_list', query_answer_list)

        # step2: evaluate q&a dataset generated by the candidate agent and generate evaluation report.
        eval_report_list = self.dataset_eval_agent(input_object)
        return {'query_answer_list': query_answer_list, 'eval_report_list': eval_report_list}

    def dataset_build_agent(self, input_object: InputObject):
        """Build q&a dataset from the candidate agent which needs to be evaluated."""

        profile_config = self.agent_model.profile
        dataset_build_agent: Agent = AgentManager().get_instance_obj(profile_config.get('dataset_builder'))
        if dataset_build_agent is None:
            raise ValueError('The agent instance corresponding to `dataset_builder` parameter is empty')
        output: OutputObject = dataset_build_agent.run(**input_object.to_dict())
        return output.get_data('query_answer_list', [])

    def dataset_eval_agent(self, input_object: InputObject):
        """Evaluate q&a dataset generated by the candidate agent and generate evaluation report."""

        profile_config = self.agent_model.profile
        dataset_eval_agent: Agent = AgentManager().get_instance_obj(profile_config.get('dataset_evaluator'))
        if dataset_eval_agent is None:
            LOGGER.warn("The agent instance corresponding to `dataset_evaluator` parameter is empty")
            return []
        output: OutputObject = dataset_eval_agent.run(**input_object.to_dict())
        return output.get_data('eval_report_json_list', [])
