# Usage:
#   python -m venv .venv
#   . .venv/bin/activate  # or . .\.venv\Scripts\Activate.ps1
#   pip install git+https://github.com/pydantic/pydantic.git#egg=pydantic
#   pip freeze  # confirm pydantic commit is current tip of main
#   python union_root_discriminator_issues_pydv2b3.py
# pylint: disable=all

from typing import Literal
from pydantic import RootModel, BaseModel, Field


class IdBase(BaseModel):
    kind: Literal["id"]


class IdVariant1(IdBase):
    id_str: str


class IdVariant2(IdBase):
    id_int: int


class Id(RootModel):
    root: IdVariant1 | IdVariant2


class Name(BaseModel):
    name: str
    kind: Literal["name"]


# This doesn't work, yielding this error:
#   TypeError: 'definition-ref' is not a valid discriminated union variant; should be a `BaseModel` or `dataclass`
class Item(RootModel):
    root: Name | Id = Field(..., discriminator="kind")


class Items(RootModel):
    root: list[Item]


# This does work though
# class Item(RootModel):
#     root: Name | Id
#
#
# class Items(RootModel):
#     root: list[Item] = Field(..., discriminator="kind")


def main():
    items = Items.model_validate(
        [
            {"id_str": "id1", "kind": "id"},
            {"id_int": 2, "kind": "id"},
            {"name": "number-three", "kind": "name"},
        ]
    )
    print(items)
    print(items.model_json_schema())
    print(items.model_dump_json())


if __name__ == "__main__":
    main()