from enum import Enum
from typing import Literal, TypeVar, Generic, Annotated
from pydantic import BaseModel, Field

T = TypeVar("T", bound=float | str)

class StageData(BaseModel, Generic[T]):
    stage_no: int = Field(...)
    base: T = Field(...)

class AStageData(StageData[T], BaseModel, Generic[T]):
    type: Literal["A"] = Field(
        ...
    )
    a1: T = Field(...)
    a2: T = Field(...)

class BStageData(StageData[T], BaseModel, Generic[T]):
    type: Literal["B"] = Field(
        ...
    )
    b1: T = Field(...)
    b2: T = Field(...)

type StageDataBase[T] = Annotated[
    AStageData[T] | BStageData[T], Field(discriminator="type")
]


class StageSeries(BaseModel, Generic[T]):
    stage_no: int = Field(...)
    base: list[T] = Field(
        []
    )

    def append_value(self, stage_data: StageDataBase[T]) -> None:
        self.base.append(stage_data.base)

class AStageSeries(StageSeries[T], Generic[T]):
    type: Literal["A"] = Field(
        ...
    )
    a1: list[T] = Field(
        []
    )
    a2: list[T] = Field(
        []
    )

    def append_value(self, stage_data: AStageData[T]) -> None:
        super().append_value(stage_data)
        self.a1.append(stage_data.a1)
        self.a2.append(stage_data.a2)

class BStageSeries(StageSeries[T], Generic[T]):
    type: Literal["B"] = Field(
        ...
    )
    b1: list[T] = Field(
        []
    )
    b2: list[T] = Field(
        []
    )

    def append_value(self, stage_data: BStageData[T]) -> None:
        super().append_value(stage_data)
        self.b1.append(stage_data.b1)
        self.b2.append(stage_data.b2)

type StageSeriesBase[T] = Annotated[
    AStageSeries[T] | BStageSeries[T], Field(discriminator="type")
]

class Series(BaseModel, Generic[T]):
    stages: list[StageSeriesBase[T]] = Field(
        []
    )

    def append_value(self, stages: list[StageDataBase[T]]) -> None:
        for stage_data in stages:
            stage = self.get_stage(stage_data)
            stage.append_value(stage_data)

    def get_stage(self, stage_data: StageDataBase[T]) -> StageSeriesBase[T]:
        # this throws "Parameterized generics cannot be used with class or instance checks"
        # but this is the only possible way to instantiate the correct generic type class
        # google didn't provide any meaningful info to that mypy error either
        try:
            stage = next(s for s in self.stages if s.stage_no == stage_data.stage_no)
        except StopIteration:
            if isinstance(stage_data, AStageData[float]):
                stage = AStageSeries[float](
                    type="A", stage_no=stage_data.stage_no
                )
            elif isinstance(stage_data, AStageData[str]):
                stage = AStageSeries[str](
                    type="A", stage_no=stage_data.stage_no
                )
            elif isinstance(trend_stage, BStageData[float]):
                stage = BStageSeries[float](
                    type="B", stage_no=stage_data.stage_no
                )
            elif isinstance(trend_stage, BStageData[std]):
                stage = BlStageSeries[str](
                    type="B", stage_no=stage_data.stage_no
                )
            else:
                raise ValueError(f"Unsupported type: {type(stage_data)}") from None

            self.stages.append(stage)
        return stage

g = Series[float]()
g.append_value(
    stages = [
        AStageData[float](
            type="A",
            stage_no=1,
            base=1,
            a1=2,
            a2=3,
        )
    ]
)
g.append_value(
    stages = [
        AStageData[float](
            type="A",
            stage_no=1,
            base=1.1,
            a1=2.1,
            a2=3.1,
        )
    ]
)

g.model_dump_json()