from dataclasses import dataclass
from typing import Optional, Union, List

from datagen.api.assets import (
    Human,
    Camera,
    Glasses,
    Mask,
    Background,
    Light,
    HumanDatapoint,
    Accessories,
)


@dataclass
class HumanDatapointBuilder:
    human: Human
    camera: Camera
    glasses: Optional[Glasses]
    mask: Optional[Mask]
    background: Optional[Background]
    lights: Optional[List[Light]]

    def get_basic_datapoint(self) -> HumanDatapoint:
        return HumanDatapoint(human=self.human.copy(), camera=self.camera.copy())

    def get_accessories(self) -> Union[Accessories, None]:
        if self.glasses is None and self.mask is None:
            return None
        else:
            accessories = Accessories()
            if self.glasses is not None:
                accessories.glasses = self.glasses.copy()
            if self.mask is not None:
                accessories.mask = self.mask.copy()
            return accessories

    def get_background(self) -> Union[Background, None]:
        if self.background is not None:
            return self.background.copy()
        else:
            return None

    def get_lights(self) -> Union[List[Light], None]:
        if self.lights is not None and len(self.lights) > 0:
            return [l_.copy() for l_ in self.lights]
        else:
            return None
