app/utils/schema_utils.py

from __future__ import annotations

import json
import logging
from typing import Any, Dict, Optional, Tuple

from jsonschema import Draft202012Validator, SchemaError, ValidationError

logger = logging.getLogger(__name__)


def validate_response_format(body: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    """Validate the response_format field from the request body.

    Accepts:
      - {"type": "json_object"}
      - {"type": "json_schema", "json_schema": {"name": "...", "strict": true, "schema": {...}}}

    Returns the validated dict or None if response_format is absent.
    Raises ValueError with a human-readable message on invalid input.
    """
    rf = body.get("response_format")
    if rf is None:
        return None

    if not isinstance(rf, dict):
        raise ValueError("response_format must be an object")

    rf_type = rf.get("type")
    if rf_type not in ("json_object", "json_schema"):
        raise ValueError(
            f"response_format.type must be 'json_object' or 'json_schema', got '{rf_type}'"
        )

    if rf_type == "json_object":
        return rf

    # json_schema path - validate the nested structure
    js = rf.get("json_schema")
    if not isinstance(js, dict):
        raise ValueError("response_format.json_schema must be an object")

    name = js.get("name")
    if not isinstance(name, str) or not name.strip():
        raise ValueError("json_schema.name must be a non-empty string")

    schema = js.get("schema")
    if not isinstance(schema, dict):
        raise ValueError("json_schema.schema must be a JSON Schema object")

    strict = js.get("strict")
    if strict is not None and not isinstance(strict, bool):
        raise ValueError("json_schema.strict must be a boolean")

    # Validate the schema itself is a valid JSON Schema
    try:
        Draft202012Validator.check_schema(schema)
    except (ValidationError, SchemaError) as e:
        raise ValueError(f"json_schema.schema is not a valid JSON Schema: {e.message}")

    return rf


def generate_schema_prompt(
    schema: Dict[str, Any], name: str = "response"
) -> str:
    """Generate a system prompt that instructs the LLM to produce output
    conforming to the given JSON Schema."""
    schema_json = json.dumps(schema, indent=2)
    field_instructions = _build_field_instructions(schema)

    return (
        f'You MUST respond with a single valid JSON object that conforms EXACTLY '
        f'to this JSON Schema named "{name}":\n\n'
        f"```json\n{schema_json}\n```\n\n"
        "CRITICAL RULES:\n"
        "1. Your entire response must be ONLY a JSON object. No text before or after.\n"
        "2. Do NOT wrap the JSON in markdown code fences or any other formatting.\n"
        "3. Every required field MUST be present in your response.\n"
        "4. Use ONLY the types specified in the schema (string, number, integer, boolean, array, object, null).\n"
        "5. Do NOT include any fields that are not defined in the schema properties.\n"
        '6. For "enum" fields, use EXACTLY one of the specified values.\n'
        '7. For "const" fields, use the exact specified value.\n'
        "8. For nested objects, follow the sub-schema recursively.\n\n"
        f"{field_instructions}\n\n"
        "Respond with ONLY the JSON object - no explanation, no markdown, no code fences."
    )


def _build_field_instructions(
    schema: Dict[str, Any], prefix: str = ""
) -> str:
    """Recursively build human-readable field instructions from a schema."""
    lines: list[str] = []
    props = schema.get("properties", {})
    required = set(schema.get("required", []))

    for field_name, field_schema in props.items():
        full_name = f"{prefix}{field_name}" if prefix else field_name
        field_type = field_schema.get("type", "any")
        is_required = field_name in required

        status = "REQUIRED" if is_required else "optional"
        desc = field_schema.get("description", "")

        if field_type == "string" and "enum" in field_schema:
            enum_vals = ", ".join(f'"{v}"' for v in field_schema["enum"])
            line = f'- Field "{full_name}" ({status}): Must be one of [{enum_vals}].'
        elif field_type == "string" and "format" in field_schema:
            fmt = field_schema["format"]
            line = f'- Field "{full_name}" ({status}): type=string, format="{fmt}".'
        elif field_type == "array" and "items" in field_schema:
            items = field_schema["items"]
            item_type = items.get("type", "any")
            line = f'- Field "{full_name}" ({status}): type=array of {item_type}.'
        elif field_type == "object" and "properties" in field_schema:
            nested = _build_field_instructions(field_schema, prefix=f"{full_name}.")
            if nested:
                lines.append(f'- Field "{full_name}" ({status}): nested object with fields:')
                lines.append(nested)
            continue
        elif field_type == "null":
            line = f'- Field "{full_name}" ({status}): can be null.'
        else:
            line = f'- Field "{full_name}" ({status}): type={field_type}.'

        if desc:
            line = line.rstrip(".") + f". {desc}"
        lines.append(line)

    if not schema.get("additionalProperties", True):
        lines.append("- Do NOT include any additional properties not listed above.")

    return "\n".join(lines) if lines else ""


def validate_against_schema(
    data: Any, schema: Dict[str, Any]
) -> Tuple[bool, Any, Optional[str]]:
    """Validate data against a JSON Schema.

    Returns:
        (is_valid, data, error_message_or_None)
    """
    try:
        validator = Draft202012Validator(schema)
        validator.validate(data)
        return True, data, None
    except ValidationError as e:
        path = (
            ".".join(str(p) for p in e.absolute_path)
            if e.absolute_path
            else "(root)"
        )
        msg = f"Validation error at '{path}': {e.message}"
        return False, data, msg