"""
Type-introspection helpers shared across backend adapters.

Both mizan-django and mizan-fastapi need to walk @client-decorated function
annotations the same way during schema export. Drift here breaks AFI parity,
so the helpers live in core.
"""

from __future__ import annotations

import types
from typing import Any, Union, get_args, get_origin

from pydantic import BaseModel


__all__ = [
    "extract_optional",
    "extract_list_element",
    "is_structured_output",
    "types_match_for_merge",
]


def extract_optional(annotation: Any) -> tuple[Any, bool]:
    """Unwrap `Optional[T]` / `T | None`.

    Returns `(T, True)` for a union containing exactly one non-None member
    and `None` itself. For anything else, returns `(annotation, False)`.

    Multi-arm unions like `A | B | None` are returned as-is — protocol-level
    discriminated unions aren't supported yet, and silently picking one arm
    would hide that.
    """
    origin = get_origin(annotation)
    if origin is Union or isinstance(annotation, types.UnionType):
        non_none = [a for a in get_args(annotation) if a is not type(None)]
        if len(non_none) == 1:
            return non_none[0], True
    return annotation, False


def extract_list_element(annotation: Any) -> Any | None:
    """If `annotation` is `list[T]` (or sibling container of one), return `T`.

    Recognizes `list`, `tuple`, `set`, `frozenset`. For `tuple[T, ...]` the
    variadic shape is treated as a homogeneous container; heterogeneous
    tuples are not unwrapped.
    """
    origin = get_origin(annotation)
    if origin not in (list, tuple, set, frozenset):
        return None
    args = get_args(annotation)
    if len(args) == 1:
        return args[0]
    if origin is tuple and len(args) == 2 and args[1] is Ellipsis:
        return args[0]
    return None


def is_structured_output(annotation: Any) -> bool:
    """Recognize return types that don't need a `{result: ...}` primitive wrap.

    Matches `BaseModel`, `Optional[BaseModel]` / `BaseModel | None`, and
    container-of-BaseModel (`list[T]`, `tuple[T, ...]`, etc.). Anything else
    (primitives, dicts, raw `Any`) is treated as primitive and gets wrapped
    so it can ride through Pydantic's typed serialization.
    """
    if isinstance(annotation, type) and issubclass(annotation, BaseModel):
        return True
    origin = get_origin(annotation)
    if origin is Union or isinstance(annotation, types.UnionType):
        return any(
            arg is not type(None) and is_structured_output(arg)
            for arg in get_args(annotation)
        )
    if origin in (list, tuple, set, frozenset):
        return any(is_structured_output(arg) for arg in get_args(annotation))
    return False


def types_match_for_merge(slot_type: Any, value_type: Any) -> bool:
    """True if a `value_type` mutation return can splice into a `slot_type` context slot.

    Used by backend dispatch to resolve `@client(merge=ctx)` to a concrete
    function-name slot inside the context bundle. Three shapes match:

    - direct: slot is `T`, value is `T` → replace
    - upsert: slot is `list[T]`, value is `T` → upsert by id
    - list replace: slot is `list[T]`, value is `list[T]`

    `Optional[T]` is unwrapped on both sides before comparison.
    """
    slot_inner, _ = extract_optional(slot_type)
    value_inner, _ = extract_optional(value_type)
    if slot_inner is value_inner:
        return True
    slot_elem = extract_list_element(slot_inner)
    if slot_elem is not None and slot_elem is value_inner:
        return True
    value_elem = extract_list_element(value_inner)
    if slot_elem is not None and value_elem is not None and slot_elem is value_elem:
        return True
    return False
