from __future__ import annotations

import asyncio
import inspect
import logging
import time
import uuid
from dataclasses import dataclass
from typing import Any, Dict, Optional

from ._expected_languages import normalize_expected_languages
from .client import (
    _build_sdk_timing_metadata,
    _is_effective_strict_mode,
    _is_guard_failure_eligible_for_local_fallback,
)
from .errors import SecurityBlockError
from .guard import ensure_capability_flags_sync
from .security import get_injection_scanner

logger = logging.getLogger("agentid")


try:
    from langchain_core.callbacks import BaseCallbackHandler  # type: ignore
except Exception:  # pragma: no cover
    BaseCallbackHandler = object  # type: ignore


@dataclass
class _RunState:
    input_text: str
    started_at: float
    model: Optional[str] = None
    client_event_id: Optional[str] = None
    guard_event_id: Optional[str] = None
    transparency: Optional[Dict[str, Any]] = None
    sdk_config_fetch_ms: int = 0
    sdk_local_scan_ms: int = 0
    sdk_guard_ms: int = 0
    local_fallback_applied: bool = False
    local_fallback_reason: Optional[str] = None


def _safe_text(val: Any) -> str:
    if isinstance(val, str):
        return val
    return ""


def _extract_prompt_from_prompts(prompts: Any) -> Optional[str]:
    if isinstance(prompts, list) and prompts:
        last = prompts[-1]
        return _safe_text(last)
    return None


def _extract_prompt_from_messages(messages: Any) -> Optional[str]:
    # LangChain chat hooks can deliver BaseMessage[][] or similar.
    flat: list[Any] = []
    if isinstance(messages, list):
        for item in messages:
            if isinstance(item, list):
                flat.extend(item)
            else:
                flat.append(item)

    last = None
    for msg in flat:
        role = getattr(msg, "type", None) or getattr(msg, "role", None)
        if role in ("human", "user"):
            last = msg
        if isinstance(msg, dict) and msg.get("role") == "user":
            last = msg

    if last is None:
        return None

    if isinstance(last, dict):
        return _safe_text(last.get("content"))

    return _safe_text(getattr(last, "content", None))


def _set_prompt_in_prompts(prompts: Any, sanitized: str) -> bool:
    if not isinstance(prompts, list) or not prompts:
        return False
    prompts[-1] = sanitized
    return True


def _set_prompt_in_messages(messages: Any, sanitized: str) -> bool:
    if not isinstance(messages, list):
        return False

    flat: list[Any] = []
    for item in messages:
        if isinstance(item, list):
            flat.extend(item)
        else:
            flat.append(item)

    for msg in reversed(flat):
        role = getattr(msg, "type", None) or getattr(msg, "role", None)
        if role not in ("human", "user") and not (
            isinstance(msg, dict) and msg.get("role") == "user"
        ):
            continue

        if isinstance(msg, dict):
            if "content" in msg:
                msg["content"] = sanitized
                return True
            if "text" in msg:
                msg["text"] = sanitized
                return True
            msg["content"] = sanitized
            return True

        if hasattr(msg, "content"):
            try:
                setattr(msg, "content", sanitized)
                return True
            except Exception:
                continue

    return False


def _extract_model_name(serialized: Any, kwargs: Any) -> Optional[str]:
    # Best-effort; varies by integration.
    if isinstance(kwargs, dict):
        model = kwargs.get("model") or kwargs.get("model_name")
        if isinstance(model, str) and model:
            return model
    if isinstance(serialized, dict):
        model = serialized.get("name") or serialized.get("id")
        if isinstance(model, str) and model:
            return model
    return None


def _as_bool(value: Any) -> Optional[bool]:
    if isinstance(value, bool):
        return value
    return None


def _extract_stream_flag(serialized: Any, kwargs: Any) -> bool:
    if isinstance(kwargs, dict):
        direct = _as_bool(kwargs.get("stream"))
        if direct is None:
            direct = _as_bool(kwargs.get("streaming"))
        if direct is not None:
            return direct

        invocation = kwargs.get("invocation_params")
        if isinstance(invocation, dict):
            inv_stream = _as_bool(invocation.get("stream"))
            if inv_stream is None:
                inv_stream = _as_bool(invocation.get("streaming"))
            if inv_stream is not None:
                return inv_stream

    if isinstance(serialized, dict):
        serialized_kwargs = serialized.get("kwargs")
        if isinstance(serialized_kwargs, dict):
            ser_stream = _as_bool(serialized_kwargs.get("stream"))
            if ser_stream is None:
                ser_stream = _as_bool(serialized_kwargs.get("streaming"))
            if ser_stream is not None:
                return ser_stream

    return False


def _extract_output_text(response: Any) -> str:
    # LLMResult.generations: list[list[Generation]]
    gens = getattr(response, "generations", None)
    if isinstance(gens, list) and gens:
        first_list = gens[0]
        if isinstance(first_list, list) and first_list:
            gen0 = first_list[0]
            text = getattr(gen0, "text", None)
            if isinstance(text, str):
                return text
            message = getattr(gen0, "message", None)
            content = getattr(message, "content", None)
            if isinstance(content, str):
                return content
    return ""


def _extract_token_usage(response: Any) -> Optional[Dict[str, Any]]:
    llm_output = getattr(response, "llm_output", None) or getattr(response, "llmOutput", None)
    if isinstance(llm_output, dict):
        usage = llm_output.get("token_usage") or llm_output.get("tokenUsage") or llm_output.get("usage")
        if isinstance(usage, dict):
            return usage

    gens = getattr(response, "generations", None)
    if isinstance(gens, list) and gens:
        first_list = gens[0]
        if isinstance(first_list, list) and first_list:
            gen0 = first_list[0]
            message = getattr(gen0, "message", None)

            usage_meta = getattr(message, "usage_metadata", None)
            if isinstance(usage_meta, dict):
                return usage_meta

            response_meta = getattr(message, "response_metadata", None)
            if isinstance(response_meta, dict):
                token_usage = response_meta.get("token_usage") or response_meta.get("tokenUsage")
                if isinstance(token_usage, dict):
                    return token_usage

            generation_info = getattr(gen0, "generation_info", None)
            if isinstance(generation_info, dict):
                token_usage = generation_info.get("token_usage") or generation_info.get("tokenUsage")
                if isinstance(token_usage, dict):
                    return token_usage
    return None


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 _coerce_transparency_metadata(value: Any) -> Optional[Dict[str, Any]]:
    if not isinstance(value, dict):
        return None
    if (
        value.get("is_ai_generated") is not True
        or value.get("disclosure") != "You are interacting with an AI."
        or value.get("article") != "EU_AI_ACT_ARTICLE_50"
        or value.get("injection_mode") != "deterministic"
    ):
        return None
    return {
        "is_ai_generated": True,
        "disclosure": "You are interacting with an AI.",
        "article": "EU_AI_ACT_ARTICLE_50",
        "injection_mode": "deterministic",
    }


class AgentIDCallbackHandler(BaseCallbackHandler):
    """
    LangChain callback handler:
    - runs shared local enforcement before request leaves the client
    - blocks execution immediately when AgentID guard returns allowed=False
    - logs results to AgentID ingest (fire-and-forget via AgentID client)
    """

    def __init__(
        self,
        agent: Any,
        system_id: str,
        api_key: Optional[str] = None,
        expected_languages: Optional[list[str]] = None,
        expectedLanguages: Optional[list[str]] = None,
    ):
        self.agent = agent
        self.system_id = system_id
        self.api_key = api_key.strip() if isinstance(api_key, str) else None
        self.expected_languages = normalize_expected_languages(
            expected_languages if expected_languages is not None else expectedLanguages
        )
        self._runs: Dict[str, _RunState] = {}

    def _scan_injection_sync(self, input_text: str) -> None:
        cfg = getattr(self.agent, "_config", None)
        if cfg is None or not bool(getattr(cfg, "check_injection", False)):
            return

        effective_api_key = self.api_key or str(getattr(cfg, "api_key", "") or "").strip()
        base_url = str(getattr(cfg, "base_url", "") or "").strip()
        if not effective_api_key or not base_url:
            return

        capability_flags = ensure_capability_flags_sync(
            api_key=effective_api_key,
            base_url=base_url,
            timeout_seconds=float(getattr(cfg, "guard_timeout_s", 2.0)),
            ttl_seconds=300.0,
        )
        should_run_local_scan = bool(getattr(cfg, "check_injection", False))
        should_run_override = getattr(self.agent, "_should_run_local_injection_scan", None)
        if callable(should_run_override):
            should_run_local_scan = bool(should_run_override(capability_flags))
        if not should_run_local_scan:
            return

        get_injection_scanner().scan(
            prompt=input_text,
            api_key=effective_api_key,
            base_url=base_url,
            ai_scan_enabled=bool(getattr(cfg, "ai_scan_enabled", True)),
            store_pii=bool(getattr(cfg, "store_pii", False)),
            pii_manager=getattr(self.agent, "_pii", None),
            source="python_langchain",
        )

    def _langchain_capabilities(self) -> Dict[str, Any]:
        pii_masking_enabled = False
        resolver = getattr(self.agent, "_resolve_effective_pii_masking", None)
        get_cached_flags = getattr(self.agent, "_get_cached_capability_flags", None)
        if callable(resolver):
            try:
                capability_flags = (
                    get_cached_flags(api_key=self.api_key) if callable(get_cached_flags) else None
                )
                pii_masking_enabled = bool(resolver(capability_flags=capability_flags))
            except Exception:
                pii_masking_enabled = False
        return {
            "capabilities": {
                "has_feedback_handler": True,
                "pii_masking_enabled": pii_masking_enabled,
                "framework": "langchain",
            }
        }

    async def _scan_injection_async(self, input_text: str) -> None:
        loop = asyncio.get_running_loop()
        await loop.run_in_executor(None, self._scan_injection_sync, input_text)

    def _client_fast_fail_enabled(self) -> bool:
        cfg = getattr(self.agent, "_config", None)
        return bool(getattr(cfg, "client_fast_fail", False))

    def _prepare_input_sync(
        self,
        *,
        input_text: str,
        stream: bool,
        client_event_id: str,
    ) -> Dict[str, Any]:
        prepare_fn = getattr(self.agent, "_prepare_input_for_dispatch", None)
        if callable(prepare_fn) and not inspect.iscoroutinefunction(prepare_fn):
            return prepare_fn(
                input_text=input_text,
                system_id=self.system_id,
                stream=stream,
                api_key=self.api_key,
                client_event_id=client_event_id,
            )

        self._scan_injection_sync(input_text)
        sanitized_input = self._enforce_local_sync(input_text, stream)
        return {
            "sanitized_input": sanitized_input,
            "capability_flags": None,
            "sdk_config_fetch_ms": 0,
            "sdk_local_scan_ms": 0,
        }

    async def _prepare_input_async(
        self,
        *,
        input_text: str,
        stream: bool,
        client_event_id: str,
    ) -> Dict[str, Any]:
        prepare_fn = getattr(self.agent, "_prepare_input_for_dispatch_async", None)
        if callable(prepare_fn) and inspect.iscoroutinefunction(prepare_fn):
            return await prepare_fn(
                input_text=input_text,
                system_id=self.system_id,
                stream=stream,
                api_key=self.api_key,
                client_event_id=client_event_id,
            )

        await self._scan_injection_async(input_text)
        sanitized_input = await self._enforce_local_async(input_text, stream)
        return {
            "sanitized_input": sanitized_input,
            "capability_flags": None,
            "sdk_config_fetch_ms": 0,
            "sdk_local_scan_ms": 0,
        }

    def _apply_local_fallback_sync(
        self,
        *,
        input_text: str,
        stream: bool,
        client_event_id: str,
        capability_flags: Any,
        sdk_config_fetch_ms: int,
    ) -> Dict[str, Any]:
        fallback_fn = getattr(self.agent, "_apply_local_fallback_for_guard_failure", None)
        if callable(fallback_fn) and not inspect.iscoroutinefunction(fallback_fn):
            return fallback_fn(
                input_text=input_text,
                system_id=self.system_id,
                stream=stream,
                api_key=self.api_key,
                client_event_id=client_event_id,
                capability_flags=capability_flags,
                sdk_config_fetch_ms=sdk_config_fetch_ms,
            )

        sanitized_input = self._enforce_local_sync(input_text, stream)
        return {
            "sanitized_input": sanitized_input,
            "capability_flags": capability_flags,
            "sdk_config_fetch_ms": sdk_config_fetch_ms,
            "sdk_local_scan_ms": 0,
        }

    async def _apply_local_fallback_async(
        self,
        *,
        input_text: str,
        stream: bool,
        client_event_id: str,
        capability_flags: Any,
        sdk_config_fetch_ms: int,
    ) -> Dict[str, Any]:
        fallback_fn = getattr(self.agent, "_apply_local_fallback_for_guard_failure_async", None)
        if callable(fallback_fn) and inspect.iscoroutinefunction(fallback_fn):
            return await fallback_fn(
                input_text=input_text,
                system_id=self.system_id,
                stream=stream,
                api_key=self.api_key,
                client_event_id=client_event_id,
                capability_flags=capability_flags,
                sdk_config_fetch_ms=sdk_config_fetch_ms,
            )

        sanitized_input = await self._enforce_local_async(input_text, stream)
        return {
            "sanitized_input": sanitized_input,
            "capability_flags": capability_flags,
            "sdk_config_fetch_ms": sdk_config_fetch_ms,
            "sdk_local_scan_ms": 0,
        }

    def _guard_sync(
        self,
        input_text: str,
        model: Optional[str] = None,
        client_event_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        return self.agent.guard(
            input=input_text,
            system_id=self.system_id,
            model=model,
            client_event_id=client_event_id,
            api_key=self.api_key,
            expected_languages=self.expected_languages,
            client_capabilities=self._langchain_capabilities(),
        )

    async def _guard_async(
        self,
        input_text: str,
        model: Optional[str] = None,
        client_event_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        guard_fn = getattr(self.agent, "guard", None)
        if guard_fn is None:
            return {"allowed": True, "fail_open": True}

        if inspect.iscoroutinefunction(guard_fn):
            return await guard_fn(
                input=input_text,
                system_id=self.system_id,
                model=model,
                client_event_id=client_event_id,
                api_key=self.api_key,
                expected_languages=self.expected_languages,
                client_capabilities=self._langchain_capabilities(),
            )

        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(
            None, self._guard_sync, input_text, model, client_event_id
        )

    def _enforce_local_sync(self, input_text: str, stream: bool) -> str:
        enforce_fn = getattr(self.agent, "_apply_local_security", None)
        if not callable(enforce_fn):
            return input_text

        sanitized, _flags = enforce_fn(
            system_id=self.system_id,
            input_text=input_text,
            stream=stream,
            api_key=self.api_key,
        )
        return sanitized

    async def _enforce_local_async(self, input_text: str, stream: bool) -> str:
        async_enforce_fn = getattr(self.agent, "_apply_local_security_async", None)
        if callable(async_enforce_fn) and inspect.iscoroutinefunction(async_enforce_fn):
            sanitized, _flags = await async_enforce_fn(
                system_id=self.system_id,
                input_text=input_text,
                stream=stream,
                api_key=self.api_key,
            )
            return sanitized

        sync_enforce_fn = getattr(self.agent, "_apply_local_security", None)
        if not callable(sync_enforce_fn):
            return input_text

        loop = asyncio.get_running_loop()

        def _run_sync() -> str:
            sanitized, _flags = sync_enforce_fn(
                system_id=self.system_id,
                input_text=input_text,
                stream=stream,
                api_key=self.api_key,
            )
            return sanitized

        return await loop.run_in_executor(None, _run_sync)

    def _strict_mode_from_config(self) -> bool:
        cfg = getattr(self.agent, "_config", None)
        if cfg is None:
            return False
        return bool(
            bool(getattr(cfg, "strict_mode", False))
            or getattr(cfg, "failure_mode", None) == "fail_close"
        )

    def _resolve_strict_mode_sync(self) -> bool:
        resolver = getattr(self.agent, "_resolve_effective_strict_mode", None)
        if not callable(resolver):
            return self._strict_mode_from_config()
        if inspect.iscoroutinefunction(resolver):
            return self._strict_mode_from_config()
        try:
            return bool(resolver(api_key=self.api_key))
        except Exception:
            return self._strict_mode_from_config()

    async def _resolve_strict_mode_async(self) -> bool:
        resolver = getattr(self.agent, "_resolve_effective_strict_mode", None)
        if not callable(resolver):
            return self._strict_mode_from_config()
        try:
            if inspect.iscoroutinefunction(resolver):
                return bool(await resolver(api_key=self.api_key))
            return bool(resolver(api_key=self.api_key))
        except Exception:
            return self._strict_mode_from_config()

    def _invoke_log_sync(self, payload: Dict[str, Any]) -> None:
        log_fn = getattr(self.agent, "log", None)
        if not callable(log_fn):
            return
        strict_for_ingest = self._resolve_strict_mode_sync()
        try:
            signature = inspect.signature(log_fn)
            if "raise_on_failure" in signature.parameters:
                payload_with_mode = dict(payload)
                payload_with_mode["raise_on_failure"] = strict_for_ingest
                log_fn(**payload_with_mode)
                return
        except (TypeError, ValueError):
            pass
        log_fn(**payload)

    async def _invoke_log_async(self, payload: Dict[str, Any]) -> None:
        strict_for_ingest = await self._resolve_strict_mode_async()

        log_async_fn = getattr(self.agent, "log_async", None)
        if callable(log_async_fn) and inspect.iscoroutinefunction(log_async_fn):
            try:
                signature = inspect.signature(log_async_fn)
                if "raise_on_failure" in signature.parameters:
                    payload_with_mode = dict(payload)
                    payload_with_mode["raise_on_failure"] = strict_for_ingest
                    await log_async_fn(**payload_with_mode)
                    return
            except (TypeError, ValueError):
                pass
            await log_async_fn(**payload)
            return

        self._invoke_log_sync(payload)

    def _block_if_needed(self, verdict: Dict[str, Any]) -> None:
        if verdict.get("allowed") is False:
            raise SecurityBlockError(verdict.get("reason"))

    @staticmethod
    def _resolve_transformed_input(verdict: Dict[str, Any], fallback_input: str) -> str:
        transformed = verdict.get("transformed_input")
        if isinstance(transformed, str) and transformed:
            return transformed
        return fallback_input

    # ---- Sync hooks ----
    def on_llm_start(self, serialized: Dict[str, Any], prompts: Any, **kwargs: Any) -> Any:
        run_id = str(kwargs.get("run_id") or kwargs.get("runId") or "")
        input_text = _extract_prompt_from_prompts(prompts) or ""
        if not input_text:
            raise ValueError("AgentID: No prompt found. Security guard requires string input.")

        stream = _extract_stream_flag(serialized, kwargs)
        request_client_event_id = str(uuid.uuid4())
        prepared = self._prepare_input_sync(
            input_text=input_text,
            stream=stream,
            client_event_id=request_client_event_id,
        )
        sanitized_input = str(prepared.get("sanitized_input") or input_text)
        capability_flags = prepared.get("capability_flags")
        sdk_config_fetch_ms = int(prepared.get("sdk_config_fetch_ms") or 0)
        sdk_local_scan_ms = int(prepared.get("sdk_local_scan_ms") or 0)
        if sanitized_input != input_text and not _set_prompt_in_prompts(prompts, sanitized_input):
            raise ValueError("AgentID: Strict PII mode requires mutable LangChain prompt payload.")

        model = _extract_model_name(serialized, kwargs.get("invocation_params") or kwargs)
        effective_strict_mode = _is_effective_strict_mode(
            getattr(self.agent, "_config", None), capability_flags
        )
        guard_started_at = time.perf_counter()
        verdict = self._guard_sync(
            sanitized_input, model=model, client_event_id=request_client_event_id
        )
        sdk_guard_ms = (
            int(verdict.get("guard_latency_ms"))
            if isinstance(verdict.get("guard_latency_ms"), int)
            else int((time.perf_counter() - guard_started_at) * 1000)
        )
        local_fallback_applied = False
        local_fallback_reason: Optional[str] = None
        if verdict.get("allowed") is False:
            if (
                effective_strict_mode
                and _is_guard_failure_eligible_for_local_fallback(verdict.get("reason"))
            ):
                local_fallback_applied = True
                local_fallback_reason = str(
                    verdict.get("reason") or "guard_unreachable"
                )
                if sdk_local_scan_ms == 0:
                    fallback = self._apply_local_fallback_sync(
                        input_text=sanitized_input,
                        stream=stream,
                        client_event_id=request_client_event_id,
                        capability_flags=capability_flags,
                        sdk_config_fetch_ms=sdk_config_fetch_ms,
                    )
                    sanitized_input = str(
                        fallback.get("sanitized_input") or sanitized_input
                    )
                    capability_flags = fallback.get("capability_flags")
                    sdk_config_fetch_ms = int(
                        fallback.get("sdk_config_fetch_ms") or sdk_config_fetch_ms
                    )
                    sdk_local_scan_ms = int(fallback.get("sdk_local_scan_ms") or 0)
                    if sanitized_input != input_text and not _set_prompt_in_prompts(
                        prompts, sanitized_input
                    ):
                        raise ValueError(
                            "AgentID: Strict PII mode requires mutable LangChain prompt payload."
                        )
            else:
                self._block_if_needed(verdict)
        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 = self._resolve_transformed_input(verdict, sanitized_input)
        if transformed_input != sanitized_input and not _set_prompt_in_prompts(prompts, transformed_input):
            raise ValueError(
                "AgentID: Guard transformed input could not be applied to LangChain prompt payload."
            )
        self._runs[run_id] = _RunState(
            input_text=transformed_input,
            started_at=time.perf_counter(),
            model=_extract_model_name(serialized, kwargs.get("invocation_params") or kwargs),
            client_event_id=canonical_client_event_id,
            guard_event_id=guard_event_id_str,
            transparency=_coerce_transparency_metadata(verdict.get("transparency")),
            sdk_config_fetch_ms=sdk_config_fetch_ms,
            sdk_local_scan_ms=sdk_local_scan_ms,
            sdk_guard_ms=sdk_guard_ms,
            local_fallback_applied=local_fallback_applied,
            local_fallback_reason=local_fallback_reason,
        )

    def on_chat_model_start(self, serialized: Dict[str, Any], messages: Any, **kwargs: Any) -> Any:
        run_id = str(kwargs.get("run_id") or kwargs.get("runId") or "")
        input_text = _extract_prompt_from_messages(messages) or ""
        if not input_text:
            raise ValueError(
                "AgentID: No user message found. Security guard requires string input."
            )

        stream = _extract_stream_flag(serialized, kwargs)
        request_client_event_id = str(uuid.uuid4())
        prepared = self._prepare_input_sync(
            input_text=input_text,
            stream=stream,
            client_event_id=request_client_event_id,
        )
        sanitized_input = str(prepared.get("sanitized_input") or input_text)
        capability_flags = prepared.get("capability_flags")
        sdk_config_fetch_ms = int(prepared.get("sdk_config_fetch_ms") or 0)
        sdk_local_scan_ms = int(prepared.get("sdk_local_scan_ms") or 0)
        if sanitized_input != input_text and not _set_prompt_in_messages(messages, sanitized_input):
            raise ValueError("AgentID: Strict PII mode requires mutable LangChain message payload.")

        model = _extract_model_name(serialized, kwargs.get("invocation_params") or kwargs)
        effective_strict_mode = _is_effective_strict_mode(
            getattr(self.agent, "_config", None), capability_flags
        )
        guard_started_at = time.perf_counter()
        verdict = self._guard_sync(
            sanitized_input, model=model, client_event_id=request_client_event_id
        )
        sdk_guard_ms = (
            int(verdict.get("guard_latency_ms"))
            if isinstance(verdict.get("guard_latency_ms"), int)
            else int((time.perf_counter() - guard_started_at) * 1000)
        )
        local_fallback_applied = False
        local_fallback_reason: Optional[str] = None
        if verdict.get("allowed") is False:
            if (
                effective_strict_mode
                and _is_guard_failure_eligible_for_local_fallback(verdict.get("reason"))
            ):
                local_fallback_applied = True
                local_fallback_reason = str(
                    verdict.get("reason") or "guard_unreachable"
                )
                if sdk_local_scan_ms == 0:
                    fallback = self._apply_local_fallback_sync(
                        input_text=sanitized_input,
                        stream=stream,
                        client_event_id=request_client_event_id,
                        capability_flags=capability_flags,
                        sdk_config_fetch_ms=sdk_config_fetch_ms,
                    )
                    sanitized_input = str(
                        fallback.get("sanitized_input") or sanitized_input
                    )
                    capability_flags = fallback.get("capability_flags")
                    sdk_config_fetch_ms = int(
                        fallback.get("sdk_config_fetch_ms") or sdk_config_fetch_ms
                    )
                    sdk_local_scan_ms = int(fallback.get("sdk_local_scan_ms") or 0)
                    if sanitized_input != input_text and not _set_prompt_in_messages(
                        messages, sanitized_input
                    ):
                        raise ValueError(
                            "AgentID: Strict PII mode requires mutable LangChain message payload."
                        )
            else:
                self._block_if_needed(verdict)
        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 = self._resolve_transformed_input(verdict, sanitized_input)
        if transformed_input != sanitized_input and not _set_prompt_in_messages(messages, transformed_input):
            raise ValueError(
                "AgentID: Guard transformed input could not be applied to LangChain message payload."
            )
        self._runs[run_id] = _RunState(
            input_text=transformed_input,
            started_at=time.perf_counter(),
            model=_extract_model_name(serialized, kwargs.get("invocation_params") or kwargs),
            client_event_id=canonical_client_event_id,
            guard_event_id=guard_event_id_str,
            transparency=_coerce_transparency_metadata(verdict.get("transparency")),
            sdk_config_fetch_ms=sdk_config_fetch_ms,
            sdk_local_scan_ms=sdk_local_scan_ms,
            sdk_guard_ms=sdk_guard_ms,
            local_fallback_applied=local_fallback_applied,
            local_fallback_reason=local_fallback_reason,
        )

    def on_llm_end(self, response: Any, **kwargs: Any) -> Any:
        run_id = str(kwargs.get("run_id") or kwargs.get("runId") or "")
        state = self._runs.pop(run_id, None)
        if not state:
            return

        latency_ms = int((time.perf_counter() - state.started_at) * 1000)
        output_text = _extract_output_text(response)
        usage = _extract_token_usage(response)
        metadata: Dict[str, Any] = {}
        if state.client_event_id:
            metadata["client_event_id"] = state.client_event_id
        if state.guard_event_id:
            metadata["guard_event_id"] = state.guard_event_id
        if state.transparency:
            metadata["transparency"] = state.transparency
        metadata.update(
            _build_sdk_timing_metadata(
                sdk_config_fetch_ms=state.sdk_config_fetch_ms,
                sdk_local_scan_ms=state.sdk_local_scan_ms,
                sdk_guard_ms=state.sdk_guard_ms,
            )
        )
        if state.local_fallback_applied:
            metadata["sdk_local_fallback_applied"] = True
            metadata["sdk_local_fallback_reason"] = state.local_fallback_reason

        self._invoke_log_sync(
            {
                "system_id": self.system_id,
                "input": state.input_text,
                "output": output_text,
                "event_id": state.client_event_id,
                "model": state.model or "unknown",
                "usage": usage if isinstance(usage, dict) else None,
                "metadata": metadata or None,
                "latency": latency_ms,
                "api_key": self.api_key,
                "client_capabilities": self._langchain_capabilities(),
            }
        )

    def on_llm_error(self, error: BaseException, **kwargs: Any) -> Any:
        run_id = str(kwargs.get("run_id") or kwargs.get("runId") or "")
        state = self._runs.pop(run_id, None)
        input_text = state.input_text if state else ""
        metadata: Dict[str, Any] = {
            "error_type": error.__class__.__name__,
            "error_message": str(error),
        }
        if state and state.client_event_id:
            metadata["client_event_id"] = state.client_event_id
        if state and state.guard_event_id:
            metadata["guard_event_id"] = state.guard_event_id
        if state and state.transparency:
            metadata["transparency"] = state.transparency
        if state:
            metadata.update(
                _build_sdk_timing_metadata(
                    sdk_config_fetch_ms=state.sdk_config_fetch_ms,
                    sdk_local_scan_ms=state.sdk_local_scan_ms,
                    sdk_guard_ms=state.sdk_guard_ms,
                )
            )
            if state.local_fallback_applied:
                metadata["sdk_local_fallback_applied"] = True
                metadata["sdk_local_fallback_reason"] = state.local_fallback_reason

        self._invoke_log_sync(
            {
                "system_id": self.system_id,
                "input": input_text,
                "output": "",
                "event_id": state.client_event_id if state else None,
                "model": state.model if state and state.model else "unknown",
                "event_type": "error",
                "severity": "error",
                "metadata": metadata,
                "api_key": self.api_key,
                "client_capabilities": self._langchain_capabilities(),
            }
        )

    # ---- Async hooks ----
    async def aon_llm_start(self, serialized: Dict[str, Any], prompts: Any, **kwargs: Any) -> Any:
        run_id = str(kwargs.get("run_id") or kwargs.get("runId") or "")
        input_text = _extract_prompt_from_prompts(prompts) or ""
        if not input_text:
            raise ValueError("AgentID: No prompt found. Security guard requires string input.")

        stream = _extract_stream_flag(serialized, kwargs)
        request_client_event_id = str(uuid.uuid4())
        prepared = await self._prepare_input_async(
            input_text=input_text,
            stream=stream,
            client_event_id=request_client_event_id,
        )
        sanitized_input = str(prepared.get("sanitized_input") or input_text)
        capability_flags = prepared.get("capability_flags")
        sdk_config_fetch_ms = int(prepared.get("sdk_config_fetch_ms") or 0)
        sdk_local_scan_ms = int(prepared.get("sdk_local_scan_ms") or 0)
        if sanitized_input != input_text and not _set_prompt_in_prompts(prompts, sanitized_input):
            raise ValueError("AgentID: Strict PII mode requires mutable LangChain prompt payload.")

        model = _extract_model_name(serialized, kwargs.get("invocation_params") or kwargs)
        effective_strict_mode = _is_effective_strict_mode(
            getattr(self.agent, "_config", None), capability_flags
        )
        guard_started_at = time.perf_counter()
        verdict = await self._guard_async(
            sanitized_input, model=model, client_event_id=request_client_event_id
        )
        sdk_guard_ms = (
            int(verdict.get("guard_latency_ms"))
            if isinstance(verdict.get("guard_latency_ms"), int)
            else int((time.perf_counter() - guard_started_at) * 1000)
        )
        local_fallback_applied = False
        local_fallback_reason: Optional[str] = None
        if verdict.get("allowed") is False:
            if (
                effective_strict_mode
                and _is_guard_failure_eligible_for_local_fallback(verdict.get("reason"))
            ):
                local_fallback_applied = True
                local_fallback_reason = str(
                    verdict.get("reason") or "guard_unreachable"
                )
                if sdk_local_scan_ms == 0:
                    fallback = await self._apply_local_fallback_async(
                        input_text=sanitized_input,
                        stream=stream,
                        client_event_id=request_client_event_id,
                        capability_flags=capability_flags,
                        sdk_config_fetch_ms=sdk_config_fetch_ms,
                    )
                    sanitized_input = str(
                        fallback.get("sanitized_input") or sanitized_input
                    )
                    capability_flags = fallback.get("capability_flags")
                    sdk_config_fetch_ms = int(
                        fallback.get("sdk_config_fetch_ms") or sdk_config_fetch_ms
                    )
                    sdk_local_scan_ms = int(fallback.get("sdk_local_scan_ms") or 0)
                    if sanitized_input != input_text and not _set_prompt_in_prompts(
                        prompts, sanitized_input
                    ):
                        raise ValueError(
                            "AgentID: Strict PII mode requires mutable LangChain prompt payload."
                        )
            else:
                self._block_if_needed(verdict)
        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 = self._resolve_transformed_input(verdict, sanitized_input)
        if transformed_input != sanitized_input and not _set_prompt_in_prompts(prompts, transformed_input):
            raise ValueError(
                "AgentID: Guard transformed input could not be applied to LangChain prompt payload."
            )
        self._runs[run_id] = _RunState(
            input_text=transformed_input,
            started_at=time.perf_counter(),
            model=_extract_model_name(serialized, kwargs.get("invocation_params") or kwargs),
            client_event_id=canonical_client_event_id,
            guard_event_id=guard_event_id_str,
            transparency=_coerce_transparency_metadata(verdict.get("transparency")),
            sdk_config_fetch_ms=sdk_config_fetch_ms,
            sdk_local_scan_ms=sdk_local_scan_ms,
            sdk_guard_ms=sdk_guard_ms,
            local_fallback_applied=local_fallback_applied,
            local_fallback_reason=local_fallback_reason,
        )

    async def aon_chat_model_start(
        self, serialized: Dict[str, Any], messages: Any, **kwargs: Any
    ) -> Any:
        run_id = str(kwargs.get("run_id") or kwargs.get("runId") or "")
        input_text = _extract_prompt_from_messages(messages) or ""
        if not input_text:
            raise ValueError(
                "AgentID: No user message found. Security guard requires string input."
            )

        stream = _extract_stream_flag(serialized, kwargs)
        request_client_event_id = str(uuid.uuid4())
        prepared = await self._prepare_input_async(
            input_text=input_text,
            stream=stream,
            client_event_id=request_client_event_id,
        )
        sanitized_input = str(prepared.get("sanitized_input") or input_text)
        capability_flags = prepared.get("capability_flags")
        sdk_config_fetch_ms = int(prepared.get("sdk_config_fetch_ms") or 0)
        sdk_local_scan_ms = int(prepared.get("sdk_local_scan_ms") or 0)
        if sanitized_input != input_text and not _set_prompt_in_messages(messages, sanitized_input):
            raise ValueError("AgentID: Strict PII mode requires mutable LangChain message payload.")

        model = _extract_model_name(serialized, kwargs.get("invocation_params") or kwargs)
        effective_strict_mode = _is_effective_strict_mode(
            getattr(self.agent, "_config", None), capability_flags
        )
        guard_started_at = time.perf_counter()
        verdict = await self._guard_async(
            sanitized_input, model=model, client_event_id=request_client_event_id
        )
        sdk_guard_ms = (
            int(verdict.get("guard_latency_ms"))
            if isinstance(verdict.get("guard_latency_ms"), int)
            else int((time.perf_counter() - guard_started_at) * 1000)
        )
        local_fallback_applied = False
        local_fallback_reason: Optional[str] = None
        if verdict.get("allowed") is False:
            if (
                effective_strict_mode
                and _is_guard_failure_eligible_for_local_fallback(verdict.get("reason"))
            ):
                local_fallback_applied = True
                local_fallback_reason = str(
                    verdict.get("reason") or "guard_unreachable"
                )
                if sdk_local_scan_ms == 0:
                    fallback = await self._apply_local_fallback_async(
                        input_text=sanitized_input,
                        stream=stream,
                        client_event_id=request_client_event_id,
                        capability_flags=capability_flags,
                        sdk_config_fetch_ms=sdk_config_fetch_ms,
                    )
                    sanitized_input = str(
                        fallback.get("sanitized_input") or sanitized_input
                    )
                    capability_flags = fallback.get("capability_flags")
                    sdk_config_fetch_ms = int(
                        fallback.get("sdk_config_fetch_ms") or sdk_config_fetch_ms
                    )
                    sdk_local_scan_ms = int(fallback.get("sdk_local_scan_ms") or 0)
                    if sanitized_input != input_text and not _set_prompt_in_messages(
                        messages, sanitized_input
                    ):
                        raise ValueError(
                            "AgentID: Strict PII mode requires mutable LangChain message payload."
                        )
            else:
                self._block_if_needed(verdict)
        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 = self._resolve_transformed_input(verdict, sanitized_input)
        if transformed_input != sanitized_input and not _set_prompt_in_messages(messages, transformed_input):
            raise ValueError(
                "AgentID: Guard transformed input could not be applied to LangChain message payload."
            )
        self._runs[run_id] = _RunState(
            input_text=transformed_input,
            started_at=time.perf_counter(),
            model=_extract_model_name(serialized, kwargs.get("invocation_params") or kwargs),
            client_event_id=canonical_client_event_id,
            guard_event_id=guard_event_id_str,
            transparency=_coerce_transparency_metadata(verdict.get("transparency")),
            sdk_config_fetch_ms=sdk_config_fetch_ms,
            sdk_local_scan_ms=sdk_local_scan_ms,
            sdk_guard_ms=sdk_guard_ms,
            local_fallback_applied=local_fallback_applied,
            local_fallback_reason=local_fallback_reason,
        )

    async def aon_llm_end(self, response: Any, **kwargs: Any) -> Any:
        run_id = str(kwargs.get("run_id") or kwargs.get("runId") or "")
        state = self._runs.pop(run_id, None)
        if not state:
            return

        latency_ms = int((time.perf_counter() - state.started_at) * 1000)
        output_text = _extract_output_text(response)
        usage = _extract_token_usage(response)
        metadata: Dict[str, Any] = {}
        if state.client_event_id:
            metadata["client_event_id"] = state.client_event_id
        if state.guard_event_id:
            metadata["guard_event_id"] = state.guard_event_id
        if state.transparency:
            metadata["transparency"] = state.transparency
        metadata.update(
            _build_sdk_timing_metadata(
                sdk_config_fetch_ms=state.sdk_config_fetch_ms,
                sdk_local_scan_ms=state.sdk_local_scan_ms,
                sdk_guard_ms=state.sdk_guard_ms,
            )
        )
        if state.local_fallback_applied:
            metadata["sdk_local_fallback_applied"] = True
            metadata["sdk_local_fallback_reason"] = state.local_fallback_reason

        await self._invoke_log_async(
            {
                "system_id": self.system_id,
                "input": state.input_text,
                "output": output_text,
                "event_id": state.client_event_id,
                "model": state.model or "unknown",
                "usage": usage if isinstance(usage, dict) else None,
                "metadata": metadata or None,
                "latency": latency_ms,
                "api_key": self.api_key,
                "client_capabilities": self._langchain_capabilities(),
            }
        )

    async def aon_llm_error(self, error: BaseException, **kwargs: Any) -> Any:
        run_id = str(kwargs.get("run_id") or kwargs.get("runId") or "")
        state = self._runs.pop(run_id, None)
        input_text = state.input_text if state else ""
        metadata: Dict[str, Any] = {
            "error_type": error.__class__.__name__,
            "error_message": str(error),
        }
        if state and state.client_event_id:
            metadata["client_event_id"] = state.client_event_id
        if state and state.guard_event_id:
            metadata["guard_event_id"] = state.guard_event_id
        if state and state.transparency:
            metadata["transparency"] = state.transparency
        if state:
            metadata.update(
                _build_sdk_timing_metadata(
                    sdk_config_fetch_ms=state.sdk_config_fetch_ms,
                    sdk_local_scan_ms=state.sdk_local_scan_ms,
                    sdk_guard_ms=state.sdk_guard_ms,
                )
            )
            if state.local_fallback_applied:
                metadata["sdk_local_fallback_applied"] = True
                metadata["sdk_local_fallback_reason"] = state.local_fallback_reason

        await self._invoke_log_async(
            {
                "system_id": self.system_id,
                "input": input_text,
                "output": "",
                "event_id": state.client_event_id if state else None,
                "model": state.model if state and state.model else "unknown",
                "event_type": "error",
                "severity": "error",
                "metadata": metadata,
                "api_key": self.api_key,
                "client_capabilities": self._langchain_capabilities(),
            }
        )
