from __future__ import annotations

from contextlib import AbstractContextManager, asynccontextmanager, contextmanager
from dataclasses import dataclass, replace
import inspect
from typing import Any, Iterable, Sequence, cast
from collections.abc import Iterator
from uuid import uuid4

from pydantic_ai import Agent
from pydantic_ai.toolsets import AbstractToolset, CombinedToolset
from pydantic_ai.toolsets._dynamic import DynamicToolset
from pydantic_ai.toolsets.function import FunctionToolset
from pydantic_ai.toolsets.wrapper import WrapperToolset

from ..config import RegulatorConfig, get_config
from ..context import context, get_current_context, restore_tags_from_messages
from ..coverage import emit_coverage
from ..decorators import rules_for_tool
from ..governance import evaluate_policies, is_visible, record_policy_event, record_visibility_event
from ..types import PolicyRequest
from ..types import ToolRules
from .mcp import RegulatedMCPAdapter
from .toolset import RegulatedToolset


@dataclass(frozen=True)
class _PreparedToolsets:
    override_toolsets: list[AbstractToolset[Any]] | None
    override_tools: Sequence[Any] | None
    run_toolsets: Sequence[AbstractToolset[Any]] | None
    lookup_toolsets: list[Any]


class RegulatedAgent:
    def __init__(
        self,
        agent: Agent[Any, Any],
        *,
        agent_id: str | None = None,
        config: RegulatorConfig | None = None,
        wrap_agent_toolset: bool = True,
        govern_builtin_tools: bool = False,
    ) -> None:
        self._agent = agent
        self._agent_id = agent_id or getattr(agent, "name", None) or f"agent-{uuid4()}"
        self._config = config
        self._coverage_emitted = False
        self._wrap_agent_toolset = wrap_agent_toolset
        self._govern_builtin_tools = govern_builtin_tools

    @property
    def agent(self) -> Agent:
        return self._agent

    @property
    def agent_id(self) -> str:
        return self._agent_id

    async def run(self, *args: Any, run_id: str | None = None, **kwargs: Any) -> Any:
        message_history = kwargs.get("message_history")
        toolsets = kwargs.get("toolsets")
        builtin_tools = kwargs.get("builtin_tools")
        cfg = self._config or get_config()
        prepared = self._prepare_toolsets(toolsets)
        async with context(agent_id=self._agent_id, run_id=run_id, config=cfg) as ctx:
            lookup = self._build_rule_lookup(prepared.lookup_toolsets)
            if message_history is not None:
                restore_tags_from_messages(message_history, lookup, target=ctx)
            if not self._coverage_emitted:
                emit_coverage(
                    ctx,
                    self._collect_tool_rules(prepared.lookup_toolsets),
                    self._collect_mcp_summaries(prepared.lookup_toolsets),
                )
                self._coverage_emitted = True
            kwargs["toolsets"] = prepared.run_toolsets
            builtin_override: AbstractContextManager[None] | None = None
            if self._govern_builtin_tools:
                kwargs["builtin_tools"] = self._wrap_builtin_tools(builtin_tools)
                builtin_override = self._override_builtin_tools()
            with self._override_agent(prepared):
                if builtin_override is None:
                    return await self._agent.run(*args, **kwargs)
                with builtin_override:
                    return await self._agent.run(*args, **kwargs)

    def run_sync(self, *args: Any, run_id: str | None = None, **kwargs: Any) -> Any:
        message_history = kwargs.get("message_history")
        toolsets = kwargs.get("toolsets")
        builtin_tools = kwargs.get("builtin_tools")
        cfg = self._config or get_config()
        prepared = self._prepare_toolsets(toolsets)
        with context(agent_id=self._agent_id, run_id=run_id, config=cfg) as ctx:
            lookup = self._build_rule_lookup(prepared.lookup_toolsets)
            if message_history is not None:
                restore_tags_from_messages(message_history, lookup, target=ctx)
            if not self._coverage_emitted:
                emit_coverage(
                    ctx,
                    self._collect_tool_rules(prepared.lookup_toolsets),
                    self._collect_mcp_summaries(prepared.lookup_toolsets),
                )
                self._coverage_emitted = True
            kwargs["toolsets"] = prepared.run_toolsets
            builtin_override: AbstractContextManager[None] | None = None
            if self._govern_builtin_tools:
                kwargs["builtin_tools"] = self._wrap_builtin_tools(builtin_tools)
                builtin_override = self._override_builtin_tools()
            with self._override_agent(prepared):
                if builtin_override is None:
                    return self._agent.run_sync(*args, **kwargs)
                with builtin_override:
                    return self._agent.run_sync(*args, **kwargs)

    @asynccontextmanager
    async def run_stream(self, *args: Any, run_id: str | None = None, **kwargs: Any) -> Any:
        message_history = kwargs.get("message_history")
        toolsets = kwargs.get("toolsets")
        builtin_tools = kwargs.get("builtin_tools")
        cfg = self._config or get_config()
        prepared = self._prepare_toolsets(toolsets)
        async with context(agent_id=self._agent_id, run_id=run_id, config=cfg) as ctx:
            lookup = self._build_rule_lookup(prepared.lookup_toolsets)
            if message_history is not None:
                restore_tags_from_messages(message_history, lookup, target=ctx)
            if not self._coverage_emitted:
                emit_coverage(
                    ctx,
                    self._collect_tool_rules(prepared.lookup_toolsets),
                    self._collect_mcp_summaries(prepared.lookup_toolsets),
                )
                self._coverage_emitted = True
            kwargs["toolsets"] = prepared.run_toolsets
            builtin_override: AbstractContextManager[None] | None = None
            if self._govern_builtin_tools:
                kwargs["builtin_tools"] = self._wrap_builtin_tools(builtin_tools)
                builtin_override = self._override_builtin_tools()
            with self._override_agent(prepared):
                if builtin_override is None:
                    async with self._agent.run_stream(*args, **kwargs) as stream:
                        yield stream
                else:
                    with builtin_override:
                        async with self._agent.run_stream(*args, **kwargs) as stream:
                            yield stream

    def _prepare_toolsets(self, toolsets: Iterable[AbstractToolset[Any]] | None) -> _PreparedToolsets:
        agent_toolsets = list(self._agent.toolsets)
        function_toolset = agent_toolsets[0] if agent_toolsets else None
        user_toolsets = agent_toolsets[1:] if len(agent_toolsets) > 1 else []

        wrapped_user_toolsets = [self._wrap_toolset(toolset) for toolset in user_toolsets]
        wrapped_run_toolsets = [self._wrap_toolset(toolset) for toolset in toolsets or []]

        wrapped_function_toolset: AbstractToolset[Any] | None = None
        if self._wrap_agent_toolset and function_toolset is not None:
            wrapped_function_toolset = self._wrap_toolset(function_toolset)

        lookup_toolsets: list[Any] = []
        if wrapped_function_toolset is not None:
            lookup_toolsets.append(wrapped_function_toolset)
        lookup_toolsets.extend(wrapped_user_toolsets)
        lookup_toolsets.extend(wrapped_run_toolsets)

        if self._wrap_agent_toolset:
            override_toolsets: list[AbstractToolset[Any]] = []
            if wrapped_function_toolset is not None:
                override_toolsets.append(wrapped_function_toolset)
            override_toolsets.extend(wrapped_user_toolsets)
            override_toolsets.extend(wrapped_run_toolsets)
            return _PreparedToolsets(
                override_toolsets=override_toolsets,
                override_tools=[],
                run_toolsets=None,
                lookup_toolsets=lookup_toolsets,
            )

        if user_toolsets:
            return _PreparedToolsets(
                override_toolsets=[*wrapped_user_toolsets, *wrapped_run_toolsets],
                override_tools=None,
                run_toolsets=None,
                lookup_toolsets=lookup_toolsets,
            )

        if toolsets:
            return _PreparedToolsets(
                override_toolsets=None,
                override_tools=None,
                run_toolsets=wrapped_run_toolsets,
                lookup_toolsets=lookup_toolsets,
            )

        return _PreparedToolsets(
            override_toolsets=None,
            override_tools=None,
            run_toolsets=None,
            lookup_toolsets=lookup_toolsets,
        )

    def _wrap_toolset(self, toolset: AbstractToolset[Any]) -> AbstractToolset[Any]:
        if isinstance(toolset, (RegulatedToolset, RegulatedMCPAdapter)):
            return toolset
        if isinstance(toolset, DynamicToolset):
            return self._wrap_dynamic_toolset(toolset)
        if isinstance(toolset, FunctionToolset):
            return RegulatedToolset(toolset=toolset)
        if isinstance(toolset, CombinedToolset):
            return replace(toolset, toolsets=[self._wrap_toolset(item) for item in toolset.toolsets])
        if isinstance(toolset, WrapperToolset):
            return replace(toolset, wrapped=self._wrap_toolset(toolset.wrapped))
        return toolset

    def _wrap_dynamic_toolset(self, toolset: DynamicToolset[Any]) -> DynamicToolset[Any]:
        async def wrapped(ctx: Any) -> AbstractToolset[Any] | None:
            result = toolset.toolset_func(ctx)
            if inspect.isawaitable(result):
                result = await result
            if result is None:
                return None
            return self._wrap_toolset(result)

        return DynamicToolset(toolset_func=wrapped, per_run_step=toolset.per_run_step, id=toolset.id)

    @contextmanager
    def _override_agent(self, prepared: _PreparedToolsets) -> Iterator[None]:
        if prepared.override_toolsets is None and prepared.override_tools is None:
            yield
            return
        if prepared.override_toolsets is None:
            with self._agent.override(tools=cast(Any, prepared.override_tools)):
                yield
            return
        if prepared.override_tools is None:
            with self._agent.override(toolsets=prepared.override_toolsets):
                yield
            return
        with self._agent.override(
            toolsets=prepared.override_toolsets,
            tools=cast(Any, prepared.override_tools),
        ):
            yield

    @contextmanager
    def _override_builtin_tools(self) -> Iterator[None]:
        agent = self._agent
        if not hasattr(agent, "_builtin_tools"):
            yield
            return
        original = getattr(agent, "_builtin_tools")
        setattr(agent, "_builtin_tools", [])
        try:
            yield
        finally:
            setattr(agent, "_builtin_tools", original)

    def _collect_tool_rules(self, toolsets: Iterable[Any]) -> dict[str, ToolRules]:
        rules: dict[str, ToolRules] = {}
        for toolset in toolsets:
            if hasattr(toolset, "tool_rules_by_name"):
                rules.update(toolset.tool_rules_by_name())
        return rules

    def _collect_mcp_summaries(self, toolsets: Iterable[Any]) -> dict[str, dict[str, Any]]:
        summaries: dict[str, dict[str, Any]] = {}
        for toolset in toolsets:
            summary = getattr(toolset, "coverage_summary", None)
            if callable(summary):
                summaries[getattr(toolset, "label", "mcp")] = summary()
        return summaries

    def _build_rule_lookup(self, toolsets: Iterable[Any]) -> Any:
        tool_rules = self._collect_tool_rules(toolsets)
        toolsets_all = list(toolsets)

        def lookup(tool_name: str) -> set[Any] | None:
            rules = tool_rules.get(tool_name)
            if rules:
                return set(rules.activates)
            for toolset in toolsets_all:
                rules_for = getattr(toolset, "rules_for_tool", None)
                if callable(rules_for):
                    resolved = rules_for(tool_name)
                    if resolved:
                        return set(resolved.activates)
            return None

        return lookup

    def _wrap_builtin_tools(self, extra_tools: Iterable[Any] | None) -> list[Any]:
        builtin_tools = list(getattr(self._agent, "_builtin_tools", []))
        if extra_tools:
            builtin_tools.extend(extra_tools)
        return [self._wrap_builtin_tool(tool) for tool in builtin_tools]

    def _wrap_builtin_tool(self, tool: Any) -> Any:
        rules = None
        if callable(tool):
            rules = rules_for_tool(tool)
        else:
            rules = getattr(tool, "__rgltr__", None)

        async def wrapped(ctx: Any) -> Any | None:
            result = tool(ctx) if callable(tool) else tool
            if inspect.isawaitable(result):
                result = await result
            if result is None:
                return None
            if rules is None or rules.is_empty():
                if get_config().require_annotations:
                    raise ValueError(f"Builtin tool '{_builtin_name(result)}' lacks rgltr annotations")
                return result
            current = get_current_context()
            active_tags = set(current.active_tags) if current else set()
            visible = is_visible(rules, active_tags, get_config().boundaries)
            record_visibility_event(
                current,
                _builtin_name(result),
                rules,
                active_tags,
                visible,
                tool_kind="builtin",
            )
            if not visible:
                return None
            decision = await evaluate_policies(
                rules,
                PolicyRequest(
                    tool_name=_builtin_name(result),
                    args={},
                    active_tags=active_tags,
                    boundary=rules.boundary,
                    agent_id=current.agent_id if current else None,
                    run_id=current.run_id if current else "no_context",
                    run_context=ctx,
                ),
            )
            if decision.outcome != "allow":
                record_policy_event(
                    current,
                    _builtin_name(result),
                    rules,
                    decision,
                    tool_kind="builtin",
                    enforced=get_config().mode != "monitor",
                    tool_call_id=getattr(ctx, "tool_call_id", None),
                )
                if get_config().mode != "monitor":
                    return None
            return result

        return wrapped


def _builtin_name(tool: Any) -> str:
    if hasattr(tool, "unique_id"):
        return str(tool.unique_id)
    if hasattr(tool, "kind"):
        return str(tool.kind)
    if hasattr(tool, "label"):
        return str(tool.label)
    return str(tool.__class__.__name__)
