"""Agent wrapper that applies middleware hooks.

Provides create_agent() as a convenience for building middleware-aware agents.
"""

from __future__ import annotations

from typing import Any

from langchain_nobulex.base import AgentMiddleware, ModelResponse, ToolCall


class MiddlewareAgent:
    """An agent wrapper that applies middleware hooks around the agent loop.

    This wraps any callable agent (or LangChain agent) and applies
    registered middleware at each stage of execution.
    """

    def __init__(
        self,
        *,
        model: str = "gpt-4.1",
        tools: list[Any] | None = None,
        middleware: list[AgentMiddleware] | None = None,
    ) -> None:
        self.model = model
        self.tools = tools or []
        self.middleware = middleware or []
        self._state: dict[str, Any] = {}

    def apply_after_model(self, response: ModelResponse) -> ModelResponse:
        """Run all middleware after_model hooks in order."""
        for mw in self.middleware:
            response = mw.after_model(response)
        return response

    def apply_wrap_tool_call(self, tool_call: ToolCall) -> ToolCall | None:
        """Run all middleware wrap_tool_call hooks in order."""
        current: ToolCall | None = tool_call
        for mw in self.middleware:
            if current is None:
                break
            current = mw.wrap_tool_call(current)
        return current

    def apply_after_agent(self, state: dict[str, Any]) -> dict[str, Any]:
        """Run all middleware after_agent hooks in order."""
        for mw in self.middleware:
            state = mw.after_agent(state)
        return state

    def finalize(self) -> dict[str, Any]:
        """Finalize the agent run, applying after_agent middleware."""
        self._state = self.apply_after_agent(self._state)
        return self._state

    def get_audit_log(self) -> list[dict[str, Any]]:
        """Convenience: get audit log from the first compliance middleware."""
        for mw in self.middleware:
            if hasattr(mw, "get_audit_log"):
                return mw.get_audit_log()
        return []


def create_agent(
    *,
    model: str = "gpt-4.1",
    tools: list[Any] | None = None,
    middleware: list[AgentMiddleware] | None = None,
) -> MiddlewareAgent:
    """Create a middleware-aware agent.

    Args:
        model: The model name to use.
        tools: List of tools available to the agent.
        middleware: List of AgentMiddleware instances.

    Returns:
        A MiddlewareAgent with the given configuration.
    """
    return MiddlewareAgent(model=model, tools=tools, middleware=middleware)
