import json
from collections.abc import Callable, Iterable
from typing import Any, Optional

from cuid2 import cuid_wrapper

from frisk_sdk.framework_adapter.framework_adapter import (
    FrameworkAdapter,
    NormalizedToolCall,
)
from frisk_sdk.redact import (
    RedactOption,
    RedactedDictionary,
    RedactedSerializedObject,
    redact_dictionary,
)
from strands import tool as strands_tool
from .wrap_tool_with_reasoning import wrap_tool_with_reasoning
from ...core.wrap_tool_options import WrapToolOptions


class StrandsFrameworkAdapter(FrameworkAdapter[dict[str, Any], Any]):
    """Adapter that normalizes Strands tool/runtime objects for Frisk evaluation."""

    def __init__(self) -> None:
        self._cuid_generator = cuid_wrapper()

    def tool_args_to_dict(
        self,
        tool_args: Optional[dict[str, Any]],
        *,
        redact: Optional[RedactOption] = None,
    ) -> RedactedDictionary:
        if tool_args is None:
            return RedactedDictionary(value={}, redacted_paths=[])
        return redact_dictionary(tool_args, redact)

    def serialize_tool_args(
        self,
        tool_args: Optional[dict[str, Any]],
        *,
        redact: Optional[RedactOption] = None,
    ) -> RedactedSerializedObject:
        if tool_args is None:
            return RedactedSerializedObject(value=None, redacted_paths=[])
        redacted_dictionary = self.tool_args_to_dict(tool_args, redact=redact)
        return RedactedSerializedObject(
            value=json.dumps(redacted_dictionary.value),
            redacted_paths=redacted_dictionary.redacted_paths,
        )

    def agent_state_to_dict(
        self,
        agent_state: dict[str, Any],
        *,
        redact: Optional[RedactOption] = None,
    ) -> RedactedDictionary:
        json_safe_agent_state = self._normalize_json_like(agent_state)
        return redact_dictionary(json_safe_agent_state, redact)

    def serialize_agent_state(
        self,
        agent_state: Optional[dict[str, Any]],
        *,
        redact: Optional[RedactOption] = None,
    ) -> RedactedSerializedObject:
        if agent_state is None:
            return RedactedSerializedObject(value=None, redacted_paths=[])
        redacted_dictionary = self.agent_state_to_dict(agent_state, redact=redact)
        return RedactedSerializedObject(
            value=json.dumps(redacted_dictionary.value),
            redacted_paths=redacted_dictionary.redacted_paths,
        )

    def normalize_tool_call(self, tool_call: dict[str, Any]) -> NormalizedToolCall:
        tool_name = self.get_tool_name(tool_call)
        if not isinstance(tool_name, str) or not tool_name:
            raise ValueError("Tool call is missing a valid name")

        tool_args = self.get_tool_args(tool_call)
        if tool_args is None:
            tool_args = {}
        if not isinstance(tool_args, dict):
            raise ValueError("Tool call input must be a dictionary")

        return NormalizedToolCall(
            id=self.get_or_create_tool_call_id(tool_call),
            name=tool_name,
            args=tool_args,
        )

    def get_or_create_tool_call_id(self, tool_call: dict[str, Any]) -> str:
        tool_use_id = self._get_attr_or_key(tool_call, "toolUseId")
        if isinstance(tool_use_id, str) and tool_use_id:
            return tool_use_id

        fallback_id = self._get_attr_or_key(tool_call, "id")
        if isinstance(fallback_id, str) and fallback_id:
            return fallback_id

        return self._cuid_generator()

    def get_tool_name(self, tool_call: dict[str, Any]) -> str | None:
        name = self._get_attr_or_key(tool_call, "name")
        return name if isinstance(name, str) else None

    def get_tool_args(self, tool_call: dict[str, Any]) -> dict[str, Any] | None:
        tool_input = self._get_attr_or_key(tool_call, "input")
        if tool_input is not None:
            return tool_input

        args = self._get_attr_or_key(tool_call, "args")
        if args is not None:
            return args

        return None

    def extract_prompt(self, agent_state: dict[str, Any]) -> str | None:
        messages = agent_state.get("messages")
        if isinstance(messages, list):
            for message in reversed(messages):
                if not isinstance(message, dict) or message.get("role") != "user":
                    continue
                content = message.get("content")
                if not isinstance(content, list):
                    continue
                text_blocks: list[str] = []
                for block in content:
                    if not isinstance(block, dict):
                        continue
                    text = block.get("text")
                    if isinstance(text, str):
                        text_blocks.append(text)
                prompt = "\n".join(text_blocks).strip()
                if prompt:
                    return prompt
        return None

    def wrap_tools(
        self, tools: Iterable[Any], options: WrapToolOptions | None = None
    ) -> list[Any]:
        return [self.wrap_tool(tool, options) for tool in tools]

    def wrap_tool(self, tool: Any, options: WrapToolOptions | None = None) -> Any:
        wrapped_tool = tool

        if (options or {}).get("collect_reasoning") is not False:
            if self._is_strands_tool(tool):
                wrapped_tool = wrap_tool_with_reasoning(tool)
            elif callable(tool):
                wrapped_tool = self.wrap_callable_with_reasoning(tool)
            else:
                raise ValueError(
                    f"Cannot convert tool of type {type(tool)} to a Strands tool"
                )
        return wrapped_tool

    def wrap_callable_with_reasoning(self, tool: Callable[..., Any]) -> Any:
        return strands_tool(wrap_tool_with_reasoning(tool))

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

    def _get_attr_or_key(self, source: Any, field: str) -> Any:
        attr_value = getattr(source, field, None)
        if attr_value is not None:
            return attr_value
        if isinstance(source, dict):
            return source.get(field)
        return None

    def _normalize_json_like(self, value: Any) -> dict[str, Any]:
        normalized = json.loads(json.dumps(value, default=self._json_default))
        if isinstance(normalized, dict):
            return normalized
        return {"value": normalized}

    def _json_default(self, value: Any) -> str:
        return f"<<non-serializable: {type(value).__qualname__}>>"
