# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from collections.abc import Callable
from types import TracebackType
from typing import AsyncContextManager, Literal, Protocol, Self, TypedDict

from nemo_relay import (
    Event,
    EventSanitizeGuardrail,
    JsonObject,
    LlmConditionalExecutionGuardrail,
    LlmExecutionIntercept,
    LlmRequestIntercept,
    LlmSanitizeRequestGuardrail,
    LlmSanitizeResponseGuardrail,
    LlmStreamExecutionIntercept,
    ToolConditionalExecutionGuardrail,
    ToolExecutionIntercept,
    ToolRequestIntercept,
    ToolSanitizeGuardrail,
)

UnsupportedBehavior = Literal["ignore", "warn", "error"]
DynamicPluginKind = Literal["rust_dynamic", "worker"]

class _ConfigDiagnosticRequired(TypedDict):
    level: Literal["warning", "error"]
    code: str
    message: str

class ConfigDiagnostic(_ConfigDiagnosticRequired, total=False):
    component: str
    field: str

class ConfigReport(TypedDict):
    diagnostics: list[ConfigDiagnostic]

class PluginContext(Protocol):
    def register_subscriber(self, name: str, callback: Callable[[Event], None]) -> None: ...
    def register_mark_sanitize_guardrail(self, name: str, priority: int, callback: EventSanitizeGuardrail) -> None: ...
    def register_scope_sanitize_start_guardrail(
        self, name: str, priority: int, callback: EventSanitizeGuardrail
    ) -> None: ...
    def register_scope_sanitize_end_guardrail(
        self, name: str, priority: int, callback: EventSanitizeGuardrail
    ) -> None: ...
    def register_tool_sanitize_request_guardrail(
        self, name: str, priority: int, callback: ToolSanitizeGuardrail
    ) -> None: ...
    def register_tool_sanitize_response_guardrail(
        self, name: str, priority: int, callback: ToolSanitizeGuardrail
    ) -> None: ...
    def register_tool_conditional_execution_guardrail(
        self, name: str, priority: int, callback: ToolConditionalExecutionGuardrail
    ) -> None: ...
    def register_llm_sanitize_request_guardrail(
        self, name: str, priority: int, callback: LlmSanitizeRequestGuardrail
    ) -> None: ...
    def register_llm_sanitize_response_guardrail(
        self, name: str, priority: int, callback: LlmSanitizeResponseGuardrail
    ) -> None: ...
    def register_llm_conditional_execution_guardrail(
        self, name: str, priority: int, callback: LlmConditionalExecutionGuardrail
    ) -> None: ...
    def register_llm_request_intercept(
        self, name: str, priority: int, break_chain: bool, callback: LlmRequestIntercept
    ) -> None: ...
    def register_llm_execution_intercept(self, name: str, priority: int, callback: LlmExecutionIntercept) -> None: ...
    def register_llm_stream_execution_intercept(
        self, name: str, priority: int, callback: LlmStreamExecutionIntercept
    ) -> None: ...
    def register_tool_request_intercept(
        self, name: str, priority: int, break_chain: bool, callback: ToolRequestIntercept
    ) -> None: ...
    def register_tool_execution_intercept(self, name: str, priority: int, callback: ToolExecutionIntercept) -> None: ...

class Plugin(Protocol):
    def validate(self, plugin_config: JsonObject) -> list[ConfigDiagnostic] | None: ...
    def register(self, plugin_config: JsonObject, context: PluginContext) -> None: ...

class ConfigPolicy:
    unknown_component: UnsupportedBehavior
    unknown_field: UnsupportedBehavior
    unsupported_value: UnsupportedBehavior

    def __init__(
        self,
        unknown_component: UnsupportedBehavior = "warn",
        unknown_field: UnsupportedBehavior = "warn",
        unsupported_value: UnsupportedBehavior = "error",
    ) -> None: ...
    def to_dict(self) -> JsonObject: ...

class ComponentSpec:
    kind: str
    enabled: bool
    config: JsonObject

    def __init__(
        self,
        kind: str,
        enabled: bool = True,
        config: JsonObject = ...,
    ) -> None: ...
    def to_dict(self) -> JsonObject: ...

class PluginConfig:
    version: int
    components: list[object]
    policy: ConfigPolicy

    def __init__(
        self,
        version: int = 1,
        components: list[object] = ...,
        policy: ConfigPolicy = ...,
    ) -> None: ...
    def to_dict(self) -> JsonObject: ...

class DynamicPluginActivationSpec:
    plugin_id: str
    kind: DynamicPluginKind
    manifest_ref: str
    environment_ref: str | None
    config: JsonObject

    def __init__(
        self,
        plugin_id: str,
        kind: DynamicPluginKind,
        manifest_ref: str,
        environment_ref: str | None = None,
        config: JsonObject = ...,
    ) -> None: ...
    def to_dict(self) -> JsonObject: ...

class PluginHostActivation:
    @property
    def report(self) -> ConfigReport: ...
    @property
    def is_active(self) -> bool: ...
    async def close(self) -> None: ...
    async def __aenter__(self) -> Self: ...
    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> None: ...

def validate(config: PluginConfig | JsonObject) -> ConfigReport: ...
async def initialize(config: PluginConfig | JsonObject) -> ConfigReport: ...
async def initialize_with_dynamic_plugins(
    config: PluginConfig | JsonObject,
    dynamic_plugins: list[DynamicPluginActivationSpec | JsonObject],
) -> PluginHostActivation: ...
def clear() -> None: ...
def plugin(config: PluginConfig | JsonObject) -> AsyncContextManager[ConfigReport]: ...
def report() -> ConfigReport | None: ...
def list_kinds() -> list[str]: ...
def register(plugin_kind: str, plugin: Plugin) -> None: ...
def deregister(plugin_kind: str) -> bool: ...