# -*- coding: utf-8 -*-
"""Location: ./plugins/schema_guard/schema_guard.py
Copyright 2025
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

Schema Guard Plugin.
Validates tool args and results against a minimal JSONSchema-like subset.
Supported: type, properties, required. Types: object, string, number, integer, boolean, array.
"""

# Future
from __future__ import annotations

# Standard
from typing import Any, Dict, Optional

# Third-Party
from pydantic import BaseModel

# Third-Party
from cpex.framework import (
    Plugin,
    PluginConfig,
    PluginContext,
    PluginViolation,
    ToolPostInvokePayload,
    ToolPostInvokeResult,
    ToolPreInvokePayload,
    ToolPreInvokeResult,
)


class SchemaGuardConfig(BaseModel):
    """Configuration for schema validation guard.

    Attributes:
        arg_schemas: Map of tool names to argument schemas.
        result_schemas: Map of tool names to result schemas.
        block_on_violation: Whether to block on validation failures.
    """

    arg_schemas: Optional[Dict[str, Dict[str, Any]]] = None
    result_schemas: Optional[Dict[str, Dict[str, Any]]] = None
    block_on_violation: bool = True


def _is_type(value: Any, typ: str) -> bool:
    """Check if value matches the specified type.

    Args:
        value: Value to check.
        typ: Type name (object, string, number, integer, boolean, array).

    Returns:
        True if value matches the type.
    """
    match typ:
        case "object":
            return isinstance(value, dict)
        case "string":
            return isinstance(value, str)
        case "number":
            return isinstance(value, (int, float))
        case "integer":
            return isinstance(value, int) and not isinstance(value, bool)
        case "boolean":
            return isinstance(value, bool)
        case "array":
            return isinstance(value, list)
    return True


def _validate(data: Any, schema: Dict[str, Any]) -> list[str]:
    """Validate data against a schema.

    Args:
        data: Data to validate.
        schema: JSONSchema-like validation schema.

    Returns:
        List of validation error messages.
    """
    errors: list[str] = []
    s_type = schema.get("type")
    if s_type and not _is_type(data, s_type):
        errors.append(f"Type mismatch: expected {s_type}")
        return errors
    if s_type == "object":
        props = schema.get("properties", {})
        required = schema.get("required", [])
        for key in required:
            if not isinstance(data, dict) or key not in data:
                errors.append(f"Missing required property: {key}")
        if isinstance(data, dict):
            for key, sub in props.items():
                if key in data:
                    errors.extend([f"{key}: {e}" for e in _validate(data[key], sub)])
    if s_type == "array":
        if isinstance(data, list) and "items" in schema:
            for idx, item in enumerate(data):
                errors.extend([f"[{idx}]: {e}" for e in _validate(item, schema["items"])])
    return errors


class SchemaGuardPlugin(Plugin):
    """Validate tool args and results using a simple schema subset."""

    def __init__(self, config: PluginConfig) -> None:
        """Initialize the schema guard plugin.

        Args:
            config: Plugin configuration.
        """
        super().__init__(config)
        self._cfg = SchemaGuardConfig(**(config.config or {}))

    async def tool_pre_invoke(self, payload: ToolPreInvokePayload, context: PluginContext) -> ToolPreInvokeResult:
        """Validate tool arguments before invocation.

        Args:
            payload: Tool invocation payload.
            context: Plugin execution context.

        Returns:
            Result indicating whether arguments pass schema validation.
        """
        schema = (self._cfg.arg_schemas or {}).get(payload.name)
        if not schema:
            return ToolPreInvokeResult(continue_processing=True)
        errors = _validate(payload.args or {}, schema)
        if errors and self._cfg.block_on_violation:
            return ToolPreInvokeResult(
                continue_processing=False,
                violation=PluginViolation(
                    reason="Schema validation failed",
                    description="Arguments do not conform to schema",
                    code="SCHEMA_GUARD_ARGS",
                    details={"errors": errors},
                ),
            )
        return ToolPreInvokeResult(metadata={"schema_errors": errors})

    async def tool_post_invoke(self, payload: ToolPostInvokePayload, context: PluginContext) -> ToolPostInvokeResult:
        """Validate tool result after invocation.

        Args:
            payload: Tool result payload.
            context: Plugin execution context.

        Returns:
            Result indicating whether tool result passes schema validation.
        """
        schema = (self._cfg.result_schemas or {}).get(payload.name)
        if not schema:
            return ToolPostInvokeResult(continue_processing=True)
        errors = _validate(payload.result, schema)
        if errors and self._cfg.block_on_violation:
            return ToolPostInvokeResult(
                continue_processing=False,
                violation=PluginViolation(
                    reason="Schema validation failed",
                    description="Result does not conform to schema",
                    code="SCHEMA_GUARD_RESULT",
                    details={"errors": errors},
                ),
            )
        return ToolPostInvokeResult(metadata={"schema_errors": errors})
