from __future__ import annotations

import asyncio
import importlib
import logging
import os
import threading
import time
import uuid
from datetime import datetime, timezone
from dataclasses import dataclass
from typing import Any, Dict, Optional, TypeVar

import httpx

from .adapters import LLMAdapter, OpenAIAdapter
from .errors import SecurityBlockError
from .guard import (
    ActionTaken,
    CapabilityFlags,
    LocalPolicyBlockError,
    LocalSecurityEnforcer,
    StrictPiiStreamingError,
    ViolationType,
    build_security_policy_violation_metadata,
    ensure_capability_flags_async,
    ensure_capability_flags_sync,
    get_cached_capability_flags,
)
from .security import get_injection_scanner
from .types import GuardParams, GuardResponse, LogParams
from .pii import PIIManager

logger = logging.getLogger("agentid")

TClient = TypeVar("TClient")

_OPENAI_VERSION_CHECKED = False
_OPENAI_VERSION_LOCK = threading.Lock()
AGENTID_SDK_VERSION_HEADER = "py-1.1.0"


def _assert_openai_v1_or_higher() -> None:
    global _OPENAI_VERSION_CHECKED

    if _OPENAI_VERSION_CHECKED:
        return

    with _OPENAI_VERSION_LOCK:
        if _OPENAI_VERSION_CHECKED:
            return

        try:
            openai_module = importlib.import_module("openai")
        except Exception:
            # If openai is not installed yet, we defer to caller-side import/runtime behavior.
            _OPENAI_VERSION_CHECKED = True
            return

        version = str(getattr(openai_module, "__version__", "")).strip()
        if version.startswith("0."):
            raise RuntimeError(
                "AgentID requires openai v1.0.0+. Run 'pip install -U openai'."
            )

        _OPENAI_VERSION_CHECKED = True


def _normalize_base_url(base_url: str) -> str:
    return base_url.rstrip("/")


def _normalize_timeout_seconds(value: float, default: float) -> float:
    try:
        timeout = float(value)
    except (TypeError, ValueError):
        return default
    if timeout <= 0:
        return default
    return timeout


def _resolve_initial_api_key(api_key: Optional[str]) -> str:
    explicit = api_key.strip() if isinstance(api_key, str) else ""
    from_env = (os.getenv("AGENTID_API_KEY") or "").strip()
    resolved = explicit or from_env
    if not resolved:
        raise ValueError("AgentID API key missing. Pass api_key or set AGENTID_API_KEY.")
    return resolved


def _is_infrastructure_guard_reason(reason: Optional[str]) -> bool:
    return reason in {
        "system_failure",
        "system_failure_db_unavailable",
        "logging_failed",
        "server_error",
        "guard_unreachable",
        "api_key_pepper_missing",
        "encryption_key_missing",
    }


def _utc_now_iso() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


def _coerce_uuid_str(value: Any) -> Optional[str]:
    if not isinstance(value, str):
        return None
    candidate = value.strip()
    if not candidate:
        return None
    try:
        return str(uuid.UUID(candidate))
    except Exception:
        return None


def _headers(api_key: str) -> Dict[str, str]:
    return {
        "Content-Type": "application/json",
        "x-agentid-api-key": api_key,
        "X-AgentID-SDK-Version": AGENTID_SDK_VERSION_HEADER,
    }


def _build_client_capabilities(
    *,
    pii_masking_enabled: bool,
    framework: str = "python_sdk",
    has_feedback_handler: bool = False,
) -> Dict[str, Any]:
    return {
        "capabilities": {
            "has_feedback_handler": bool(has_feedback_handler),
            "pii_masking_enabled": bool(pii_masking_enabled),
            "framework": framework.strip()[:64] if framework else "python_sdk",
        }
    }


@dataclass(frozen=True)
class _ClientConfig:
    api_key: str
    base_url: str
    guard_timeout_s: float
    ingest_timeout_s: float
    strict_mode: bool
    pii_masking: bool
    check_injection: bool
    ai_scan_enabled: bool
    store_pii: bool


def _mask_openai_messages(messages: Any, masked_text: str) -> Any:
    if not isinstance(messages, list):
        return messages

    last_user_idx: Optional[int] = None
    for idx, msg in enumerate(messages):
        if isinstance(msg, dict) and msg.get("role") == "user":
            last_user_idx = idx

    if last_user_idx is None:
        return messages

    new_messages = list(messages)
    last = new_messages[last_user_idx]
    if isinstance(last, dict):
        new_last = dict(last)
        new_last["content"] = masked_text
        new_messages[last_user_idx] = new_last
    return new_messages


def _with_masked_openai_request(
    args: tuple[Any, ...], kwargs: Dict[str, Any], masked_text: str
) -> tuple[tuple[Any, ...], Dict[str, Any]]:
    if args and isinstance(args[0], dict):
        req = dict(args[0])
        req["messages"] = _mask_openai_messages(req.get("messages"), masked_text)
        return (req,) + args[1:], kwargs

    if "messages" in kwargs:
        new_kwargs = dict(kwargs)
        new_kwargs["messages"] = _mask_openai_messages(kwargs.get("messages"), masked_text)
        return args, new_kwargs

    return args, kwargs


def _deanonymize_openai_response(
    response: Any, pii_manager: PIIManager, mapping: Dict[str, str]
) -> None:
    if not mapping:
        return

    try:
        choices = getattr(response, "choices", None)
        if not isinstance(choices, list):
            return

        for choice in choices:
            if isinstance(choice, dict):
                message = choice.get("message")
                if isinstance(message, dict) and isinstance(message.get("content"), str):
                    message["content"] = pii_manager.deanonymize(message["content"], mapping)
                delta = choice.get("delta")
                if isinstance(delta, dict) and isinstance(delta.get("content"), str):
                    delta["content"] = pii_manager.deanonymize(delta["content"], mapping)
                continue

            message = getattr(choice, "message", None)
            if message is not None:
                content = getattr(message, "content", None)
                if isinstance(content, str):
                    setattr(message, "content", pii_manager.deanonymize(content, mapping))

            delta = getattr(choice, "delta", None)
            if delta is not None:
                delta_content = getattr(delta, "content", None)
                if isinstance(delta_content, str):
                    setattr(delta, "content", pii_manager.deanonymize(delta_content, mapping))
    except Exception:
        # Fail-safe: never break completion flow due to post-processing.
        return


class _CompletionsCreateSyncWrapper:
    def __init__(
        self,
        agent: "AgentID",
        system_id: str,
        user_id: Optional[str],
        completions_obj: Any,
        adapter: LLMAdapter,
        api_key_override: Optional[str] = None,
    ):
        self._agent = agent
        self._system_id = system_id
        self._user_id = user_id
        self._completions = completions_obj
        self._adapter = adapter
        self._api_key_override = api_key_override.strip() if api_key_override else None

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

    def create(self, *args: Any, **kwargs: Any) -> Any:
        req = kwargs if not (args and isinstance(args[0], dict)) else args[0]
        stream = bool(req.get("stream")) if isinstance(req, dict) else bool(kwargs.get("stream"))
        effective_api_key = self._agent._resolve_api_key(self._api_key_override)

        user_text = self._adapter.extract_input(args, kwargs)
        capability_flags = self._agent._get_cached_capability_flags(api_key=effective_api_key)

        # 1) SECURITY (raw input): local prompt-injection scan before any masking.
        if (
            user_text is not None
            and user_text != ""
            and self._agent._config.check_injection
        ):
            get_injection_scanner().scan(
                prompt=user_text,
                api_key=effective_api_key,
                base_url=self._agent._config.base_url,
                ai_scan_enabled=self._agent._config.ai_scan_enabled,
                store_pii=self._agent._config.store_pii,
                pii_manager=self._agent._pii,
                source="python_sdk",
            )

        masked_text = user_text
        pii_mapping: Dict[str, str] = {}
        should_deanonymize = False
        if user_text:
            masked_text, capability_flags = self._agent._apply_local_security(
                system_id=self._system_id,
                input_text=user_text,
                stream=stream,
                api_key=effective_api_key,
            )
            if masked_text != user_text:
                args, kwargs = _with_masked_openai_request(args, kwargs, masked_text)

            if capability_flags.block_pii_leakage is False and self._agent._config.pii_masking:
                if stream:
                    logger.warning(
                        "AgentID: PII masking is disabled for streaming responses."
                    )
                else:
                    masked_text, pii_mapping = self._agent._pii.anonymize(user_text)
                    should_deanonymize = bool(pii_mapping)
                    args, kwargs = _with_masked_openai_request(args, kwargs, masked_text)

        if user_text is None:
            raise ValueError(
                "AgentID: No user message found. Security guard requires string input."
            )

        request_client_event_id = str(uuid.uuid4())
        requested_model = self._adapter.get_model_name(args, kwargs)
        verdict = self._agent.guard(
            input=masked_text or "",
            system_id=self._system_id,
            model=requested_model,
            user_id=self._user_id,
            client_event_id=request_client_event_id,
            api_key=effective_api_key,
            client_capabilities=self._agent._default_client_capabilities(
                framework="openai",
            ),
        )
        if not verdict.get("allowed", True):
            raise SecurityBlockError(verdict.get("reason"))
        canonical_client_event_id = (
            _coerce_uuid_str(verdict.get("client_event_id")) or request_client_event_id
        )
        guard_event_id = verdict.get("guard_event_id")
        guard_event_id_str = guard_event_id if isinstance(guard_event_id, str) and guard_event_id else None

        transformed_input = verdict.get("transformed_input")
        if isinstance(transformed_input, str) and transformed_input:
            if transformed_input != (masked_text or ""):
                masked_text = transformed_input
                args, kwargs = _with_masked_openai_request(args, kwargs, transformed_input)

        if stream:
            logger.warning(
                "AgentID: Automatic logging for streaming responses is not yet supported."
            )
            return self._completions.create(*args, **kwargs)

        start = time.perf_counter()
        resp = self._completions.create(*args, **kwargs)
        duration_ms = int((time.perf_counter() - start) * 1000)

        if masked_text is not None:
            usage = self._adapter.get_token_usage(resp)
            model = self._adapter.get_model_name(args, kwargs, resp) or "unknown"
            output_text = self._adapter.extract_output(resp)
            self._agent.log(
                system_id=self._system_id,
                user_id=self._user_id,
                input=masked_text,
                output=output_text,
                event_id=canonical_client_event_id,
                model=str(model),
                usage=usage,
                metadata={
                    "client_event_id": canonical_client_event_id,
                    "guard_event_id": guard_event_id_str,
                },
                latency=duration_ms,
                api_key=effective_api_key,
                client_capabilities=self._agent._default_client_capabilities(
                    framework="openai",
                ),
            )

        if capability_flags.block_pii_leakage is False and should_deanonymize:
            _deanonymize_openai_response(resp, self._agent._pii, pii_mapping)

        return resp


class _ChatSyncWrapper:
    def __init__(
        self,
        agent: "AgentID",
        system_id: str,
        user_id: Optional[str],
        chat_obj: Any,
        adapter: LLMAdapter,
        api_key_override: Optional[str] = None,
    ):
        self._chat = chat_obj
        self._completions_wrapper = _CompletionsCreateSyncWrapper(
            agent=agent,
            system_id=system_id,
            user_id=user_id,
            completions_obj=getattr(chat_obj, "completions", None),
            adapter=adapter,
            api_key_override=api_key_override,
        )

    def __getattr__(self, name: str) -> Any:
        if name == "completions":
            return self._completions_wrapper
        return getattr(self._chat, name)


class SyncOpenAIWrapper:
    """
    Transparent proxy for openai.OpenAI that intercepts chat.completions.create.
    """

    def __init__(
        self,
        agent: "AgentID",
        client: Any,
        system_id: str,
        user_id: Optional[str],
        adapter: LLMAdapter,
        api_key_override: Optional[str] = None,
    ):
        self._agent = agent
        self._client = client
        self._system_id = system_id
        self._user_id = user_id
        self._adapter = adapter
        self._api_key_override = api_key_override.strip() if api_key_override else None

    def __getattr__(self, name: str) -> Any:
        if name == "chat":
            return _ChatSyncWrapper(
                self._agent,
                self._system_id,
                self._user_id,
                getattr(self._client, "chat"),
                self._adapter,
                api_key_override=self._api_key_override,
            )
        return getattr(self._client, name)


class _CompletionsCreateAsyncWrapper:
    def __init__(
        self,
        agent: "AsyncAgentID",
        system_id: str,
        user_id: Optional[str],
        completions_obj: Any,
        adapter: LLMAdapter,
        api_key_override: Optional[str] = None,
    ):
        self._agent = agent
        self._system_id = system_id
        self._user_id = user_id
        self._completions = completions_obj
        self._adapter = adapter
        self._api_key_override = api_key_override.strip() if api_key_override else None

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

    async def create(self, *args: Any, **kwargs: Any) -> Any:
        req = kwargs if not (args and isinstance(args[0], dict)) else args[0]
        stream = bool(req.get("stream")) if isinstance(req, dict) else bool(kwargs.get("stream"))
        effective_api_key = self._agent._resolve_api_key(self._api_key_override)

        user_text = self._adapter.extract_input(args, kwargs)
        capability_flags = self._agent._get_cached_capability_flags(api_key=effective_api_key)

        # 1) SECURITY (raw input): local prompt-injection scan before any masking.
        if (
            user_text is not None
            and user_text != ""
            and self._agent._config.check_injection
        ):
            get_injection_scanner().scan(
                prompt=user_text,
                api_key=effective_api_key,
                base_url=self._agent._config.base_url,
                ai_scan_enabled=self._agent._config.ai_scan_enabled,
                store_pii=self._agent._config.store_pii,
                pii_manager=self._agent._pii,
                source="python_sdk",
            )

        masked_text = user_text
        pii_mapping: Dict[str, str] = {}
        should_deanonymize = False
        if user_text:
            masked_text, capability_flags = await self._agent._apply_local_security_async(
                system_id=self._system_id,
                input_text=user_text,
                stream=stream,
                api_key=effective_api_key,
            )
            if masked_text != user_text:
                args, kwargs = _with_masked_openai_request(args, kwargs, masked_text)

            if capability_flags.block_pii_leakage is False and self._agent._config.pii_masking:
                if stream:
                    logger.warning(
                        "AgentID: PII masking is disabled for streaming responses."
                    )
                else:
                    masked_text, pii_mapping = self._agent._pii.anonymize(user_text)
                    should_deanonymize = bool(pii_mapping)
                    args, kwargs = _with_masked_openai_request(args, kwargs, masked_text)

        if user_text is None:
            raise ValueError(
                "AgentID: No user message found. Security guard requires string input."
            )

        request_client_event_id = str(uuid.uuid4())
        requested_model = self._adapter.get_model_name(args, kwargs)
        verdict = await self._agent.guard(
            input=masked_text or "",
            system_id=self._system_id,
            model=requested_model,
            user_id=self._user_id,
            client_event_id=request_client_event_id,
            api_key=effective_api_key,
            client_capabilities=self._agent._default_client_capabilities(
                framework="openai",
            ),
        )
        if not verdict.get("allowed", True):
            raise SecurityBlockError(verdict.get("reason"))
        canonical_client_event_id = (
            _coerce_uuid_str(verdict.get("client_event_id")) or request_client_event_id
        )
        guard_event_id = verdict.get("guard_event_id")
        guard_event_id_str = guard_event_id if isinstance(guard_event_id, str) and guard_event_id else None

        transformed_input = verdict.get("transformed_input")
        if isinstance(transformed_input, str) and transformed_input:
            if transformed_input != (masked_text or ""):
                masked_text = transformed_input
                args, kwargs = _with_masked_openai_request(args, kwargs, transformed_input)

        if stream:
            logger.warning(
                "AgentID: Automatic logging for streaming responses is not yet supported."
            )
            return await self._completions.create(*args, **kwargs)

        start = time.perf_counter()
        resp = await self._completions.create(*args, **kwargs)
        duration_ms = int((time.perf_counter() - start) * 1000)

        if masked_text is not None:
            usage = self._adapter.get_token_usage(resp)
            model = self._adapter.get_model_name(args, kwargs, resp) or "unknown"
            output_text = self._adapter.extract_output(resp)
            self._agent.log(
                system_id=self._system_id,
                user_id=self._user_id,
                input=masked_text,
                output=output_text,
                event_id=canonical_client_event_id,
                model=str(model),
                usage=usage,
                metadata={
                    "client_event_id": canonical_client_event_id,
                    "guard_event_id": guard_event_id_str,
                },
                latency=duration_ms,
                api_key=effective_api_key,
                client_capabilities=self._agent._default_client_capabilities(
                    framework="openai",
                ),
            )

        if capability_flags.block_pii_leakage is False and should_deanonymize:
            _deanonymize_openai_response(resp, self._agent._pii, pii_mapping)

        return resp


class _ChatAsyncWrapper:
    def __init__(
        self,
        agent: "AsyncAgentID",
        system_id: str,
        user_id: Optional[str],
        chat_obj: Any,
        adapter: LLMAdapter,
        api_key_override: Optional[str] = None,
    ):
        self._chat = chat_obj
        self._completions_wrapper = _CompletionsCreateAsyncWrapper(
            agent=agent,
            system_id=system_id,
            user_id=user_id,
            completions_obj=getattr(chat_obj, "completions", None),
            adapter=adapter,
            api_key_override=api_key_override,
        )

    def __getattr__(self, name: str) -> Any:
        if name == "completions":
            return self._completions_wrapper
        return getattr(self._chat, name)


class AsyncOpenAIWrapper:
    """
    Transparent proxy for openai.AsyncOpenAI that intercepts chat.completions.create.
    """

    def __init__(
        self,
        agent: "AsyncAgentID",
        client: Any,
        system_id: str,
        user_id: Optional[str],
        adapter: LLMAdapter,
        api_key_override: Optional[str] = None,
    ):
        self._agent = agent
        self._client = client
        self._system_id = system_id
        self._user_id = user_id
        self._adapter = adapter
        self._api_key_override = api_key_override.strip() if api_key_override else None

    def __getattr__(self, name: str) -> Any:
        if name == "chat":
            return _ChatAsyncWrapper(
                self._agent,
                self._system_id,
                self._user_id,
                getattr(self._client, "chat"),
                self._adapter,
                api_key_override=self._api_key_override,
            )
        return getattr(self._client, name)


def _is_openai_client(obj: Any, expected_name: str) -> bool:
    cls = obj.__class__
    return cls.__name__ == expected_name and (cls.__module__ or "").startswith("openai")


class AgentID:
    """
    Synchronous AgentID client.

    - guard(): blocking call (default fail-open on connectivity failure)
    - log(): fire-and-forget telemetry via background thread
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: str = "https://app.getagentid.com/api/v1",
        pii_masking: bool = False,
        check_injection: bool = True,
        ai_scan_enabled: bool = True,
        store_pii: bool = False,
        strict_mode: bool = False,
        guard_timeout_s: float = 2.0,
        ingest_timeout_s: float = 5.0,
    ):
        resolved_api_key = _resolve_initial_api_key(api_key)
        self._config = _ClientConfig(
            api_key=resolved_api_key,
            base_url=_normalize_base_url(base_url),
            guard_timeout_s=_normalize_timeout_seconds(guard_timeout_s, 2.0),
            ingest_timeout_s=_normalize_timeout_seconds(ingest_timeout_s, 5.0),
            strict_mode=bool(strict_mode),
            pii_masking=bool(pii_masking),
            check_injection=bool(check_injection),
            ai_scan_enabled=bool(ai_scan_enabled),
            store_pii=bool(store_pii),
        )
        self._client = httpx.Client(
            base_url=self._config.base_url,
            headers=_headers(self._config.api_key),
        )
        self._closed = False
        self._pii = PIIManager()
        self._local_enforcer = LocalSecurityEnforcer(self._pii)

    def _resolve_api_key(self, api_key: Optional[str]) -> str:
        trimmed = api_key.strip() if isinstance(api_key, str) else ""
        return trimmed or self._config.api_key

    def _default_client_capabilities(
        self,
        *,
        framework: str = "python_sdk",
        has_feedback_handler: bool = False,
    ) -> Dict[str, Any]:
        return _build_client_capabilities(
            pii_masking_enabled=self._config.pii_masking,
            framework=framework,
            has_feedback_handler=has_feedback_handler,
        )

    def _log_security_policy_violation(
        self,
        *,
        system_id: str,
        violation_type: ViolationType,
        action_taken: ActionTaken,
        api_key: Optional[str] = None,
    ) -> None:
        metadata = build_security_policy_violation_metadata(
            system_id=system_id,
            violation_type=violation_type,
            action_taken=action_taken,
        )
        self.log(
            system_id=system_id,
            input="[REDACTED_SAMPLE]",
            output="",
            model="agentid.policy.enforcer",
            event_type="security_policy_violation",
            severity="high",
            metadata=metadata,
            api_key=api_key,
        )

    def _get_cached_capability_flags(self, *, api_key: Optional[str] = None) -> CapabilityFlags:
        effective_api_key = self._resolve_api_key(api_key)
        return get_cached_capability_flags(
            api_key=effective_api_key,
            base_url=self._config.base_url,
        )

    def _resolve_effective_strict_mode(self, *, api_key: str) -> bool:
        if self._config.strict_mode:
            return True
        flags = ensure_capability_flags_sync(
            api_key=api_key,
            base_url=self._config.base_url,
            timeout_seconds=self._config.guard_timeout_s,
            ttl_seconds=300.0,
        )
        return bool(flags.strict_security_mode or flags.failure_mode == "fail_close")

    def _apply_local_security(
        self,
        *,
        system_id: str,
        input_text: str,
        stream: bool,
        api_key: Optional[str] = None,
    ) -> tuple[str, CapabilityFlags]:
        effective_api_key = self._resolve_api_key(api_key)
        capability_flags = ensure_capability_flags_sync(
            api_key=effective_api_key,
            base_url=self._config.base_url,
            timeout_seconds=2.0,
            ttl_seconds=300.0,
        )

        try:
            result = self._local_enforcer.enforce(
                input_text=input_text,
                stream=stream,
                flags=capability_flags,
            )
        except (LocalPolicyBlockError, StrictPiiStreamingError) as exc:
            self._log_security_policy_violation(
                system_id=system_id,
                violation_type=exc.violation_type,
                action_taken=exc.action_taken,
                api_key=effective_api_key,
            )
            raise

        for event in result.events:
            self._log_security_policy_violation(
                system_id=system_id,
                violation_type=event.violation_type,
                action_taken=event.action_taken,
                api_key=effective_api_key,
            )

        return result.sanitized_input, capability_flags

    def close(self) -> None:
        if self._closed:
            return
        self._client.close()
        self._closed = True

    def __enter__(self) -> "AgentID":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.close()

    def wrap_openai(
        self,
        client: TClient,
        system_id: str,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
    ) -> TClient:
        _assert_openai_v1_or_higher()
        return self.wrap_client(
            client,
            system_id=system_id,
            user_id=user_id,
            adapter=OpenAIAdapter(),
            api_key=api_key,
        )

    def wrap_client(
        self,
        client: TClient,
        system_id: str,
        adapter: LLMAdapter,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
    ) -> TClient:
        """
        Generic wrapper entrypoint for any supported provider adapter.
        """
        effective_api_key = self._resolve_api_key(api_key)
        if _is_openai_client(client, "OpenAI"):
            _assert_openai_v1_or_higher()
            return SyncOpenAIWrapper(
                self,
                client,
                system_id,
                user_id,
                adapter,
                api_key_override=effective_api_key,
            )  # type: ignore[return-value]
        if _is_openai_client(client, "AsyncOpenAI"):
            _assert_openai_v1_or_higher()
            async_agent = AsyncAgentID(
                self._config.api_key,
                self._config.base_url,
                pii_masking=self._config.pii_masking,
                check_injection=self._config.check_injection,
                ai_scan_enabled=self._config.ai_scan_enabled,
                store_pii=self._config.store_pii,
                strict_mode=self._config.strict_mode,
                guard_timeout_s=self._config.guard_timeout_s,
                ingest_timeout_s=self._config.ingest_timeout_s,
            )
            return AsyncOpenAIWrapper(
                async_agent,
                client,
                system_id,
                user_id,
                adapter,
                api_key_override=effective_api_key,
            )  # type: ignore[return-value]
        raise TypeError("wrap_client() currently supports OpenAI clients only")

    def guard(
        self,
        input: str,
        system_id: str,
        model: Optional[str] = None,
        user_id: Optional[str] = None,
        client_event_id: Optional[str] = None,
        api_key: Optional[str] = None,
        client_capabilities: Optional[Dict[str, Any]] = None,
    ) -> GuardResponse:
        params: GuardParams = {"input": input, "system_id": system_id}
        if isinstance(model, str) and model.strip():
            params["model"] = model.strip()
        if user_id:
            params["user_id"] = user_id
        coerced_client_event_id = _coerce_uuid_str(client_event_id)
        if coerced_client_event_id is not None:
            params["client_event_id"] = coerced_client_event_id
        params["client_capabilities"] = (
            client_capabilities
            if isinstance(client_capabilities, dict)
            else self._default_client_capabilities(framework="python_sdk")
        )
        effective_api_key = self._resolve_api_key(api_key)
        effective_strict_mode = self._resolve_effective_strict_mode(api_key=effective_api_key)
        use_default_client = effective_api_key == self._config.api_key

        try:
            if use_default_client:
                res = self._client.post(
                    "/guard",
                    json=params,
                    timeout=httpx.Timeout(self._config.guard_timeout_s),
                )
            else:
                with httpx.Client(
                    base_url=self._config.base_url,
                    headers=_headers(effective_api_key),
                ) as override_client:
                    res = override_client.post(
                        "/guard",
                        json=params,
                        timeout=httpx.Timeout(self._config.guard_timeout_s),
                    )
            data = res.json()
            if isinstance(data, dict) and isinstance(data.get("allowed"), bool):
                infrastructure_failure = (
                    data.get("allowed") is False
                    and (
                        _is_infrastructure_guard_reason(data.get("reason"))
                        or (not data.get("reason") and res.status_code >= 500)
                    )
                )
                if infrastructure_failure:
                    reason = str(data.get("reason") or f"http_{res.status_code}")
                    if effective_strict_mode:
                        logger.warning(
                            "[AgentID] Guard API infrastructure failure in strict mode (%s). Blocking request.",
                            reason,
                        )
                        return {"allowed": False, "reason": str(data.get("reason") or "network_error_strict_mode")}
                    logger.warning(
                        "[AgentID] Guard API infrastructure fallback in fail-open mode (%s).",
                        reason,
                    )
                    return {"allowed": True, "reason": "system_failure_fail_open"}
                return data  # type: ignore[return-value]

            if res.is_error:
                res.raise_for_status()

            raise ValueError("Invalid guard response")
        except Exception as exc:
            if isinstance(exc, httpx.TimeoutException):
                logger.warning("AgentID API Warning: Connection timeout exceeded.")
                if effective_strict_mode:
                    return {"allowed": False, "reason": "network_error_strict_mode"}
                return {"allowed": True, "reason": "timeout_fallback"}

            if isinstance(exc, httpx.HTTPStatusError) and exc.response.status_code >= 500:
                if effective_strict_mode:
                    return {"allowed": False, "reason": "server_error"}
                return {"allowed": True, "reason": "system_failure_fail_open"}

            logger.warning(
                "[AgentID] Guard check failed (%s): %s",
                "Strict mode active" if effective_strict_mode else "Fail-Open active",
                exc.__class__.__name__,
            )
            if effective_strict_mode:
                return {"allowed": False, "reason": "network_error_strict_mode"}
            return {"allowed": True, "reason": "guard_unreachable"}

    def log(
        self,
        system_id: str,
        input: str,
        output: str,
        event_id: Optional[str] = None,
        model: Optional[str] = None,
        usage: Optional[Dict[str, Any]] = None,
        tokens: Optional[Dict[str, Any]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        latency: Optional[int] = None,
        user_id: Optional[str] = None,
        event_type: Optional[str] = None,
        severity: Optional[str] = None,
        timestamp: Optional[str] = None,
        api_key: Optional[str] = None,
        client_capabilities: Optional[Dict[str, Any]] = None,
    ) -> None:
        merged_metadata: Dict[str, Any] = dict(metadata or {})
        merged_metadata.setdefault("agentid_base_url", self._config.base_url)
        payload_event_id = _coerce_uuid_str(event_id) or str(uuid.uuid4())
        payload: LogParams = {
            "event_id": payload_event_id,
            "system_id": system_id,
            "input": input,
            "output": output,
            # Server requires model; default to a safe placeholder.
            "model": model or "unknown",
            "timestamp": timestamp or _utc_now_iso(),
            "metadata": merged_metadata,
            "client_capabilities": (
                client_capabilities
                if isinstance(client_capabilities, dict)
                else self._default_client_capabilities(framework="python_sdk")
            ),
        }
        if usage is not None:
            payload["usage"] = usage
        if tokens is not None:
            payload["tokens"] = tokens
        if latency is not None:
            payload["latency"] = int(latency)
        if user_id is not None:
            payload["user_id"] = user_id
        if event_type is not None:
            payload["event_type"] = event_type
        if severity is not None:
            payload["severity"] = severity
        if timestamp is not None:
            payload["timestamp"] = timestamp
        effective_api_key = self._resolve_api_key(api_key)

        def _send() -> None:
            try:
                ensure_capability_flags_sync(
                    api_key=effective_api_key,
                    base_url=self._config.base_url,
                    timeout_seconds=self._config.guard_timeout_s,
                    ttl_seconds=300.0,
                )
                with httpx.Client(
                    base_url=self._config.base_url,
                    headers=_headers(effective_api_key),
                ) as client:
                    client.post(
                        "/ingest",
                        json=payload,
                        timeout=httpx.Timeout(self._config.ingest_timeout_s),
                    )
            except Exception as exc:
                logger.error("[AgentID] Log failed: %s", exc.__class__.__name__)

        t = threading.Thread(target=_send, daemon=True, name="agentid-ingest")
        t.start()


class AsyncAgentID:
    """
    Asynchronous AgentID client.

    - guard(): awaitable (default fail-open on connectivity failure)
    - log(): schedules a background task (fire-and-forget)
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: str = "https://app.getagentid.com/api/v1",
        pii_masking: bool = False,
        check_injection: bool = True,
        ai_scan_enabled: bool = True,
        store_pii: bool = False,
        strict_mode: bool = False,
        guard_timeout_s: float = 2.0,
        ingest_timeout_s: float = 5.0,
    ):
        resolved_api_key = _resolve_initial_api_key(api_key)
        self._config = _ClientConfig(
            api_key=resolved_api_key,
            base_url=_normalize_base_url(base_url),
            guard_timeout_s=_normalize_timeout_seconds(guard_timeout_s, 2.0),
            ingest_timeout_s=_normalize_timeout_seconds(ingest_timeout_s, 5.0),
            strict_mode=bool(strict_mode),
            pii_masking=bool(pii_masking),
            check_injection=bool(check_injection),
            ai_scan_enabled=bool(ai_scan_enabled),
            store_pii=bool(store_pii),
        )
        self._client: Optional[httpx.AsyncClient] = None
        self._closed = False
        self._pii = PIIManager()
        self._local_enforcer = LocalSecurityEnforcer(self._pii)

    def _resolve_api_key(self, api_key: Optional[str]) -> str:
        trimmed = api_key.strip() if isinstance(api_key, str) else ""
        return trimmed or self._config.api_key

    def _default_client_capabilities(
        self,
        *,
        framework: str = "python_sdk",
        has_feedback_handler: bool = False,
    ) -> Dict[str, Any]:
        return _build_client_capabilities(
            pii_masking_enabled=self._config.pii_masking,
            framework=framework,
            has_feedback_handler=has_feedback_handler,
        )

    def _log_security_policy_violation(
        self,
        *,
        system_id: str,
        violation_type: ViolationType,
        action_taken: ActionTaken,
        api_key: Optional[str] = None,
    ) -> None:
        metadata = build_security_policy_violation_metadata(
            system_id=system_id,
            violation_type=violation_type,
            action_taken=action_taken,
        )
        self.log(
            system_id=system_id,
            input="[REDACTED_SAMPLE]",
            output="",
            model="agentid.policy.enforcer",
            event_type="security_policy_violation",
            severity="high",
            metadata=metadata,
            api_key=api_key,
        )

    def _get_cached_capability_flags(self, *, api_key: Optional[str] = None) -> CapabilityFlags:
        effective_api_key = self._resolve_api_key(api_key)
        return get_cached_capability_flags(
            api_key=effective_api_key,
            base_url=self._config.base_url,
        )

    async def _resolve_effective_strict_mode(self, *, api_key: str) -> bool:
        if self._config.strict_mode:
            return True
        flags = await ensure_capability_flags_async(
            api_key=api_key,
            base_url=self._config.base_url,
            timeout_seconds=self._config.guard_timeout_s,
            ttl_seconds=300.0,
        )
        return bool(flags.strict_security_mode or flags.failure_mode == "fail_close")

    async def _apply_local_security_async(
        self,
        *,
        system_id: str,
        input_text: str,
        stream: bool,
        api_key: Optional[str] = None,
    ) -> tuple[str, CapabilityFlags]:
        effective_api_key = self._resolve_api_key(api_key)
        capability_flags = await ensure_capability_flags_async(
            api_key=effective_api_key,
            base_url=self._config.base_url,
            timeout_seconds=2.0,
            ttl_seconds=300.0,
        )

        try:
            result = self._local_enforcer.enforce(
                input_text=input_text,
                stream=stream,
                flags=capability_flags,
            )
        except (LocalPolicyBlockError, StrictPiiStreamingError) as exc:
            self._log_security_policy_violation(
                system_id=system_id,
                violation_type=exc.violation_type,
                action_taken=exc.action_taken,
                api_key=effective_api_key,
            )
            raise

        for event in result.events:
            self._log_security_policy_violation(
                system_id=system_id,
                violation_type=event.violation_type,
                action_taken=event.action_taken,
                api_key=effective_api_key,
            )

        return result.sanitized_input, capability_flags

    async def aclose(self) -> None:
        if self._closed:
            return
        if self._client is not None:
            await self._client.aclose()
        self._closed = True

    async def __aenter__(self) -> "AsyncAgentID":
        if self._client is None:
            self._client = httpx.AsyncClient(
                base_url=self._config.base_url,
                headers=_headers(self._config.api_key),
            )
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        await self.aclose()

    def wrap_openai(
        self,
        client: TClient,
        system_id: str,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
    ) -> TClient:
        _assert_openai_v1_or_higher()
        return self.wrap_client(
            client,
            system_id=system_id,
            user_id=user_id,
            adapter=OpenAIAdapter(),
            api_key=api_key,
        )

    def wrap_client(
        self,
        client: TClient,
        system_id: str,
        adapter: LLMAdapter,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
    ) -> TClient:
        effective_api_key = self._resolve_api_key(api_key)
        if _is_openai_client(client, "AsyncOpenAI"):
            _assert_openai_v1_or_higher()
            return AsyncOpenAIWrapper(
                self,
                client,
                system_id,
                user_id,
                adapter,
                api_key_override=effective_api_key,
            )  # type: ignore[return-value]
        if _is_openai_client(client, "OpenAI"):
            _assert_openai_v1_or_higher()
            sync_agent = AgentID(
                self._config.api_key,
                self._config.base_url,
                pii_masking=self._config.pii_masking,
                check_injection=self._config.check_injection,
                ai_scan_enabled=self._config.ai_scan_enabled,
                store_pii=self._config.store_pii,
                strict_mode=self._config.strict_mode,
                guard_timeout_s=self._config.guard_timeout_s,
                ingest_timeout_s=self._config.ingest_timeout_s,
            )
            return SyncOpenAIWrapper(
                sync_agent,
                client,
                system_id,
                user_id,
                adapter,
                api_key_override=effective_api_key,
            )  # type: ignore[return-value]
        raise TypeError("wrap_client() currently supports OpenAI clients only")

    async def guard(
        self,
        input: str,
        system_id: str,
        model: Optional[str] = None,
        user_id: Optional[str] = None,
        client_event_id: Optional[str] = None,
        api_key: Optional[str] = None,
        client_capabilities: Optional[Dict[str, Any]] = None,
    ) -> GuardResponse:
        params: GuardParams = {"input": input, "system_id": system_id}
        if isinstance(model, str) and model.strip():
            params["model"] = model.strip()
        if user_id:
            params["user_id"] = user_id
        coerced_client_event_id = _coerce_uuid_str(client_event_id)
        if coerced_client_event_id is not None:
            params["client_event_id"] = coerced_client_event_id
        params["client_capabilities"] = (
            client_capabilities
            if isinstance(client_capabilities, dict)
            else self._default_client_capabilities(framework="python_sdk")
        )

        effective_api_key = self._resolve_api_key(api_key)
        effective_strict_mode = await self._resolve_effective_strict_mode(api_key=effective_api_key)
        client = self._client
        close_after = False
        if client is None or effective_api_key != self._config.api_key:
            client = httpx.AsyncClient(
                base_url=self._config.base_url,
                headers=_headers(effective_api_key),
            )
            close_after = True

        try:
            res = await client.post(
                "/guard",
                json=params,
                timeout=httpx.Timeout(self._config.guard_timeout_s),
            )
            data = res.json()
            if isinstance(data, dict) and isinstance(data.get("allowed"), bool):
                infrastructure_failure = (
                    data.get("allowed") is False
                    and (
                        _is_infrastructure_guard_reason(data.get("reason"))
                        or (not data.get("reason") and res.status_code >= 500)
                    )
                )
                if infrastructure_failure:
                    reason = str(data.get("reason") or f"http_{res.status_code}")
                    if effective_strict_mode:
                        logger.warning(
                            "[AgentID] Guard API infrastructure failure in strict mode (%s). Blocking request.",
                            reason,
                        )
                        return {"allowed": False, "reason": str(data.get("reason") or "network_error_strict_mode")}
                    logger.warning(
                        "[AgentID] Guard API infrastructure fallback in fail-open mode (%s).",
                        reason,
                    )
                    return {"allowed": True, "reason": "system_failure_fail_open"}
                return data  # type: ignore[return-value]

            if res.is_error:
                res.raise_for_status()

            raise ValueError("Invalid guard response")
        except Exception as exc:
            if isinstance(exc, httpx.TimeoutException):
                logger.warning("AgentID API Warning: Connection timeout exceeded.")
                if effective_strict_mode:
                    return {"allowed": False, "reason": "network_error_strict_mode"}
                return {"allowed": True, "reason": "timeout_fallback"}

            if isinstance(exc, httpx.HTTPStatusError) and exc.response.status_code >= 500:
                if effective_strict_mode:
                    return {"allowed": False, "reason": "server_error"}
                return {"allowed": True, "reason": "system_failure_fail_open"}

            logger.warning(
                "[AgentID] Guard check failed (%s): %s",
                "Strict mode active" if effective_strict_mode else "Fail-Open active",
                exc.__class__.__name__,
            )
            if effective_strict_mode:
                return {"allowed": False, "reason": "network_error_strict_mode"}
            return {"allowed": True, "reason": "guard_unreachable"}
        finally:
            if close_after:
                await client.aclose()

    def log(
        self,
        system_id: str,
        input: str,
        output: str,
        event_id: Optional[str] = None,
        model: Optional[str] = None,
        usage: Optional[Dict[str, Any]] = None,
        tokens: Optional[Dict[str, Any]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        latency: Optional[int] = None,
        user_id: Optional[str] = None,
        event_type: Optional[str] = None,
        severity: Optional[str] = None,
        timestamp: Optional[str] = None,
        api_key: Optional[str] = None,
        client_capabilities: Optional[Dict[str, Any]] = None,
    ) -> None:
        merged_metadata: Dict[str, Any] = dict(metadata or {})
        merged_metadata.setdefault("agentid_base_url", self._config.base_url)
        payload_event_id = _coerce_uuid_str(event_id) or str(uuid.uuid4())
        payload: LogParams = {
            "event_id": payload_event_id,
            "system_id": system_id,
            "input": input,
            "output": output,
            "model": model or "unknown",
            "timestamp": timestamp or _utc_now_iso(),
            "metadata": merged_metadata,
            "client_capabilities": (
                client_capabilities
                if isinstance(client_capabilities, dict)
                else self._default_client_capabilities(framework="python_sdk")
            ),
        }
        if usage is not None:
            payload["usage"] = usage
        if tokens is not None:
            payload["tokens"] = tokens
        if latency is not None:
            payload["latency"] = int(latency)
        if user_id is not None:
            payload["user_id"] = user_id
        if event_type is not None:
            payload["event_type"] = event_type
        if severity is not None:
            payload["severity"] = severity
        if timestamp is not None:
            payload["timestamp"] = timestamp
        effective_api_key = self._resolve_api_key(api_key)

        async def _send() -> None:
            client = self._client
            close_after = False
            if client is None or effective_api_key != self._config.api_key:
                client = httpx.AsyncClient(
                    base_url=self._config.base_url,
                    headers=_headers(effective_api_key),
                )
                close_after = True

            try:
                await ensure_capability_flags_async(
                    api_key=effective_api_key,
                    base_url=self._config.base_url,
                    timeout_seconds=self._config.guard_timeout_s,
                    ttl_seconds=300.0,
                )
                await client.post(
                    "/ingest",
                    json=payload,
                    timeout=httpx.Timeout(self._config.ingest_timeout_s),
                )
            except Exception as exc:
                logger.error("[AgentID] Log failed: %s", exc.__class__.__name__)
            finally:
                if close_after:
                    await client.aclose()

        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            logger.warning("[AgentID] Async log skipped: no running event loop")
            return

        loop.create_task(_send())
