Utils

class EncodingSafeProcessor:
    r"""Sanitize strings for console encoding compatibility.

    A structlog processor that ensures all string values in the event dict
    can be safely written to the console regardless of its encoding (cp1252
    on Windows, UTF-8 on macOS/Linux).

    The sanitization strategy:
    1. Apply smart character replacements (preserve meaning)
    2. Encode to console encoding with 'replace' error handler
    3. Decode back to string

    This produces strings that are guaranteed to be writable to the console
    without raising UnicodeEncodeError.

    Attributes:
        REPLACEMENTS (dict[str, str]): Class constant mapping Unicode characters
            to ASCII equivalents that preserve semantic meaning.
        encoding (str): The target console encoding detected at initialization.
        sanitize_string (method): Sanitize a single string for console encoding.

    Examples:
        Use as a structlog processor:

        ```python
        processor = EncodingSafeProcessor()
        event = {"event": "User said \u2018hello\u2019"}
        result = processor(None, "info", event)
        # result["event"] == "User said 'hello'"
        ```

    Note:
        All sanitization is idempotent - processing already-sanitized strings
        produces identical output.
    """

    # Smart character replacements (preserve meaning)
    # These map common Unicode characters that cause issues on cp1252
    # to their closest ASCII equivalents
    REPLACEMENTS: dict[str, str] = {
        "\u2018": "'",  # Left single quote → apostrophe
        "\u2019": "'",  # Right single quote → apostrophe
        "\u201c": '"',  # Left double quote → quotation mark
        "\u201d": '"',  # Right double quote → quotation mark
        "\u2011": "-",  # Non-breaking hyphen → hyphen-minus
        "\u2013": "-",  # En dash → hyphen-minus
        "\u2014": "--",  # Em dash → double hyphen
        "\u2026": "...",  # Horizontal ellipsis → three periods
        "\u00a0": " ",  # Non-breaking space → regular space
    }

    def __init__(self) -> None:
        """Initialize the processor with console encoding detection.

        Detects the console encoding from sys.stdout.encoding, falling back
        to UTF-8 if detection fails (e.g., when stdout is redirected or
        unavailable).

        Note:
            Console encoding is cached at initialization time to avoid
            repeated attribute lookups during log processing.
        """
        self.encoding: str = getattr(sys.stdout, "encoding", None) or "utf-8"

    def __call__(
        self,
        logger: Any,
        method_name: str,
        event_dict: MutableMapping[str, Any],
    ) -> MutableMapping[str, Any]:
        """Process event dictionary, sanitizing all string values.

        Implements the structlog processor protocol. Recursively sanitizes
        all string values in the event dictionary, including nested dicts
        and lists.

        Args:
            logger: The wrapped logger instance (unused but required by
                protocol).
            method_name: Name of the log method called (e.g., "info", "debug").
                Unused but required by protocol.
            event_dict: Mutable mapping of log event data to sanitize.

        Returns:
            The event_dict with all string values sanitized for console
            encoding compatibility.

        Note:
            Original event_dict is not modified; a new dict is returned
            with sanitized values.
        """
        return self._sanitize_dict(event_dict)

    def sanitize_string(self, s: str) -> str:
        """Sanitize a single string for console encoding.

        Applies smart character replacements first to preserve meaning,
        then uses encode/decode with 'replace' error handler for any
        remaining unencodable characters.

        Args:
            s (str): The string to sanitize.

        Returns:
            A string that is guaranteed to be encodable to the console
            encoding without raising UnicodeEncodeError.

        Examples:
            Sanitize a string with smart quotes:

            ```python
            processor = EncodingSafeProcessor()
            result = processor.sanitize_string("User said \u2018hello\u2019")
            assert result == "User said 'hello'"
            ```

        Note:
            Order of operations: smart replacements run first to preserve
            semantic meaning, then encode/decode handles remaining characters.
        """
        # Apply smart replacements first (preserve meaning)
        for char, replacement in self.REPLACEMENTS.items():
            s = s.replace(char, replacement)

        # Encode/decode with replace for any remaining unencodable chars
        # This handles any Unicode characters not in our explicit mapping
        return s.encode(self.encoding, errors="replace").decode(self.encoding)

    def _sanitize_value(self, value: Any) -> Any:
        """Recursively sanitize any value.

        Handles strings, dicts, lists, and tuples. Other types pass through
        unchanged.

        Args:
            value: The value to sanitize.

        Returns:
            The sanitized value, with the same type as the input (except
            strings may have different content).
        """
        if isinstance(value, str):
            return self.sanitize_string(value)
        elif isinstance(value, MutableMapping):
            return self._sanitize_dict(value)
        elif isinstance(value, (list, tuple)):
            # Preserve the original type (list or tuple)
            return type(value)(self._sanitize_value(v) for v in value)
        # Non-string, non-collection types pass through unchanged
        return value

    def _sanitize_dict(self, d: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
        """Sanitize all values in a dictionary.

        Args:
            d: The dictionary or mutable mapping to sanitize.

        Returns:
            A new dictionary with all string values sanitized.
        """
        return {k: self._sanitize_value(v) for k, v in d.items()}

@dataclass(frozen=True)
class SchemaValidationResult:
    """Result of validating schema text.

    This dataclass contains both the deserialized Pydantic model class
    and metadata about the schema structure.

    Attributes:
        schema_class (type[BaseModel]): The deserialized Pydantic BaseModel subclass.
        class_name (str): Name of the class found in the schema text.
        field_count (int): Number of fields defined in the schema.
        field_names (tuple[str, ...]): Tuple of field names in the schema.

    Examples:
        Validating schema text and inspecting results:

        ```python
        result = validate_schema_text(schema_text)
        print(f"Class: {result.class_name}")
        print(f"Fields: {result.field_names}")
        instance = result.schema_class(name="test", value=42)
        ```

        Creating a result directly (typically done by validate_schema_text):

        ```python
        result = SchemaValidationResult(
            schema_class=MySchema,
            class_name="MySchema",
            field_count=2,
            field_names=("name", "value"),
        )
        ```
    """

    schema_class: type[BaseModel]
    class_name: str
    field_count: int
    field_names: tuple[str, ...]

class StateGuard:
    r"""Validates and repairs mutated component_text to preserve ADK state tokens.

    StateGuard ensures that required state injection tokens are preserved
    during component_text evolution, and escapes unauthorized new tokens introduced
    by reflection. Supports simple tokens ({name}), prefixed tokens ({app:settings}),
    optional tokens ({name?}), and combined formats ({app:config?}).

    Attributes:
        required_tokens (list[str]): List of tokens that must always be present,
            including braces (e.g., ["{user_id}", "{app:settings}", "{name?}"]).
        repair_missing (bool): Whether to re-append missing tokens. Defaults to True.
        escape_unauthorized (bool): Whether to escape new unauthorized tokens.
            Defaults to True.
        _token_pattern (re.Pattern[str]): Compiled regex for token detection (private).

    Examples:
        Basic usage with token repair:

        ```python
        guard = StateGuard(required_tokens=["{user_id}", "{context}"])
        original = "Hello {user_id}, context: {context}"
        mutated = "Hello {user_id}, welcome!"
        result = guard.validate(original, mutated)
        # result == "Hello {user_id}, welcome!\n\n{context}"
        ```

        Escaping unauthorized tokens:

        ```python
        guard = StateGuard(required_tokens=["{user_id}"])
        original = "Process for {user_id}"
        mutated = "Process for {user_id} with {malicious}"
        result = guard.validate(original, mutated)
        # result == "Process for {user_id} with {{malicious}}"
        ```

    Note:
        All validation logic is stateless and operates on string inputs only.
        No external dependencies or I/O operations are performed.

        Supported token formats:
            - Simple tokens: {name}, {user_id}, {context}
            - Prefixed tokens: {app:settings}, {user:api_key}, {temp:session}
            - Optional tokens: {name?}, {user_id?}
            - Combined: {app:config?}, {user:pref?}
    """

    def __init__(
        self,
        required_tokens: list[str] | None = None,
        repair_missing: bool = True,
        escape_unauthorized: bool = True,
    ) -> None:
        """Initialize StateGuard with configuration.

        Args:
            required_tokens: List of tokens that must be preserved,
                including braces (e.g., ["{user_id}", "{context}"]).
                Defaults to empty list.
            repair_missing: If True, re-append missing required tokens.
                Defaults to True.
            escape_unauthorized: If True, escape new unauthorized tokens.
                Defaults to True.

        Note:
            Configuration determines which tokens are protected and which
            behaviors are enabled. Both repair and escape are enabled by default
            for maximum safety.
        """
        self.required_tokens = required_tokens or []
        self.repair_missing = repair_missing
        self.escape_unauthorized = escape_unauthorized
        self._token_pattern = re.compile(r"(?<!\{)\{(\w+(?::\w+)?(?:\?)?)\}(?!\})")

    def _extract_tokens(self, text: str) -> set[str]:
        r"""Extract token names from text using regex.

        Args:
            text: Text to extract tokens from.

        Returns:
            Set of token names (without braces), including full token content
            for prefixed and optional tokens (e.g., "app:settings", "name?").

        Note:
            Scans text using the regex pattern
            `(?<!\{)\{(\w+(?::\w+)?(?:\?)?)\}(?!\})`.

            This pattern matches:
            - Simple tokens: {name} → "name"
            - Prefixed tokens: {app:settings} → "app:settings"
            - Optional tokens: {name?} → "name?"
            - Combined: {app:config?} → "app:config?"

            The negative lookbehind `(?<!\{)` and lookahead `(?!\})` ensure
            already-escaped tokens like `{{token}}` are NOT matched.

            Artifact references (e.g., {artifact.name}) are not matched as they
            contain dots and have different semantics.
        """
        matches = self._token_pattern.findall(text)
        return set(matches)

    def get_validation_summary(
        self, original: str, mutated: str
    ) -> tuple[list[str], list[str]]:
        """Summarize missing token repairs and unauthorized token escapes.

        Args:
            original: The component_text before mutation (reference for tokens).
            mutated: The component_text after mutation (pre-validation).

        Returns:
            Tuple of (repaired_tokens, escaped_tokens) as token strings with braces.

        Examples:
            Summarize repairs and escapes before validation:

            ```python
            guard = StateGuard(required_tokens=["{user_id}"])
            repaired, escaped = guard.get_validation_summary(
                "Hello {user_id}",
                "Hello {user_id} {malicious}",
            )
            # repaired == []
            # escaped == ["{malicious}"]
            ```

        Note:
            Outputs what validate() would repair or escape given the
            current configuration and inputs, using the same token detection
            logic as validate(), without modifying the component_text.
        """
        original_tokens = self._extract_tokens(original)
        mutated_tokens = self._extract_tokens(mutated)
        required_token_names = {token.strip("{}") for token in self.required_tokens}

        repaired_tokens: list[str] = []
        if self.repair_missing:
            missing_required = (original_tokens & required_token_names) - mutated_tokens
            repaired_tokens = [f"{{{token}}}" for token in sorted(missing_required)]

        escaped_tokens: list[str] = []
        if self.escape_unauthorized:
            new_tokens = mutated_tokens - original_tokens
            unauthorized_tokens = new_tokens - required_token_names
            escaped_tokens = [f"{{{token}}}" for token in sorted(unauthorized_tokens)]

        return repaired_tokens, escaped_tokens

    def validate(self, original: str, mutated: str) -> str:
        r"""Validate and repair mutated component_text.

        Compares the original and mutated component_text to:
        1. Re-append missing required tokens (if `repair_missing=True`)
        2. Escape unauthorized new tokens (if `escape_unauthorized=True`)

        Args:
            original: The component_text before mutation (reference for tokens).
                Used to determine which tokens were present initially.
            mutated: The component_text after mutation (to be validated).
                This is the component_text that may have missing or unauthorized tokens.

        Returns:
            The mutated component_text with repairs and escapes applied.
            Missing required tokens are appended at the end with `\n\n{token}`.
            Unauthorized new tokens are escaped by doubling braces: `{token}` → `{{token}}`.

        Examples:
            Repair missing token:

            ```python
            guard = StateGuard(required_tokens=["{user_id}"])
            result = guard.validate("Hello {user_id}", "Hello")
            # result == "Hello\n\n{user_id}"
            ```

            Escape unauthorized token:

            ```python
            guard = StateGuard(required_tokens=["{user_id}"])
            result = guard.validate(
                "Process {user_id}", "Process {user_id} {malicious}"
            )
            # result == "Process {user_id} {{malicious}}"
            ```

        Note:
            Only tokens present in both the original component_text and the
            required_tokens list are eligible for repair. New tokens not in
            required_tokens are escaped by default.
        """
        result = mutated

        # Extract tokens from both component_text values
        original_tokens = self._extract_tokens(original)
        mutated_tokens = self._extract_tokens(mutated)

        # Normalize required_tokens (strip braces for comparison)
        required_token_names = {token.strip("{}") for token in self.required_tokens}

        # Repair missing tokens
        if self.repair_missing:
            # Find tokens in original AND required_tokens AND missing from mutated
            missing_required = (original_tokens & required_token_names) - mutated_tokens

            # Append missing tokens
            for token_name in sorted(missing_required):
                full_token = f"{{{token_name}}}"
                result += f"\n\n{full_token}"

        # Escape unauthorized new tokens
        if self.escape_unauthorized:
            # Find tokens that are new (in mutated but not in original)
            new_tokens = mutated_tokens - original_tokens

            # Only escape tokens that are NOT in required_tokens
            unauthorized_tokens = new_tokens - required_token_names

            # Escape each unauthorized token
            for token_name in unauthorized_tokens:
                full_token = f"{{{token_name}}}"
                escaped_token = f"{{{{{token_name}}}}}"
                result = result.replace(full_token, escaped_token)

        return result