# Copyright (c) 2026 Luigi Paone <ppc.luigi.paone@gmail.com>
# SPDX-License-Identifier: BSD-3-Clause

"""Helper functions for Pydantic model operations."""

from typing import Any, TypeVar, Type
from pydantic import BaseModel, ValidationError

T = TypeVar("T", bound=BaseModel)


def find_all_submodels(data: dict[str, Any] | list[Any], model_class: Type[T]) -> list[T]:
    """Find all occurrences of a submodel in nested JSON data.

    Recursively searches through nested dictionaries and lists to find
    all data structures that can be validated as the given Pydantic model.

    Args:
        data: The JSON data structure (dict or list) to search.
        model_class: The Pydantic model class to search for.

    Returns:
        A list of all matching model instances found in the data.

    Example:
        >>> class Inner(BaseModel):
        ...     x: int
        ...     y: str
        >>> data = {"items": [{"x": 1, "y": "a"}, {"x": 2, "y": "b"}]}
        >>> matches = find_all_submodels(data, Inner)
        >>> len(matches)
        2
    """
    results: list[T] = []
    if isinstance(data, dict):
        try:
            results.append(model_class.model_validate(data))
        except ValidationError:
            pass
        # Always recurse into nested values to find all matches
        for value in data.values():
            results.extend(find_all_submodels(value, model_class))
    elif isinstance(data, list):
        for item in data:
            results.extend(find_all_submodels(item, model_class))
    return results


def find_submodel(data: dict[str, Any] | list[Any], model_class: Type[T]) -> T | None:
    """Find the first occurrence of a submodel in nested JSON data.

    Recursively searches through nested dictionaries and lists to find
    the first data structure that can be validated as the given Pydantic model.

    Args:
        data: The JSON data structure (dict or list) to search.
        model_class: The Pydantic model class to search for.

    Returns:
        The first matching model instance, or None if not found.

    Example:
        >>> class Inner(BaseModel):
        ...     x: int
        >>> data = {"nested": {"x": 42}}
        >>> match = find_submodel(data, Inner)
        >>> match.x
        42
    """
    if isinstance(data, dict):
        try:
            return model_class.model_validate(data)
        except ValidationError:
            pass
        for value in data.values():
            result = find_submodel(value, model_class)
            if result:
                return result
    elif isinstance(data, list):
        for item in data:
            result = find_submodel(item, model_class)
            if result:
                return result
    return None
