import asyncio
import inspect
from typing import Any

import pytest

from frisk_sdk.adapters.strands.wrap_tool_with_reasoning import wrap_tool_with_reasoning
from frisk_sdk.core.llm_reasoning import LLM_REASONING_ARG_NAME


@pytest.mark.unit
def test_wrap_plain_callable_sync_happy_path():
    calls: list[dict[str, Any]] = []

    def tool_fn(x: int, y: int = 1, **kwargs: Any) -> int:
        calls.append({"x": x, "y": y, "kwargs": kwargs})
        return x + y

    wrapped = wrap_tool_with_reasoning(tool_fn)

    signature = inspect.signature(wrapped)
    assert LLM_REASONING_ARG_NAME in signature.parameters
    assert wrapped(2, y=3, **{LLM_REASONING_ARG_NAME: "because"}) == 5
    assert calls == [{"x": 2, "y": 3, "kwargs": {}}]
    assert getattr(wrapped, "__frisk_reasoning_wrapped__", False) is True


@pytest.mark.unit
def test_wrap_plain_callable_async_happy_path():
    calls: list[dict[str, Any]] = []

    async def tool_fn(x: int, **kwargs: Any) -> int:
        calls.append({"x": x, "kwargs": kwargs})
        return x + 10

    wrapped = wrap_tool_with_reasoning(tool_fn)
    signature = inspect.signature(wrapped)
    assert LLM_REASONING_ARG_NAME in signature.parameters
    result = asyncio.run(wrapped(5, **{LLM_REASONING_ARG_NAME: "step by step"}))
    assert result == 15
    assert calls == [{"x": 5, "kwargs": {}}]


@pytest.mark.unit
def test_wrap_is_idempotent_for_already_wrapped_callable():
    def tool_fn(x: int) -> int:
        return x

    once = wrap_tool_with_reasoning(tool_fn)
    twice = wrap_tool_with_reasoning(once)
    assert twice is once


@pytest.mark.unit
def test_wrap_existing_strands_tool_happy_path():
    class _Tool:
        tool_name = "my_tool"
        tool_spec = {"name": "my_tool", "description": "d", "inputSchema": {"json": {}}}

        def __init__(self):
            self.seen: list[dict[str, Any]] = []

        async def stream(
            self,
            tool_use: dict[str, Any],
            invocation_state: dict[str, Any],
            **kwargs: Any,
        ):
            self.seen.append(tool_use)
            yield {"status": "ok"}

    tool = _Tool()
    wrapped = wrap_tool_with_reasoning(tool)

    assert wrapped is tool
    assert (
        LLM_REASONING_ARG_NAME in wrapped.tool_spec["inputSchema"]["json"]["properties"]
    )
    assert "IMPORTANT" in wrapped.tool_spec["description"]

    async def _run() -> None:
        async for _ in wrapped.stream(
            {"name": "my_tool", "input": {"a": 1, LLM_REASONING_ARG_NAME: "hidden"}}, {}
        ):
            pass

    asyncio.run(_run())
    assert wrapped.seen == [{"name": "my_tool", "input": {"a": 1}}]


@pytest.mark.unit
def test_wrap_existing_strands_tool_raises_when_spec_not_dict():
    class _Tool:
        tool_name = "bad"
        tool_spec = "not-a-dict"

        async def stream(
            self,
            tool_use: dict[str, Any],
            invocation_state: dict[str, Any],
            **kwargs: Any,
        ):
            yield {"ok": True}

    with pytest.raises(ValueError, match="tool_spec"):
        wrap_tool_with_reasoning(_Tool())


@pytest.mark.unit
def test_wrap_rejects_non_callable_non_tool():
    with pytest.raises(ValueError, match="Cannot wrap tool"):
        wrap_tool_with_reasoning(42)
