"""Type stubs for variant type hints."""

from typing import TypeVar, Protocol, Any
from typing_extensions import ParamSpec
from pydantic import BaseModel

T_co = TypeVar("T_co", bound=BaseModel, covariant=True)
P = ParamSpec("P")


class VariantClass(Protocol[T_co]):
    """Protocol for a class with dynamically created variants.

    This helps type checkers understand variant attributes.

    Example:
        @variants(basic_variant_pipeline('Input'))
        class User(BaseModel):
            name: str

        # Now type checkers understand:
        User.Input  # type: type[User]
    """

    # Common variant names - add more as needed
    Input: type[Any]
    Output: type[Any]
    Update: type[Any]
    Create: type[Any]
    Public: type[Any]
    API: type[Any]


def cast_variant_class(cls: type[T]) -> type[T]:
    """Helper to cast a variants-decorated class for type checking.

    Use this to help Pylance understand variant attributes:

    Example:
        User = cast_variant_class(User)
        user_input = User.Input(name='John')  # Now properly typed!
    """
    return cls
