from typing import Optional
from loguru import logger
from lmos_openai_types import CreateChatCompletionRequest
import mcp.types
import json

from mcp_bridge.mcp_clients.McpClientManager import ClientManager
from mcp_bridge.tool_mappers import mcp2openai


async def chat_completion_add_tools(request: CreateChatCompletionRequest):
    # Check if the last user message contains the magic word
    magic_word = "#use-tools"  # Lowercase for case-insensitive comparison
    should_add_tools = False
    logger.debug(json.dumps(request.model_dump_json(), indent=2))
    # Find the last user message
    for message in reversed(request.messages):
        # Use model_dump() to access the fields as a dictionary
        message_dict = message.model_dump()
        if message_dict.get("role") == "user":
            # Check if the magic word is in the message content
            content = message_dict.get("content")
            if content:
                if isinstance(content, str) and magic_word in content.lower():
                    should_add_tools = True
                elif isinstance(content, list):
                    # Handle content as a list of content parts
                    for part in content:
                        if isinstance(part, dict) and part.get('type') == 'text' and magic_word in part.get('text', '').lower():
                            should_add_tools = True
                            break
            break

    if not should_add_tools:
        return request

    logger.info("> using #MCP tools ")
    request.tools = []

    for _, session in ClientManager.get_clients():
        # if session is None, then the client is not running
        if session.session is None:
            logger.error(f"session is `None` for {session.name}")
            continue

        tools = await session.session.list_tools()
        for tool in tools.tools:
            request.tools.append(mcp2openai(tool))

    if len(request.tools) == 0:
        logger.info("> no #MCP tools found")
        request.tools = None

    return request


async def call_tool(
    tool_call_name: str, tool_call_json: str, timeout: Optional[int] = None
) -> Optional[mcp.types.CallToolResult]:
    if tool_call_name == "" or tool_call_name is None:
        logger.error("tool call name is empty")
        return None

    if tool_call_json is None:
        logger.error("tool call json is empty")
        return None

    session = await ClientManager.get_client_from_tool(tool_call_name)

    if session is None:
        logger.error(f"session is `None` for {tool_call_name}")
        return None

    try:
        tool_call_args = json.loads(tool_call_json)
    except json.JSONDecodeError:
        logger.error(f"failed to decode json for {tool_call_name}")
        return None

    return await session.call_tool(tool_call_name, tool_call_args, timeout)