
import re
import os
from typing import Dict, List, Optional, Union

import json5

from qwen_agent.tools.base import BaseToolWithFileAccess, register_tool

from qwen_agent.utils.utils import extract_code

from sandbox_fusion import run_code, RunCodeRequest, RunStatus
from requests.exceptions import Timeout

SANDBOX_FUSION_ENDPOINT=os.environ.get('SANDBOX_FUSION_ENDPOINT')

class PythonInterpreter(BaseToolWithFileAccess):
    name = "PythonInterpreter"
    description = 'Execute Python code in a sandboxed environment. Use this to run Python code and get the execution results.\n**Make sure to use print() for any output you want to see in the results.**\nFor code parameters, use placeholders first, and then put the code within <code></code> XML tags, such as:\n<tool_call>\n{"purpose": <detailed-purpose-of-this-tool-call>, "name": <tool-name>, "arguments": {"code": ""}}\n<code>\nHere is the code.\n</code>\n</tool_call>\n'

    parameters = {
        "type": "object",
        "properties": {
            "code": {
                "type": "string",
                "description": "The Python code to execute. Must be provided within <code></code> XML tags. Remember to use print() statements for any output you want to see.",
            }
        },
        "required": ["code"],
    }

    def __init__(self, cfg: Optional[Dict] = None):
        super().__init__(cfg)
    
    @property
    def args_format(self) -> str:
        fmt = self.cfg.get('args_format')
        if fmt is None:
            if has_chinese_chars([self.name_for_human, self.name, self.description, self.parameters]):
                fmt = '此工具的输入应为Markdown代码块。'
            else:
                fmt = 'Enclose the code within triple backticks (`) at the beginning and end of the code.'
        return fmt

    def observation(self, tool: dict, tool_dict: dict, tool_results, empty_mode: bool=False, readpage: bool=False, max_observation_length: int=None, tokenizer=None):
        assert isinstance(tool_results, str), f"result of python code should be str, instead of {type(tool_results)}. {tool_results}"
        return tool_results
    
    @property
    def function(self) -> dict:  # Bad naming. It should be `function_info`.
        return {
            # 'name_for_human': self.name_for_human,
            'name': self.name,
            'description': self.description,
            'parameters': self.parameters,
            # 'args_format': self.args_format
        }

    def call(self, params: Union[str, dict], files: List[str] = None, timeout: Optional[int] = 50, **kwargs) -> str:
        try:
            super().call(params=params, files=files)  # copy remote files to work_dir
            try:
                if type(params) is str:
                    params = json5.loads(params)
                code = params.get('code', '')
                triple_match = re.search(r'```[^\n]*\n(.+?)```', code, re.DOTALL)
                if triple_match:
                    code = triple_match.group(1)
            except Exception:
                code = extract_code(params)

            if not code.strip():
                return {
                    "success": False,
                    "error_message": '[Python Interpreter Error]: Empty code.'
                }

            try:
                code_result = run_code(RunCodeRequest(code=code, language='python', run_timeout=timeout), max_attempts=1, client_timeout=timeout, endpoint=SANDBOX_FUSION_ENDPOINT)

                result = []
                if code_result.run_result.stdout:
                    result.append(f"stdout:\n{code_result.run_result.stdout}")
                if code_result.run_result.stderr:
                    result.append(f"stderr:\n{code_result.run_result.stderr}")
                
                result = '\n'.join(result)

            except Timeout:
                result = '[Python Interpreter Error] TimeoutError: Execution timed out.'
            
            except Exception as e:
                result = f'[Python Interpreter Error]: {str(e)}'

            return {
                "success": True,
                "results": result if result.strip() else 'Finished execution.',
                "params": params
            }

        except Exception as e:

            return {
                "success": False,
                "error_message": f"[Python Interpreter Error]: {str(e)}"
            }
