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

"""Type stubs for ``nemo_relay.pii_redaction``."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Literal, TypedDict

from nemo_relay import JsonObject, UnsupportedBehavior

class ConfigDiagnostic(TypedDict, total=False):
    level: Literal["warning", "error"]
    code: str
    message: str
    component: str
    field: str

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

@dataclass(slots=True)
class ConfigPolicy:
    unknown_component: UnsupportedBehavior = ...
    unknown_field: UnsupportedBehavior = ...
    unsupported_value: UnsupportedBehavior = ...
    def to_dict(self) -> JsonObject: ...

@dataclass(slots=True)
class BuiltinConfig:
    action: Literal["remove", "redact", "regex_replace", "hash", "mask"] = ...
    target_paths: list[str] = field(default_factory=list)
    pattern: str | None = ...
    detector: str | None = ...
    replacement: str | None = ...
    mask_char: str | None = ...
    unmasked_prefix: int | None = ...
    unmasked_suffix: int | None = ...
    def to_dict(self) -> JsonObject: ...

@dataclass(slots=True)
class LocalModelConfig:
    backend: str | None = ...
    model_id: str | None = ...
    detector_profile: str | None = ...
    allow_network: bool | None = ...
    max_latency_ms: int | None = ...
    def to_dict(self) -> JsonObject: ...

@dataclass(slots=True)
class PiiRedactionConfig:
    version: int = ...
    mode: Literal["builtin", "local_model"] = ...
    input: bool = ...
    output: bool = ...
    tool_input: bool = ...
    tool_output: bool = ...
    mark: bool = ...
    priority: int = ...
    codec: Literal["openai_chat", "openai_responses", "anthropic_messages"] | str | None = ...
    builtin: BuiltinConfig | None = ...
    local: LocalModelConfig | None = ...
    policy: ConfigPolicy = field(default_factory=ConfigPolicy)
    def to_dict(self) -> JsonObject: ...

PII_REDACTION_PLUGIN_KIND: Literal["pii_redaction"]

@dataclass(slots=True)
class ComponentSpec:
    config: PiiRedactionConfig | JsonObject
    enabled: bool = ...
    def to_dict(self) -> JsonObject: ...

def validate_config(config: PiiRedactionConfig | JsonObject) -> ConfigReport: ...