from gallama.logger.logger import logger
import json
import re
from .....data_classes.data_class import (
    ChoiceDeltaToolCall,
    ChoiceDeltaToolCallFunction,
    TagDefinition
)
from .....utils.utils import get_response_tool_uid
from typing import List, Dict, Optional

def qwen3_tool_parser(tool_text: str, extra_vars: dict = None) -> List[Dict]:
    """
    Parse Qwen 3 tool call format:

    <function=tool_name>
    <parameter=param_key_1>
    param_value_1
    </parameter>
    <parameter=param_key_2>
    param_value_2
    </parameter>
    </function>

    Multiple <function>...</function> blocks may appear within a single
    <tool_call>...</tool_call> pair.
    """

    results = []

    # Initialize state
    if extra_vars is None:
        extra_vars = {"state": {}}
    if not extra_vars.get("state"):
        extra_vars["state"] = {}
    if "tool_call" not in extra_vars["state"]:
        extra_vars["state"]["tool_call"] = 0

    if not tool_text or not tool_text.strip():
        return []

    try:
        # Find all <function=name>...</function> blocks
        function_pattern = re.compile(
            r'<function=([^>]+)>\s*(.*?)\s*</function>',
            re.DOTALL
        )

        for func_match in function_pattern.finditer(tool_text):
            tool_name = func_match.group(1).strip()
            func_body = func_match.group(2)

            # Extract all <parameter=key>value</parameter> pairs from the function body
            param_pattern = re.compile(
                r'<parameter=([^>]+)>\s*(.*?)\s*</parameter>',
                re.DOTALL
            )

            arguments_dict = {}
            for param_match in param_pattern.finditer(func_body):
                key = param_match.group(1).strip()
                value_text = param_match.group(2).strip()

                # Attempt to parse as JSON (handles lists, dicts, numbers, booleans)
                try:
                    arguments_dict[key] = json.loads(value_text)
                except (json.JSONDecodeError, TypeError):
                    arguments_dict[key] = value_text

            logger.info(f"tool: {tool_name} args: {arguments_dict}")

            chunk_data = ChoiceDeltaToolCall(
                index=extra_vars["state"]["tool_call"],
                id=get_response_tool_uid(),
                function=ChoiceDeltaToolCallFunction(
                    name=tool_name,
                    arguments=json.dumps(arguments_dict)
                ),
                type="function"
            )

            results.append(chunk_data.model_dump(exclude_unset=True))
            extra_vars["state"]["tool_call"] += 1

    except Exception as e:
        raise e

    return results


def tool_prompt(tool_name: Optional[str] = None):
    if tool_name is None:
        return '\n<tool_call>\n<function='
    else:
        return f'\n<tool_call>\n<function={tool_name}>\n'


qwen3_moe = {
    "tool": TagDefinition(
        start_marker="<tool_call>",
        end_marker="</tool_call>",
        tag_type="tool_calls",
        api_tag="tool_calls",
        role="assistant",
        post_processor=qwen3_tool_parser,
        prompt_init=tool_prompt,
        wait_till_complete=True
    ),
    "thinking": TagDefinition(
        start_marker="<think>",
        end_marker="</think>",
        tag_type="thinking",
        role="assistant",
        allowed_roles={"assistant", "tool"},
        api_tag="reasoning"
    )
}