import os
import sys
import json
import logging
from typing import Dict, List, Any

# 设置日志格式
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)


from meta_researcher.tool.base import BaseTool
from meta_researcher.tool.tools.code.executor import PythonExecutor


class ExeCutorTool(BaseTool):
    name = "executor"
    description = """
    Code Execution Tool (Code Interpreter). Capable of executing Python code within a sandboxed environment and returning the results to the model, thereby enhancing its computational, data processing, and visualization capabilities. This tool should be invoked for questions involving the following scenarios:
    1. When accurate and efficient numerical computation or large-scale data processing is required, and manual calculations are prone to errors.
    2. When dealing with complex mathematical operations (e.g., matrix computations, statistical analysis, integration, differentiation) or extensive repetitive calculations.
    3. When generating charts, tables, or other visualizations is necessary to aid in explaining and presenting data.
    4. When reading, parsing, converting, or analyzing user-provided data files (e.g., CSV, Excel, JSON) is required.
    5. When automated text processing, natural language analysis, or operations relying on third-party Python libraries are needed.
    6. When verifying or demonstrating algorithmic steps, ensuring logical correctness through actual code execution.
    7. When aiming to improve the rigor and reproducibility of answers by automatically generating results via scripts and presenting them to the user.
    """

    parameters = {
        "type": "object",
        "properties": {
            "code": {
                "type": "string",
                "description": "A complete Python script intended for execution in a sandboxed environment, which must comply with the following requirements:\n"
                            "1. Self-contained: The script must include all necessary import statements and must not rely on external files or user interaction (e.g., use of input(), stdin, etc. is prohibited).\n"
                            "2. Output display: All results (numerical values, text, tables, images, etc.) must be output via print(); if generating charts, plt.savefig() must be called at the end of the script to save the file.\n"
                            "3. Library restrictions: Only standard libraries and pre-approved third-party libraries available in the sandbox are permitted; installing new packages via the internet is prohibited.\n"
                            "4. Error handling: If script execution results in an error, the model must adjust the code based on the error message to ensure successful execution.\n"
                            "5. Reproducibility: Script execution must not depend on random seeds; if randomness is used, a random seed must be explicitly set to ensure reproducibility."
            }
        },
        "required": ["code"]
    }

    def __init__(self):
        super().__init__()

    def execute(self, args: Dict) -> Dict[str, Any]:
        try:
            code = args["code"]
            executor = PythonExecutor()
            prediction = executor.apply(code)
            # print(str(prediction))
            # if str(prediction) == str("Done"):
            return {"content": prediction[0], "success": True}
        except Exception as e:
            return {"content": str(e), "success": False}

    def batch_execute(self, args_list: List[Dict]) -> List[Dict[str, Any]]:
        try:
            batch_code = [x["code"] for x in args_list]
            executor = PythonExecutor()
            predictions = executor.batch_apply(batch_code)
            return [{"content": str(content[0]), "success": True} for content in predictions]
        except Exception as e:
            return [{"content": str(e), "success": False} for _ in args_list]