import inspect

from langchain_core.load import dumps, dumpd
from langchain_core.tools import BaseTool, StructuredTool
from typing import Any, Optional, Iterable, Callable

from langchain_core.messages import ToolCall, BaseMessage

from cuid2 import cuid_wrapper

from .find_last_match import find_last_match
from frisk_sdk.redact import (
    RedactOption,
    redact_dictionary,
    RedactedDictionary,
    RedactedSerializedObject,
    combine_redact_options,
)
from frisk_sdk.framework_adapter.framework_adapter import (
    FrameworkAdapter,
    NormalizedToolCall,
)
from .wrap_tool_with_reasoning import wrap_tool_with_reasoning
from ...core.wrap_tool_options import WrapToolOptions


class LangchainFrameworkAdapter(FrameworkAdapter[ToolCall, StructuredTool]):
    """Adapter that normalizes LangChain tool/runtime objects for Frisk evaluation."""

    # Keys to exclude from agent state when serializing for telemetry
    EXCLUDED_STATE_KEYS = ["messages"]

    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=[])

        raw_dict = dumpd(tool_args)
        return redact_dictionary(raw_dict, 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=[])

        tool_args_redacted_result = self.tool_args_to_dict(tool_args, redact=redact)
        return RedactedSerializedObject(
            value=dumps(tool_args_redacted_result.value),
            redacted_paths=tool_args_redacted_result.redacted_paths,
        )

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

    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=dumps(redacted_dictionary.value),
            redacted_paths=redacted_dictionary.redacted_paths,
        )

    def normalize_tool_call(self, tool_call: ToolCall) -> 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")
        return NormalizedToolCall(
            id=self.get_or_create_tool_call_id(tool_call),
            name=tool_name,
            args=tool_call["args"],
        )

    def get_or_create_tool_call_id(self, tool_call: ToolCall) -> str:
        attr_id = getattr(tool_call, "id", None)
        if isinstance(attr_id, str) and attr_id:
            return attr_id

        if isinstance(tool_call, dict):
            dict_id = tool_call.get("id")
            if isinstance(dict_id, str) and dict_id:
                return dict_id

        return self._cuid_generator()

    def get_tool_name(self, tool_call: ToolCall) -> str | None:
        return getattr(tool_call, "name", None) or (
            tool_call.get("name") if isinstance(tool_call, dict) else None
        )

    def extract_prompt(self, agent_state: dict[str, Any]) -> str | None:
        messages: list[BaseMessage] = agent_state.get("messages", [])
        if not messages:
            return None

        last_human_message = find_last_match(messages, lambda m: m.type == "human")
        return last_human_message.text if last_human_message else None

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

    def wrap_tool(
        self,
        tool: Any,
        options: WrapToolOptions | None = None,
    ) -> StructuredTool:
        structured_tool = self.convert_callable_to_tool(tool)
        wrapped_tool = structured_tool
        if (options or {}).get("collect_reasoning") is not False:
            wrapped_tool = wrap_tool_with_reasoning(structured_tool)
        return wrapped_tool

    def convert_callable_to_tool(
        self, original_tool: Callable[..., Any]
    ) -> StructuredTool:
        if isinstance(original_tool, StructuredTool):
            return original_tool
        if isinstance(original_tool, BaseTool):
            return StructuredTool.from_function(
                func=original_tool.invoke,
                coroutine=original_tool.ainvoke,
                name=original_tool.name,
                description=original_tool.description,
                return_direct=getattr(original_tool, "return_direct", False),
                args_schema=original_tool.args_schema,
                response_format=original_tool.response_format,
            )
        if inspect.isfunction(original_tool):
            return StructuredTool.from_function(original_tool)
        raise ValueError(
            f"Cannot convert tool of type {type(original_tool)} to StructuredTool"
        )
