from typing import Any, cast
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:
    """Return a StructuredTool that accepts optional LLM reasoning metadata."""
    if tool.args_schema is None:
        raise ValueError("Expected a StructuredTool with args_schema")

    if not isinstance(tool.args_schema, type) or not issubclass(
        tool.args_schema, BaseModel
    ):
        raise ValueError("Expected args_schema to be a pydantic BaseModel type")
    original_model: type[BaseModel] = tool.args_schema

    # New model extends the original args schema
    wrapped_model = cast(
        type[BaseModel],
        cast(Any, create_model)(
            f"{original_model.__name__}WithReasoning",
            **{
                LLM_REASONING_ARG_NAME: (
                    str | None,
                    Field(
                        default=None,
                        description="The LLM's reasoning process for this tool call, captured as a string. This should include the LLM's thought process, any intermediate steps, and how it arrived at the final arguments for the tool call.",
                    ),
                )
            },
            __base__=original_model,
        ),
    )

    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=cast(Any, wrapped_model),
        infer_schema=False,
        response_format=getattr(tool, "response_format", "content"),
    )
