import json
import asyncio
from typing import Any, cast

import pytest

from frisk_sdk.core.llm_reasoning import LLM_REASONING_ARG_NAME
from frisk_sdk.adapters.strands.strands_framework_adapter import StrandsFrameworkAdapter
from frisk_sdk.framework_adapter.framework_adapter import NormalizedToolCall


class _ToolCallObj:
    def __init__(
        self,
        *,
        toolUseId: str | None = None,
        name: str | None = None,
        input: dict[str, Any] | None = None,
    ):
        self.toolUseId = toolUseId
        self.name = name
        self.input = input


@pytest.mark.unit
def test_serialize_tool_args_json_roundtrip():
    adapter = StrandsFrameworkAdapter()

    args = {"a": 1, "b": {"c": True}}
    result = adapter.serialize_tool_args(args)

    assert json.loads(cast(str, result.value)) == args


@pytest.mark.unit
def test_serialize_tool_args_redact_true_redacts_all_and_tracks_star_path():
    adapter = StrandsFrameworkAdapter()

    args = {"a": 1, "b": {"c": True, "d": "secret"}}
    result = adapter.serialize_tool_args(args, redact=True)

    assert json.loads(cast(str, result.value)) == {}
    assert "*" in result.redacted_paths


@pytest.mark.unit
def test_normalize_tool_call_from_strands_tool_use(monkeypatch: pytest.MonkeyPatch):
    adapter = StrandsFrameworkAdapter()
    monkeypatch.setattr(adapter, "_cuid_generator", lambda: "generated-id")

    tool_use = {"name": "weather", "input": {"city": "sf"}}
    info = adapter.normalize_tool_call(tool_use)

    assert isinstance(info, NormalizedToolCall)
    assert info == NormalizedToolCall(
        id="generated-id", name="weather", args={"city": "sf"}
    )


@pytest.mark.unit
def test_normalize_tool_call_prefers_tool_use_id():
    adapter = StrandsFrameworkAdapter()
    tool_use = {"toolUseId": "tooluse-1", "name": "weather", "input": {"city": "sf"}}
    info = adapter.normalize_tool_call(tool_use)
    assert info.id == "tooluse-1"


@pytest.mark.unit
def test_normalize_tool_call_falls_back_to_id_and_args():
    adapter = StrandsFrameworkAdapter()
    tool_use = {"id": "fallback-id", "name": "weather", "args": {"city": "sf"}}
    info = adapter.normalize_tool_call(tool_use)
    assert info == NormalizedToolCall(
        id="fallback-id", name="weather", args={"city": "sf"}
    )


@pytest.mark.unit
def test_normalize_tool_call_from_object_tool_use():
    adapter = StrandsFrameworkAdapter()
    tool_use = _ToolCallObj(toolUseId="tooluse-obj", name="math", input={"x": 1})

    info = adapter.normalize_tool_call(cast(dict[str, Any], tool_use))
    assert info == NormalizedToolCall(id="tooluse-obj", name="math", args={"x": 1})


@pytest.mark.unit
def test_normalize_tool_call_raises_for_non_dict_input():
    adapter = StrandsFrameworkAdapter()
    with pytest.raises(ValueError, match="input must be a dictionary"):
        adapter.normalize_tool_call(
            cast(dict[str, Any], {"toolUseId": "id", "name": "t", "input": "bad"})
        )


@pytest.mark.unit
def test_normalize_tool_call_raises_for_missing_name():
    adapter = StrandsFrameworkAdapter()
    with pytest.raises(ValueError, match="missing a valid name"):
        adapter.normalize_tool_call(cast(dict[str, Any], {"input": {"x": 1}}))


@pytest.mark.unit
def test_extract_prompt_returns_latest_user_text_block():
    adapter = StrandsFrameworkAdapter()
    agent_state = {
        "messages": [
            {"role": "user", "content": [{"text": "first"}]},
            {"role": "assistant", "content": [{"text": "response"}]},
            {"role": "user", "content": [{"text": "latest prompt"}]},
        ]
    }

    assert adapter.extract_prompt(agent_state) == "latest prompt"


@pytest.mark.unit
def test_extract_prompt_joins_multiple_text_blocks():
    adapter = StrandsFrameworkAdapter()
    agent_state = {
        "messages": [
            {
                "role": "user",
                "content": [
                    {"text": "line 1"},
                    {"toolUse": {"name": "x"}},
                    {"text": "line 2"},
                ],
            }
        ]
    }

    assert adapter.extract_prompt(agent_state) == "line 1\nline 2"


@pytest.mark.unit
def test_extract_prompt_returns_none_for_missing_prompt():
    adapter = StrandsFrameworkAdapter()
    assert (
        adapter.extract_prompt({"messages": [{"role": "assistant", "content": []}]})
        is None
    )


@pytest.mark.unit
def test_serialize_agent_state_handles_non_json_objects():
    adapter = StrandsFrameworkAdapter()

    class _Agent:
        pass

    state = {
        "frisk_session_id": "123",
        "agent": _Agent(),
        "messages": [{"role": "user", "content": [{"text": "hello"}]}],
    }

    result = adapter.serialize_agent_state(state)
    payload = json.loads(cast(str, result.value))

    assert payload["frisk_session_id"] == "123"
    assert payload["messages"][0]["role"] == "user"
    assert payload["agent"].startswith("<<non-serializable: ")


@pytest.mark.unit
def test_wrap_tools_converts_plain_callables():
    adapter = StrandsFrameworkAdapter()

    def tool_a(x: int) -> int:
        return x + 1

    wrapped = adapter.wrap_tools([tool_a])
    assert len(wrapped) == 1

    wrapped_tool = wrapped[0]
    wrapped_tool_spec = cast(dict[str, Any], wrapped_tool.tool_spec)
    assert wrapped_tool_spec["name"] == "tool_a"
    assert (
        LLM_REASONING_ARG_NAME in wrapped_tool_spec["inputSchema"]["json"]["properties"]
    )
    assert wrapped_tool(x=2, **{LLM_REASONING_ARG_NAME: "because"}) == 3


@pytest.mark.unit
def test_wrap_tools_wraps_existing_strands_tool_with_reasoning():
    adapter = StrandsFrameworkAdapter()

    class _AlreadyTool:
        tool_name = "x"
        tool_spec = {"name": "x", "inputSchema": {"json": {}}, "description": "x"}

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

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

    tool = _AlreadyTool()
    wrapped = adapter.wrap_tools([tool])

    assert wrapped == [tool]
    tool_spec = cast(dict[str, Any], tool.tool_spec)
    assert LLM_REASONING_ARG_NAME in tool_spec["inputSchema"]["json"]["properties"]

    asyncio.run(
        _consume(
            tool.stream(
                {
                    "toolUseId": "id",
                    "name": "x",
                    "input": {"a": 1, LLM_REASONING_ARG_NAME: "secret"},
                },
                {},
            )
        )
    )
    assert tool.calls[0]["input"] == {"a": 1}


@pytest.mark.unit
def test_wrap_tool_raises_for_unsupported_tool_type():
    adapter = StrandsFrameworkAdapter()
    with pytest.raises(ValueError, match="Cannot convert tool of type"):
        adapter.wrap_tool(123)


@pytest.mark.unit
def test_wrap_tool_options_none_wraps_callable_with_reasoning():
    adapter = StrandsFrameworkAdapter()

    def plus_one(x: int) -> int:
        return x + 1

    wrapped = adapter.wrap_tool(plus_one, None)

    wrapped_tool_spec = cast(dict[str, Any], wrapped.tool_spec)
    assert (
        LLM_REASONING_ARG_NAME in wrapped_tool_spec["inputSchema"]["json"]["properties"]
    )
    assert wrapped(x=2, **{LLM_REASONING_ARG_NAME: "because"}) == 3


@pytest.mark.unit
def test_wrap_tool_collect_reasoning_false_does_not_wrap_callable():
    adapter = StrandsFrameworkAdapter()

    def plus_one(x: int) -> int:
        return x + 1

    wrapped = adapter.wrap_tool(plus_one, {"collect_reasoning": False})

    # With collect_reasoning=False, wrap_tool should return the original tool unchanged.
    assert wrapped is plus_one


@pytest.mark.unit
def test_wrap_tool_collect_reasoning_true_wraps_callable_with_reasoning():
    adapter = StrandsFrameworkAdapter()

    def plus_one(x: int) -> int:
        return x + 1

    wrapped = adapter.wrap_tool(plus_one, {"collect_reasoning": True})

    wrapped_tool_spec = cast(dict[str, Any], wrapped.tool_spec)
    assert (
        LLM_REASONING_ARG_NAME in wrapped_tool_spec["inputSchema"]["json"]["properties"]
    )
    assert wrapped(x=2, **{LLM_REASONING_ARG_NAME: "because"}) == 3


async def _consume(generator: Any) -> None:
    async for _ in generator:
        pass
