# Copyright (c) ModelScope Contributors. All rights reserved.
import json
import re
from typing import List, Optional, Union

from swift.infer_engine import Function
from .hermes import HermesAgentTemplate


def render_extra_keys(obj, handled_keys):
    """Helper function to render extra keys not explicitly handled"""
    result = ''
    if isinstance(obj, dict):
        for key, value in obj.items():
            if key not in handled_keys:
                result += f'\n<{key}>{json.dumps(value, ensure_ascii=False)}</{key}>'
    return result


TOOL_DESC_SUFFIX = (
    '</tools>\n\nIf you choose to call a function ONLY reply in the following format with '
    'NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\n'
    'value_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\n'
    'that can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\n'
    'Reminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> '
    'block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n'
    '- You may provide optional reasoning for your function call in natural language BEFORE the function call, '
    'but NOT after\n- If there is no function call available, '
    'answer the question like normal with your current '
    'knowledge and do not tell the user about function calls\n</IMPORTANT>')


class Qwen3CoderAgentTemplate(HermesAgentTemplate):

    @staticmethod
    def _find_function_call(single_content: str) -> Optional[Function]:
        single_content = single_content.strip()
        # Check whether the complete function tag is included
        if not single_content.startswith('<function=') or not single_content.endswith('</function>'):
            return None

        # Extract function name
        func_name_match = re.search(r'<function=([^>]+)>', single_content)
        if not func_name_match:
            return None

        func_name = func_name_match.group(1).strip()
        parameters = {}

        # Use regular expressions to match parameters
        # Match any content of <parameter=name>content</parameter>
        param_pattern = r'<parameter=([^>]+)>\s*(.*?)\s*</parameter>'
        param_matches = re.findall(param_pattern, single_content, re.DOTALL)

        for param_name, param_value in param_matches:
            # Clear the parameter values and remove any possible additional whitespace
            clean_value = param_value.strip()
            parameters[param_name.strip()] = clean_value

        return Function(name=func_name, arguments=json.dumps(parameters, ensure_ascii=False))

    def get_toolcall(self, response: str) -> List[Function]:
        # Extract the tool call parameters from the model's response
        toolcall_list = re.findall(r'<tool_call>(.*?)</tool_call>', response, re.DOTALL)
        functions = []
        for toolcall in toolcall_list:
            function = self._find_function_call(toolcall)
            if function:
                functions.append(function)
        if len(functions) == 0:
            # Compat react_en
            return super(HermesAgentTemplate, self).get_toolcall(response)
        return functions

    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
        if system is None:
            system = 'You are Qwen, a helpful AI assistant that can interact with a computer to solve tasks.'
        tool_descs = [f'{system}\n\n# Tools\n\nYou have access to the following functions:\n\n<tools>']
        for tool in tools:
            tool_desc = ''

            # Check function key
            if isinstance(tool, dict) and 'function' in tool:
                tool = tool['function']

            # Add function name
            tool_desc += f"<function>\n<name>{tool['name']}</name>"

            # Add description if available
            if 'description' in tool:
                tool_desc += f"\n<description>{tool['description'].strip()}</description>"

            # Add parameters section
            tool_desc += '\n<parameters>'

            # Process parameters if they exist in the expected structure
            if ('parameters' in tool and isinstance(tool['parameters'], dict) and 'properties' in tool['parameters']
                    and isinstance(tool['parameters']['properties'], dict)):

                for param_name, param_fields in tool['parameters']['properties'].items():
                    tool_desc += '\n<parameter>'
                    tool_desc += f'\n<name>{param_name}</name>'

                    if 'type' in param_fields:
                        tool_desc += f"\n<type>{str(param_fields['type'])}</type>"

                    if 'description' in param_fields:
                        tool_desc += f"\n<description>{param_fields['description'].strip()}</description>"

                    # Add any extra parameter fields
                    handled_param_keys = ['name', 'type', 'description']
                    tool_desc += render_extra_keys(param_fields, handled_param_keys)

                    tool_desc += '\n</parameter>'
            # Add any extra parameter section fields
            handled_keys = ['type', 'properties']
            if 'parameters' in tool:
                tool_desc += render_extra_keys(tool['parameters'], handled_keys)

            tool_desc += '\n</parameters>'

            # Add any extra function fields
            handled_keys = ['type', 'name', 'description', 'parameters']
            tool_desc += render_extra_keys(tool, handled_keys)

            tool_desc += '\n</function>'

            tool_descs.append(tool_desc)

        tool_descs.append(TOOL_DESC_SUFFIX)
        tool_descs = '\n'.join(tool_descs)
        return tool_descs

    def _format_tool_calls(self, tool_call_messages):
        result_parts = []
        for idx, message in enumerate(tool_call_messages):
            tool_call = self._parse_tool_call(message['content'])
            result_parts.append(f"<tool_call>\n<function={tool_call['name']}>\n")
            # Processing parameters (if present)
            if 'arguments' in tool_call and tool_call['arguments']:
                for args_name, args_value in tool_call['arguments'].items():
                    result_parts.append(f'<parameter={args_name}>\n')
                    # Handle different types of parameter values
                    if isinstance(args_value, (dict, list)):
                        # For dictionaries or lists, use json formatting
                        args_value = json.dumps(args_value, ensure_ascii=False)
                    else:
                        # For other types, convert to strings
                        args_value = str(args_value)
                    result_parts.append(f'{args_value}\n</parameter>\n')
            # Close tags
            result_parts.append('</function>\n</tool_call>')
            # ref: https://github.com/QwenLM/Qwen3-Coder/blob/0ae30f55e9d6c47ff763c334f99c135ad68915dd/qwencoder-eval/tool_calling_eval/berkeley-function-call-leaderboard/bfcl_eval/model_handler/local_inference/qwen_fc.py#L21  # noqa
            if idx != len(tool_call_messages) - 1:
                result_parts.append('\n')
        return ''.join(result_parts)

    def _get_tool_responses(self, tool_messages):
        res_tool = []
        for tool_message in tool_messages:
            tool_content = tool_message['content']
            res_tool.append(f'<tool_response>\n{tool_content}\n</tool_response>\n')
        return ''.join(res_tool)


class Qwen3_5AgentTemplate(Qwen3CoderAgentTemplate):

    def _add_tool_call_prefix(self, tool_content: str, pre_message=None) -> str:
        """Qwen3.5/3.6 jinja inserts \n\n between assistant content and <tool_call>
        only when effective content (after think removal) is non-empty."""
        if not pre_message or pre_message.get('role') != 'assistant':
            return tool_content
        content = pre_message.get('content', '')
        if not isinstance(content, str):
            return tool_content
        # Mirror jinja: content.split('</think>')[-1].lstrip('\n') then content|trim
        if '</think>' in content:
            effective = content.split('</think>')[-1].lstrip('\n')
        else:
            effective = content
        if effective.strip():
            return '\n\n' + tool_content
        return tool_content

    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
        tool_descs = [json.dumps(self.wrap_tool(tool), ensure_ascii=False) for tool in tools]
        tools_prompt = """# Tools

You have access to the following functions:\n\n<tools>
""" + '\n'.join(tool_descs) + f'\n{TOOL_DESC_SUFFIX}'
        if system:
            tools_prompt += f'\n\n{system}'
        return tools_prompt

    def _get_tool_responses(self, tool_messages):
        res_tool = []
        for tool_message in tool_messages:
            tool_content = tool_message['content']
            res_tool.append(f'<tool_response>\n{tool_content}\n</tool_response>')
        return '\n'.join(res_tool)
