from jsonschema import validate, ValidationError
from loguru import logger

A2A_SCHEMA = {
    "type": "object",
    "required": ["sender_agent", "payload", "timestamp"],
    "properties": {
        "sender_agent": {"enum": ["input_guard", "skin_analysis", "rag_retrieval",
                                   "ingredient_safety", "reasoning", "output_validation"]},
        "payload": {"type": "object"},
        "timestamp": {"type": "string"}
    },
    "additionalProperties": False
}

def validate_agent_message(msg: dict) -> dict:
    """Validate all inter-agent A2A messages against schema."""
    try:
        validate(instance=msg, schema=A2A_SCHEMA)
        return msg
    except ValidationError as e:
        logger.error(f"A2A injection attempt blocked: {e.message}")
        raise ValueError(f"Invalid agent message: {e.message}")