import inspect
from collections.abc import AsyncIterator

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
from langchain_core.tools import StructuredTool
from langchain_openai import ChatOpenAI

from tests.fixtures.agents.prompt import SYSTEM_PROMPT
from tests.fixtures.datasource import datasource


class LangchainAgent:
    MAX_ITERATIONS = 3
    MAX_STREAM_ITERATIONS = 4
    MODEL_NAME = "gpt-4o-mini"
    TEMPERATURE = 0

    def __init__(self):
        self.tools = self._create_tools()
        self.function_map = {func.__name__: func for func in datasource}

    def ask(self, question: str) -> str:
        model = self._create_model()
        messages = [
            SystemMessage(content=SYSTEM_PROMPT),
            HumanMessage(content=question),
        ]

        for _ in range(self.MAX_ITERATIONS):
            response = model.invoke(messages)

            if response.tool_calls:
                messages.append(response)
                self._execute_tool_calls(response.tool_calls, messages)
            else:
                return response.content

        return response.content if hasattr(response, "content") else str(response)

    async def stream(self, question: str) -> AsyncIterator[str]:
        model = self._create_model(streaming=True)
        messages = [
            SystemMessage(content=SYSTEM_PROMPT),
            HumanMessage(content=question),
        ]

        for _ in range(self.MAX_STREAM_ITERATIONS):
            tool_calls_by_id = {}
            content_buffer = ""
            ai_message_chunks = []

            async for chunk in model.astream(messages):
                ai_message_chunks.append(chunk)
                if chunk.content:
                    content_buffer += chunk.content
                    yield chunk.content

                if hasattr(chunk, "tool_calls") and chunk.tool_calls:
                    self._accumulate_tool_calls(chunk.tool_calls, tool_calls_by_id)

            tool_calls_accumulated = list(tool_calls_by_id.values())

            if ai_message_chunks:
                messages.append(self._build_ai_message(content_buffer, tool_calls_accumulated))

            if tool_calls_accumulated:
                self._execute_tool_calls(tool_calls_accumulated, messages)

                async for content in self._handle_following_responses(model, messages):
                    yield content
                return
            else:
                return

    def _accumulate_tool_calls(self, tool_calls: list, tool_calls_by_id: dict) -> None:
        for tc in tool_calls:
            if not (isinstance(tc, dict) and tc.get("id")):
                continue

            tool_call_id = tc["id"]
            if tool_call_id not in tool_calls_by_id:
                tool_calls_by_id[tool_call_id] = tc
            else:
                self._merge_tool_call(tool_calls_by_id[tool_call_id], tc)

    def _build_ai_message(self, content: str, tool_calls: list) -> AIMessage:
        message_kwargs = {"content": content}
        if tool_calls:
            message_kwargs["tool_calls"] = tool_calls
        return AIMessage(**message_kwargs)

    def _create_model(self, streaming: bool = False) -> ChatOpenAI:
        return ChatOpenAI(
            model=self.MODEL_NAME,
            temperature=self.TEMPERATURE,
            streaming=streaming
        ).bind_tools(self.tools)

    def _create_tools(self) -> list[StructuredTool]:
        return [StructuredTool.from_function(func) for func in datasource]

    def _execute_tool_call(self, tool_call: dict) -> ToolMessage:
        tool_name = tool_call.get("name", "")
        tool_args = tool_call.get("args", {})
        tool_call_id = tool_call.get("id", "unknown")

        if not tool_name or not isinstance(tool_args, dict):
            return ToolMessage(
                content="Error: Invalid tool call format",
                tool_call_id=tool_call_id,
            )

        if not self._tool_is_ready(tool_name, tool_args):
            return ToolMessage(
                content=f"Error: Tool {tool_name} requires additional parameters",
                tool_call_id=tool_call_id,
            )

        try:
            result = self.function_map[tool_name](**tool_args)
            return ToolMessage(content=str(result), tool_call_id=tool_call_id)
        except Exception as e:
            return ToolMessage(content=f"Error: {e!r}", tool_call_id=tool_call_id)

    def _execute_tool_calls(self, tool_calls: list, messages: list) -> None:
        for tool_call in tool_calls:
            messages.append(self._execute_tool_call(tool_call))

    async def _handle_following_responses(self, model: ChatOpenAI, messages: list) -> AsyncIterator[str]:
        max_follow_ups = 5
        for _ in range(max_follow_ups):
            following = model.invoke(messages)
            if following.content:
                yield following.content
                return
            elif following.tool_calls:
                messages.append(following)
                self._execute_tool_calls(following.tool_calls, messages)
            else:
                break

    def _merge_tool_call(self, existing: dict, new: dict) -> None:
        for key, value in new.items():
            if key == "name":
                if value and value != "":
                    existing[key] = value
            elif key == "args":
                if isinstance(value, dict) and value:
                    if isinstance(existing.get(key), dict):
                        existing[key].update({k: v for k, v in value.items() if v is not None})
                    else:
                        existing[key] = value
            elif value is not None:
                existing[key] = value

    def _tool_is_ready(self, tool_name: str, tool_args: dict) -> bool:
        if tool_name not in self.function_map:
            return False

        func = self.function_map[tool_name]
        sig = inspect.signature(func)

        required_params = [
            param_name
            for param_name, param in sig.parameters.items()
            if param.default == inspect.Parameter.empty
        ]

        if not required_params:
            return True

        return all(
            param_name in tool_args and tool_args[param_name] is not None
            for param_name in required_params
        )
