# Copyright (c) 2021-2022, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#

from omni.ui import scene as sc
from typing import List, Sequence, Union


class SelectionShapeModel(sc.AbstractManipulatorModel):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.__items = {
            "ndc_start": (sc.AbstractManipulatorItem(), 2),
            "ndc_current": (sc.AbstractManipulatorItem(), 2),
            "ndc_rect": (sc.AbstractManipulatorItem(), 4),
            "mode": (sc.AbstractManipulatorItem(), 1),
            "live_update": (sc.AbstractManipulatorItem(), 1)
        }
        self.__ndc_ret_item = self.__items.get('ndc_rect')[0]
        self.__values = {item: [] for item, _ in self.__items.values()}

    def __validate_arguments(self, name: Union[str, sc.AbstractManipulatorItem], values: Sequence[Union[int, float]] = None) -> sc.AbstractManipulatorItem:
        if isinstance(name, sc.AbstractManipulatorItem):
            return name
        item, expected_len = self.__items.get(name, (None, None))
        if item is None:
            raise KeyError(f"SelectionShapeModel doesn't understand values of {name}")
        if values and len(values) != expected_len:
            raise ValueError(f"SelectionShapeModel {name} takes {expected_len} values, got {len(values)}")
        return item

    def get_item(self, name: str) -> sc.AbstractManipulatorItem():
        return self.__items.get(name, (None, None))[0]

    def set_ints(self, name: str, values: Sequence[int]):
        item = self.__validate_arguments(name, values)
        self.__values[item] = values

    def set_floats(self, name: str, values: Sequence[int]):
        item = self.__validate_arguments(name, values)
        self.__values[item] = values

    def get_as_ints(self, name: str) -> List[int]:
        item = self.__validate_arguments(name)
        return self.__values[item]

    def get_as_floats(self, name: str) -> List[float]:
        item = self.__validate_arguments(name)
        if item == self.__ndc_ret_item:
            ndc_start, ndc_end = self.__values[self.get_item('ndc_start')], self.__values[self.get_item('ndc_current')]
            if not ndc_start:
                return []
            if not ndc_end:
                ndc_end = ndc_start
            min_x = min(ndc_start[0], ndc_end[0])
            max_x = max(ndc_start[0], ndc_end[0])
            min_y = min(ndc_start[1], ndc_end[1])
            max_y = max(ndc_start[1], ndc_end[1])
            return [min_x, min_y, max_x, max_y]

        return self.__values[item]
