Source code for luna_handlers.classes.schemas.handler

"""
Module contains schemas for handler
"""
from typing import Any, List, Optional, Union
from uuid import UUID

from pydantic import ValidationError, conlist, root_validator

from classes.schemas import types
from classes.schemas.base_schema import BaseSchema, HandlerSettings
from classes.schemas.match_policy import MatchPolicy
from classes.schemas.policies import Policies
from classes.schemas.simple_schemas import BoundingBoxSchema, UrlBboxSchema, UrlWithFaceBbox
from classes.schemas.types import OptionalNotNullable, userConstrainedList
from classes.schemas.validators import checkRequiredImageMimetype, validateOneOfStrParams


class ValidatorHandlerModel(BaseSchema):
    """Handler validator model"""

    # handler account id
    accountId: UUID
    # handler policies
    policies: Policies


class CreateHandlerModel(BaseSchema):
    """Handler model for handler creation"""

    # handler account id
    accountId: UUID
    # handler policies
    policies: Optional[Policies] = None
    # handler description
    description: types.Str128 = ""
    # handler dynamic flag
    isDynamic: bool = False

    @staticmethod
    def _validateDynamicHandlerPolicies(values: Any):
        """
        Validate not allowed policies for dynamic handler and required policies for non-dynamic handler
        Args:
            values: values from request
        Raises:
            ValueError("Dynamic handler can not have any policies.") if dynamic handler flag is True
                                                                                                and policies specified
            ValueError("Non-dynamic handler must have policies.") if dynamic handler flag is False
                                                                                            and policies not specified
        """
        isDynamic = values["isDynamic"]
        isPoliciesExists = values["policies"] is not None

        if isDynamic is isPoliciesExists is True:
            raise ValueError("Dynamic handler can not have any policies.")
        if isDynamic is isPoliciesExists is False:
            raise ValueError("Non-dynamic handler must have policies.")

    @root_validator(pre=False, skip_on_failure=True)
    def validateNPostInit(cls, values):
        """Validate handler, post init"""
        cls._validateDynamicHandlerPolicies(values)
        return values


class CachedHandlerModel(BaseSchema):
    """Cached handler model | contains minimal required fields"""

    # handler account id
    accountId: UUID
    # handler policies
    policies: Optional[Policies] = None
    # handler dynamic flag
    isDynamic: bool = False
    # handler last update time
    lastUpdateTime: types.CustomDatetime

    @staticmethod
    def postInitMatchPolicy(matchPolicies: List[MatchPolicy], accountId: str) -> None:
        """
        Post init matching policy - apply handler account id for matching candidates
        Args:
            matchPolicies: matching policies
            accountId: handler account id
        """
        for matchPolicy in matchPolicies:
            matchPolicy.candidates.accountId = accountId

    @root_validator(pre=False, skip_on_failure=True)
    def postInit(cls, values):
        """Post init"""
        cls.postInitMatchPolicy(values["policies"].matchPolicy, values["accountId"])
        return values


class UrlForInputEvent(UrlBboxSchema):
    """Url schema with detection time and image origin"""

    # sample detect time
    detectTime: types.CustomDatetime = OptionalNotNullable()
    # image origin
    imageOrigin: types.Str256 = OptionalNotNullable()
    # user-defined timestamp relative to something, such as the start of a video
    detectTs: Optional[types.DetectTs] = None


class SampleForInputEvent(BaseSchema):
    """Sample schema"""

    # sample id
    sampleId: UUID
    # sample detect time
    detectTime: types.CustomDatetime = OptionalNotNullable()
    # image origin
    imageOrigin: types.Str256 = OptionalNotNullable()
    # user-defined timestamp relative to something, such as the start of a video
    detectTs: Optional[types.DetectTs] = None


UrlsWithFaceBBox = userConstrainedList(
    UrlWithFaceBbox, minItems=1, maxItemsGetter=lambda: HandlerSettings.receivedImagesLimit
)
UrlsWithFaceAndBodyBBox = userConstrainedList(
    UrlBboxSchema, minItems=1, maxItemsGetter=lambda: HandlerSettings.receivedImagesLimit
)
SamplesIDs = userConstrainedList(UUID, minItems=1, maxItemsGetter=lambda: HandlerSettings.receivedImagesLimit)
UrlsForInputEvent = userConstrainedList(
    UrlForInputEvent, minItems=1, maxItemsGetter=lambda: HandlerSettings.receivedImagesLimit
)
SamplesForInputEvent = userConstrainedList(
    SampleForInputEvent, minItems=1, maxItemsGetter=lambda: HandlerSettings.receivedImagesLimit
)


class FaceInputEstimationsModel(BaseSchema):
    """Model for incoming face estimations: urls, image or samples"""

    # image (base64)
    image: Optional[str]
    # image mimetype
    mimetype: Optional[str]
    # face bounding box list
    faceBoundingBoxes: conlist(BoundingBoxSchema, min_items=1) = OptionalNotNullable()
    # url list
    urls: UrlsWithFaceBBox = OptionalNotNullable()
    # sample ids
    samples: SamplesIDs = OptionalNotNullable()

    @root_validator(skip_on_failure=True)
    def validateHandler(cls, values: dict) -> dict:
        """Validate handler model"""
        validateOneOfStrParams(values, params=["urls", "samples", "image"])
        checkRequiredImageMimetype(values)
        return values


class HandlerInputEstimationsModel(BaseSchema):
    """Handler model for incoming estimations: urls, image(with face and/or body bounding boxes) or samples"""

    # image (base64)
    image: Optional[str]
    # image mimetype
    mimetype: Optional[str]
    # face bounding box list
    faceBoundingBoxes: conlist(BoundingBoxSchema, min_items=1) = OptionalNotNullable()
    # body bounding box list
    bodyBoundingBoxes: conlist(BoundingBoxSchema, min_items=1) = OptionalNotNullable()
    # list of urls with detection time
    urls: UrlsForInputEvent = OptionalNotNullable()
    # samples with ids or detection time and ids
    samples: Union[SamplesForInputEvent, SamplesIDs] = OptionalNotNullable()
    # image detection time
    detectTime: types.CustomDatetime = OptionalNotNullable()
    # user-defined timestamp relative to something, such as the start of a video
    detectTs: Optional[types.DetectTs] = None
    # image origin
    imageOrigin: types.Str256 = OptionalNotNullable()

    def __init__(self, **kwargs):
        try:
            super().__init__(**kwargs)
        except ValidationError as exc:
            """
            Try get correct error for 'oneOf' validation for `samples` validation
            """
            if "samples" not in kwargs or not isinstance(kwargs["samples"], list) or not kwargs["samples"]:
                raise exc

            if isinstance(kwargs["samples"][0], dict):

                class Model(HandlerInputEstimationsModel):
                    def __init__(self, **kwargs):
                        super(BaseSchema, self).__init__(**kwargs)

                    samples: SamplesForInputEvent

            else:

                class Model(HandlerInputEstimationsModel):
                    def __init__(self, **kwargs):
                        super(BaseSchema, self).__init__(**kwargs)

                    samples: SamplesIDs

            try:
                Model.validate(kwargs)
            except ValidationError as exc:
                raise exc

    @root_validator(skip_on_failure=True)
    def validateHandler(cls, values: dict) -> dict:
        """Validate handler model"""
        validateOneOfStrParams(values, params=["urls", "samples", "image"])
        checkRequiredImageMimetype(values)
        return values