"""AutoPIL LlamaIndexGuard. See packages/core/autopil/llamaindex_guard.py for full docs."""

from __future__ import annotations
from typing import Any, Optional
from .guard import ContextGuard
from .models import SensitivityLevel


class LlamaIndexGuard(ContextGuard):
    """ContextGuard for LlamaIndex agents. Tags audit events with source_type='llamaindex'."""

    _source_type: str = "llamaindex"

    def wrap_query_engine(
        self,
        engine: Any,
        *,
        agent_role: str,
        user_id: str,
        source_id: str,
        sensitivity_level: SensitivityLevel = SensitivityLevel.MEDIUM,
        session_id: Optional[str] = None,
    ) -> Any:
        """Wrap a LlamaIndex QueryEngine so every .query() is policy-evaluated."""
        guard = self

        class _GuardedQueryEngine:
            def __init__(self):
                self._engine = engine
                self._protected_query = guard.protect(
                    agent_role=agent_role,
                    user_id=user_id,
                    source_id=source_id,
                    sensitivity_level=sensitivity_level,
                    session_id=session_id,
                )(engine.query)

            def query(self, query_str: str) -> Any:
                return self._protected_query(query_str)

            def __getattr__(self, name: str) -> Any:
                return getattr(self._engine, name)

        return _GuardedQueryEngine()

    def wrap_query_engine_async(
        self,
        engine: Any,
        *,
        agent_role: str,
        user_id: str,
        source_id: str,
        sensitivity_level: SensitivityLevel = SensitivityLevel.MEDIUM,
        session_id: Optional[str] = None,
    ) -> Any:
        """Async version of wrap_query_engine — wraps .aquery() with AutoPIL enforcement."""
        guard = self

        class _AsyncGuardedQueryEngine:
            def __init__(self):
                self._engine = engine
                self._protected_aquery = guard.protect_async(
                    agent_role=agent_role,
                    user_id=user_id,
                    source_id=source_id,
                    sensitivity_level=sensitivity_level,
                    session_id=session_id,
                )(self._aquery_impl)

            async def _aquery_impl(self, query_str: str) -> Any:
                if hasattr(self._engine, "aquery"):
                    return await self._engine.aquery(query_str)
                return self._engine.query(query_str)

            async def aquery(self, query_str: str) -> Any:
                return await self._protected_aquery(query_str)

            def query(self, query_str: str) -> Any:
                return guard.protect(
                    agent_role=agent_role,
                    user_id=user_id,
                    source_id=source_id,
                    sensitivity_level=sensitivity_level,
                    session_id=session_id,
                )(engine.query)(query_str)

            def __getattr__(self, name: str) -> Any:
                return getattr(self._engine, name)

        return _AsyncGuardedQueryEngine()
