"""Flatten ``allOf``, ``oneOf``, and ``anyOf`` constructs in OpenAPI schemas."""

from __future__ import annotations

from typing import Any

from lapis_spec.models.openapi import OpenAPIDict


def flatten_schemas(doc: OpenAPIDict) -> OpenAPIDict:
    """Walk the entire document and flatten composition keywords in-place.

    Returns the same *doc* object (mutated).
    """
    _walk(doc)
    return doc


def flatten_schema(schema: dict[str, Any]) -> dict[str, Any]:
    """Flatten a single schema dict and return it."""
    _walk(schema)
    return schema


def _walk(node: Any) -> None:
    if isinstance(node, dict):
        _flatten_node(node)
        for value in node.values():
            _walk(value)
    elif isinstance(node, list):
        for item in node:
            _walk(item)


def _flatten_node(node: dict[str, Any]) -> None:
    if "allOf" in node:
        _flatten_all_of(node)
    if "oneOf" in node:
        _flatten_one_of(node)
    if "anyOf" in node:
        _flatten_any_of(node)


def _flatten_all_of(node: dict[str, Any]) -> None:
    """Merge all ``allOf`` sub-schemas into the parent node."""
    all_of = node.pop("allOf", [])
    if not isinstance(all_of, list):
        return
    merged_props: dict[str, Any] = {}
    merged_required: list[str] = []
    extra: dict[str, Any] = {}
    for sub in all_of:
        if not isinstance(sub, dict):
            continue
        _walk(sub)
        props = sub.get("properties", {})
        if isinstance(props, dict):
            merged_props.update(props)
        req = sub.get("required", [])
        if isinstance(req, list):
            merged_required.extend(req)
        for k, v in sub.items():
            if k not in ("properties", "required", "type"):
                extra.setdefault(k, v)

    existing_props = node.get("properties", {})
    if isinstance(existing_props, dict):
        merged_props.update(existing_props)

    existing_req = node.get("required", [])
    if isinstance(existing_req, list):
        merged_required.extend(existing_req)

    if merged_props:
        node["properties"] = merged_props
    if merged_required:
        node["required"] = list(dict.fromkeys(merged_required))
    node.setdefault("type", "object")
    for k, v in extra.items():
        node.setdefault(k, v)


def _flatten_one_of(node: dict[str, Any]) -> None:
    """Simplify ``oneOf``: pick the first variant or fall back to ``any``."""
    variants = node.pop("oneOf", [])
    if not isinstance(variants, list) or not variants:
        return
    for sub in variants:
        if isinstance(sub, dict):
            _walk(sub)
    if len(variants) == 1 and isinstance(variants[0], dict):
        for k, v in variants[0].items():
            node.setdefault(k, v)
    # For complex unions, leave the node as-is (will map to ``any``).


def _flatten_any_of(node: dict[str, Any]) -> None:
    """Simplify ``anyOf``: pick the first variant or fall back to ``any``."""
    variants = node.pop("anyOf", [])
    if not isinstance(variants, list) or not variants:
        return
    for sub in variants:
        if isinstance(sub, dict):
            _walk(sub)
    if len(variants) == 1 and isinstance(variants[0], dict):
        for k, v in variants[0].items():
            node.setdefault(k, v)
