from typing import Annotated, Any, Literal

from pydantic import BaseModel, Discriminator, Tag


class Pie(BaseModel):
    time_to_cook: int
    num_ingredients: int


class ApplePie(Pie):
    fruit: Literal['apple'] = 'apple'


class PumpkinPie(Pie):
    filling: Literal['pumpkin'] = 'pumpkin'


def get_discriminator_value(v: Any) -> str:
    print("I'm discriminating!")
    if isinstance(v, dict):
        return v.get('fruit', v.get('filling'))
    return getattr(v, 'fruit', getattr(v, 'filling', None))


# This works
type Dessert = Annotated[
    (Annotated[ApplePie, Tag('apple')] | Annotated[PumpkinPie, Tag('pumpkin')]),
    Discriminator(get_discriminator_value),
]


class ThanksgivingDinner(BaseModel):
    dessert: Dessert


# This doesn't work
# Works if using implicit alias: TaggedApplePie = Annotated[ApplePie, Tag('apple')]
type TaggedApplePie = Annotated[ApplePie, Tag('apple')]


class ThanksgivingDinnerSeconds(BaseModel):
    dessert: Annotated[
        (TaggedApplePie | Annotated[PumpkinPie, Tag('pumpkin')]),
        Discriminator(get_discriminator_value),
    ]