from typing import Type
from pydantic import BaseModel, Field, create_model
from langchain_core.tools import StructuredTool, BaseTool

from frisk_sdk.core.llm_reasoning import LLM_REASONING_ARG_NAME


def wrap_tool_with_reasoning(tool: BaseTool) -> StructuredTool:
    if tool.args_schema is None:
        raise ValueError("Expected a StructuredTool with args_schema")

    Original: Type[BaseModel] = tool.args_schema

    # New model extends the original args schema
    Wrapped = create_model(
        f"{Original.__name__}WithReasoning",
        **{
            LLM_REASONING_ARG_NAME: (
                str | None,
                Field(
                    default=None,
                    description="Briefly explain why you're calling this tool and what you expect to get.",
                ),
            )
        },
        __base__=Original,
    )

    def wrapped_func(**kwargs):
        kwargs.pop(LLM_REASONING_ARG_NAME, None)
        return tool.invoke(kwargs)

    async def wrapped_coroutine(**kwargs):
        kwargs.pop(LLM_REASONING_ARG_NAME, None)
        return await tool.ainvoke(kwargs)

    wrapped_description = (
        f"{tool.description}\n\n"
        f"IMPORTANT: Please include the {LLM_REASONING_ARG_NAME} argument, "
        "but do not include its contents in your final response to the human end user."
    )
    return StructuredTool.from_function(
        func=wrapped_func,
        coroutine=wrapped_coroutine,
        name=tool.name,
        description=wrapped_description,
        return_direct=getattr(tool, "return_direct", False),
        args_schema=Wrapped,
        infer_schema=False,
        response_format=getattr(tool, "response_format", "content"),
    )
