from __future__ import annotations

import inspect
import json
from functools import wraps
from typing import Any

from frisk_sdk.core.llm_reasoning import LLM_REASONING_ARG_NAME


def wrap_tool_with_reasoning(tool: Any) -> Any:
    """Wrap a Strands tool/callable so it accepts optional llm_reasoning__ metadata."""
    if getattr(tool, "__frisk_reasoning_wrapped__", False):
        return tool

    if _is_strands_tool(tool):
        return _wrap_existing_strands_tool(tool)

    if callable(tool):
        return _wrap_plain_callable(tool)

    raise ValueError(f"Cannot wrap tool of type {type(tool)} with reasoning")


def _wrap_plain_callable(tool: Any) -> Any:
    signature = inspect.signature(tool)
    wrapped_signature = _signature_with_reasoning(signature)

    if inspect.iscoroutinefunction(tool):

        @wraps(tool)
        async def wrapped_async_tool(*args: Any, **kwargs: Any) -> Any:
            kwargs.pop(LLM_REASONING_ARG_NAME, None)
            return await tool(*args, **kwargs)

        wrapped_tool = wrapped_async_tool

    else:

        @wraps(tool)
        def wrapped_sync_tool(*args: Any, **kwargs: Any) -> Any:
            kwargs.pop(LLM_REASONING_ARG_NAME, None)
            return tool(*args, **kwargs)

        wrapped_tool = wrapped_sync_tool

    wrapped_tool.__signature__ = wrapped_signature  # type: ignore[attr-defined]
    wrapped_tool.__doc__ = _augment_description(getattr(tool, "__doc__", None))
    setattr(wrapped_tool, "__frisk_reasoning_wrapped__", True)
    return wrapped_tool


def _wrap_existing_strands_tool(tool: Any) -> Any:
    original_spec = getattr(tool, "tool_spec", None)
    if not isinstance(original_spec, dict):
        raise ValueError("Expected tool_spec to be a dictionary")

    wrapped_spec = _augment_tool_spec(original_spec)
    if hasattr(tool, "_tool_spec"):
        setattr(tool, "_tool_spec", wrapped_spec)
    else:
        setattr(tool, "tool_spec", wrapped_spec)

    original_stream = tool.stream

    async def wrapped_stream(
        tool_use: dict[str, Any], invocation_state: dict[str, Any], **kwargs: Any
    ):
        normalized_tool_use = dict(tool_use)
        tool_input = normalized_tool_use.get("input")
        if isinstance(tool_input, dict):
            filtered_input = dict(tool_input)
            filtered_input.pop(LLM_REASONING_ARG_NAME, None)
            normalized_tool_use["input"] = filtered_input
        async for event in original_stream(
            normalized_tool_use, invocation_state, **kwargs
        ):
            yield event

    setattr(tool, "stream", wrapped_stream)
    setattr(tool, "__frisk_reasoning_wrapped__", True)
    return tool


def _augment_tool_spec(tool_spec: dict[str, Any]) -> dict[str, Any]:
    wrapped_spec = json.loads(json.dumps(tool_spec))
    wrapped_spec.setdefault("description", "")
    wrapped_spec["description"] = _augment_description(wrapped_spec.get("description"))

    input_schema = wrapped_spec.setdefault("inputSchema", {})
    if "json" not in input_schema or not isinstance(input_schema.get("json"), dict):
        input_schema["json"] = {"type": "object", "properties": {}}

    schema = input_schema["json"]
    schema.setdefault("type", "object")
    properties = schema.setdefault("properties", {})
    if not isinstance(properties, dict):
        properties = {}
        schema["properties"] = properties

    properties[LLM_REASONING_ARG_NAME] = {
        "type": "string",
        "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."
        ),
    }

    return wrapped_spec


def _augment_description(description: str | None) -> str:
    base = description or ""
    suffix = (
        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 f"{base}\n\n{suffix}".strip()


def _signature_with_reasoning(signature: inspect.Signature) -> inspect.Signature:
    parameters = list(signature.parameters.values())
    reasoning_parameter = inspect.Parameter(
        name=LLM_REASONING_ARG_NAME,
        kind=inspect.Parameter.KEYWORD_ONLY,
        default=None,
        annotation=str | None,
    )

    if any(p.name == LLM_REASONING_ARG_NAME for p in parameters):
        return signature

    insert_index = len(parameters)
    for i, parameter in enumerate(parameters):
        if parameter.kind == inspect.Parameter.VAR_KEYWORD:
            insert_index = i
            break

    parameters.insert(insert_index, reasoning_parameter)
    return signature.replace(parameters=parameters)


def _is_strands_tool(tool: Any) -> bool:
    return (
        hasattr(tool, "tool_name")
        and hasattr(tool, "tool_spec")
        and hasattr(tool, "stream")
    )
