"""Utilities for post-processing OpenAPI specs generated by FastAPI/Pydantic.

Pydantic v2 inlines discriminated unions (``oneOf`` + ``discriminator``) at
every usage site.  When the union is recursive (e.g. ``Container.children``
references ``ResourceDataModels`` which includes ``Container``), documentation
renderers like Redoc and Swagger UI choke on the circular inline schemas.

The single public function :func:`deduplicate_discriminated_unions` fixes this
by extracting repeated inline unions into named ``$ref`` schemas.
"""

from collections import defaultdict


def _collect_inline_unions(spec: dict) -> list[tuple]:
    """Walk *spec* and return every inline discriminated union found.

    Each entry is ``(parent_obj, key_in_parent, union_obj, fingerprint,
    first_ref_name)`` where *fingerprint* is the sorted tuple of ``$ref``
    targets and *first_ref_name* is the first variant name (typically the
    base type).
    """
    unions: list[tuple] = []

    def _walk(obj: object, parent: object = None, parent_key: object = None) -> None:
        if isinstance(obj, dict):
            if "oneOf" in obj and "discriminator" in obj:
                refs = tuple(item["$ref"] for item in obj["oneOf"] if "$ref" in item)
                if len(refs) == len(obj["oneOf"]):
                    fingerprint = tuple(sorted(refs))
                    first_ref_name = refs[0].rsplit("/", 1)[-1]
                    unions.append(
                        (parent, parent_key, obj, fingerprint, first_ref_name)
                    )
            for key, value in obj.items():
                _walk(value, obj, key)
        elif isinstance(obj, list):
            for idx, value in enumerate(obj):
                _walk(value, obj, idx)

    _walk(spec)
    return unions


def _extract_named_schemas(schemas: dict, groups: dict[tuple, list]) -> None:
    """Create named schemas for duplicated union groups and replace inlines."""
    for _fingerprint, occurrences in groups.items():
        if len(occurrences) < 2:
            continue

        # Use the first $ref name (typically the base type) for a readable name.
        base = occurrences[0][3]
        schema_name = f"Any{base}Model"
        suffix = 2
        original_name = schema_name
        while schema_name in schemas:
            schema_name = f"{original_name}{suffix}"
            suffix += 1

        # The canonical definition is a copy of the first occurrence.
        _, _, canonical, _ = occurrences[0]
        schemas[schema_name] = dict(canonical)

        # Replace every inline occurrence with a $ref.
        ref_value = {"$ref": f"#/components/schemas/{schema_name}"}
        for parent, key, _, _ in occurrences:
            if parent is not None:
                parent[key] = ref_value


def deduplicate_discriminated_unions(spec: dict) -> dict:
    """Extract repeated inline discriminated unions into named schemas.

    Walks an OpenAPI spec dict and:

    1. Finds all inline discriminated unions (objects with both ``oneOf`` and
       ``discriminator`` keys whose ``oneOf`` entries are all ``$ref`` s).
    2. Groups them by fingerprint (the sorted set of ``$ref`` targets).
    3. For groups with 2+ occurrences, creates a named schema under
       ``components/schemas`` and replaces every inline occurrence with a
       single ``$ref``.

    The spec is mutated **in place** and also returned for convenience.
    """
    schemas = spec.setdefault("components", {}).setdefault("schemas", {})

    unions = _collect_inline_unions(spec)

    groups: dict[tuple, list] = defaultdict(list)
    for parent, key, union_obj, fp, first_name in unions:
        groups[fp].append((parent, key, union_obj, first_name))

    _extract_named_schemas(schemas, groups)

    return spec
