enviPy-bayer/epapi/utils/schema_transformers.py

"""
Schema transformation utilities for converting Pydantic models to RJSF format.

This module provides functions to extract UI configuration from Pydantic models
and transform them into React JSON Schema Form (RJSF) compatible format.
"""

from typing import Type, Optional, Any

import jsonref
from pydantic import BaseModel

from envipy_additional_information.ui_config import UIConfig
from envipy_additional_information import registry


def extract_groups(model_cls: Type[BaseModel]) -> list[str]:
    """
    Extract groups from registry-stored group information.

    Args:
        model_cls: The model class

    Returns:
        List of group names the model belongs to
    """
    return registry.get_groups(model_cls)


def extract_ui_metadata(model_cls: Type[BaseModel]) -> dict[str, Any]:
    """
    Extract model-level UI metadata from UI class.

    Returns metadata attributes that are NOT UIConfig instances.
    Common metadata includes: unit, description, title.
    """
    metadata: dict[str, Any] = {}

    if not hasattr(model_cls, "UI"):
        return metadata

    ui_class = getattr(model_cls, "UI")

    # Iterate over all attributes in the UI class
    for attr_name in dir(ui_class):
        # Skip private attributes
        if attr_name.startswith("_"):
            continue

        # Get the attribute value
        try:
            attr_value = getattr(ui_class, attr_name)
        except AttributeError:
            continue

        # Skip callables but keep types/classes
        if callable(attr_value) and not isinstance(attr_value, type):
            continue

        # Skip UIConfig instances (these are field-level configs, not metadata)
        # This includes both UIConfig and IntervalConfig
        if isinstance(attr_value, UIConfig):
            continue

        metadata[attr_name] = attr_value

    return metadata


def extract_ui_config_from_model(model_cls: Type[BaseModel]) -> dict[str, Any]:
    """
    Extract UI configuration from model's UI class.

    Returns a dictionary mapping field names to their UI schema configurations.
    Trusts the config classes to handle their own transformation logic.
    """
    ui_configs: dict[str, Any] = {}

    if not hasattr(model_cls, "UI"):
        return ui_configs

    ui_class = getattr(model_cls, "UI")
    schema = model_cls.model_json_schema()
    field_names = schema.get("properties", {}).keys()

    # Extract config for each field
    for field_name in field_names:
        # Skip if UI config doesn't exist for this field (field may be hidden from UI)
        if not hasattr(ui_class, field_name):
            continue

        ui_config = getattr(ui_class, field_name)

        if isinstance(ui_config, UIConfig):
            ui_configs[field_name] = ui_config.to_ui_schema_field()

    return ui_configs


def build_ui_schema(model_cls: Type[BaseModel]) -> dict:
    """Generate RJSF uiSchema from model's UI class."""
    ui_schema = {}

    # Extract field-level UI configs
    field_configs = extract_ui_config_from_model(model_cls)

    for field_name, config in field_configs.items():
        ui_schema[field_name] = config

    return ui_schema


def build_schema(model_cls: Type[BaseModel]) -> dict[str, Any]:
    """
    Build JSON schema from Pydantic model, applying UI metadata.

    Dereferences all $ref pointers to produce fully inlined schema.
    This ensures the frontend receives schemas with enum values and nested
    properties fully resolved, without needing client-side ref resolution.

    Extracts model-level metadata from UI class (title, unit, etc.) and applies
    it to the generated schema. This ensures UI metadata is the single source of truth.
    """
    schema = model_cls.model_json_schema()

    # Dereference $ref pointers (inlines $defs) using jsonref
    # This ensures the frontend receives schemas with enum values and nested
    # properties fully resolved, currently necessary for client-side rendering.
    # FIXME: This is a hack to get the schema to work with alpine schema-form.js replace once we migrate to client-side framework.
    schema = jsonref.replace_refs(schema, proxies=False)

    # Remove $defs section as all refs are now inlined
    if "$defs" in schema:
        del schema["$defs"]

    # Extract and apply UI metadata (title, unit, description, etc.)
    ui_metadata = extract_ui_metadata(model_cls)

    # Apply all metadata consistently as custom properties with x- prefix
    # This ensures consistency and avoids conflicts with standard JSON Schema properties
    for key, value in ui_metadata.items():
        if value is not None:
            schema[f"x-{key}"] = value

    # Set standard title property from UI metadata for JSON Schema compliance
    if "title" in ui_metadata:
        schema["title"] = ui_metadata["title"]
    elif "label" in ui_metadata:
        schema["title"] = ui_metadata["label"]

    return schema


def build_rjsf_output(model_cls: Type[BaseModel], initial_data: Optional[dict] = None) -> dict:
    """
    Main function that returns complete RJSF format.

    Trusts the config classes to handle their own transformation logic.
    No special-case handling - if a config knows how to transform itself, it will.

    Returns:
        dict with keys: schema, uiSchema, formData, groups
    """
    # Build schema with UI metadata applied
    schema = build_schema(model_cls)

    # Build UI schema - config classes handle their own transformation
    ui_schema = build_ui_schema(model_cls)

    # Extract groups from marker interfaces
    groups = extract_groups(model_cls)

    # Use provided initial_data or empty dict
    form_data = initial_data if initial_data is not None else {}

    return {
        "schema": schema,
        "uiSchema": ui_schema,
        "formData": form_data,
        "groups": groups,
    }