#!/usr/bin/env python3
#
# Copyright 2020 Graviti. All Rights Reserved.
#

"""This file defines the class Labelset.
"""

import os
from typing import Any, Dict, List, Optional, Union

from ..dataset import Data, Frame, FusionSegment, Segment
from ..label import LabelTable, LabelType
from .path import SegmentNameWrapper
from .requests import Client, multithread_upload

_FULL_DATA_LIST_TYPE = 2


class LabelSegmentClientBase:  # pylint: disable=too-few-public-methods
    """This class defines the concept of a labelset segment client and some opertions on it.

    :param name: Segment name, unique for a labelset
    :param labelset_id: Labelset ID
    :param client: The client used for sending request to TensorBay
    """

    def __init__(self, name: str, labelset_id: str, client: Client) -> None:
        self._name = name
        self._labelset_id = labelset_id
        self._client = client
        self._name_wrapper = SegmentNameWrapper(name)

    def get_segment_name(self) -> str:
        """Return the name of this labelset segment."""
        return self._name

    def _upload_label(
        self,
        data: Data,
        sensor_name: str = "",
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Upload label to the labelset.

        :param data: The Data contains the labels to be uploaded
        :param sensor_name: Sensor name of the Data
        :param metadata: Some additional data of the label
        """
        remote_path = data.remote_path if data.remote_path else os.path.basename(data.fileuri)

        post_data: Dict[str, Any] = {}
        post_data["ObjectPath"] = self._name_wrapper.get_object_path(remote_path, sensor_name)
        post_data["labelValues"] = data.dump_labels()
        post_data["labelSetId"] = self._labelset_id
        if metadata:
            post_data["labelMeta"] = metadata
        self._client.labelset_post("putLabel", post_data)


class LabelSegmentClient(LabelSegmentClientBase):
    """A class used to represent a segment client of a labelset."""

    def upload_data_object(self, data: Data) -> None:
        """Upload the input data's label to the labelset segment.

        :param data: The `Data` object which contains the labels to be uploaded
        """
        self._upload_label(data)

    def list_data(self) -> List[str]:
        """List all data remote path in a labelset segment.

        :return: A list of data path
        """
        post_data = {
            "labelSetId": self._labelset_id,
            "listType": _FULL_DATA_LIST_TYPE,
            "segmentName": self._name_wrapper.post_segment_name,
        }
        objects = self._client.labelset_post("listObjectsBySegment", post_data)["objects"]
        return [self._name_wrapper.get_remote_path(info["objectName"]) for info in objects]


class FusionLabelSegmentClient(LabelSegmentClientBase):
    """A class used to represent a segment client of a fusion labelset."""

    def upload_frame_object(self, frame: Frame) -> None:
        """Upload the input frame's label to the fusion labelset segment.

        :param frame: The `Frame` object which contains the labels to be uploaded
        """
        for sensor_name, data in frame.items():
            self._upload_label(data, sensor_name)

    def list_frame_objects(self) -> List[Frame]:
        """List all frames in the labelset segment without the label info.

        :return: A list `Frame` object
        """
        post_data = {
            "labelSetId": self._labelset_id,
            "listType": _FULL_DATA_LIST_TYPE,
            "pageSize": -1,
            "needLabel": False,
            "segmentName": self._name_wrapper.post_segment_name,
        }
        response = self._client.labelset_post("listFrameObjects", post_data)["frameObjects"]
        frames = []
        for frame_info in response:
            frame = Frame()
            for data_info in frame_info:
                sensor_name = data_info["sensorName"]
                object_path = data_info["objectPath"]
                remote_path = self._name_wrapper.get_remote_path(object_path, sensor_name)
                data = Data(remote_path, timestamp=data_info.get("timestamp", None))
                frame[sensor_name] = data

            frames.append(frame)

        return frames


class LabelsetClientBase:
    """A base class which respresents a client of labelset.

    :param labelset_id: The ID of the labelset
    :param client: The client used for sending request to TensorBay
    """

    TYPE_GROUND_TRUTH = 3
    _PUBLISH_STATUS = 2
    _STATISTICSED_STATUS = 4
    _TASK_TYPE = {
        LabelType.CLASSIFICATION: {"code": 7, "nameEn": "2D CLASSIFICATION"},
        LabelType.BOX2D: {"code": 10, "nameEn": "2D BOX"},
        LabelType.BOX3D: {"code": 4, "nameEn": "3D BOX"},
        LabelType.POLYGON: {"code": 22, "nameEn": "2D POLYGON"},
        LabelType.POLYLINE: {"code": 8, "nameEn": "2D POLYLINE"},
    }

    def __init__(self, labelset_id: str, client: Client) -> None:
        self._labelset_id = labelset_id
        self._client = client

    def get_labelset_id(self) -> str:
        """Return the id of the labelset.

        :return: The id of the labelset
        """
        return self._labelset_id

    def upload_label_table(self, label_table: Union[LabelTable, List[LabelTable]]) -> None:
        """Upload label table to the labelset.

        :param label_table: The label table to be uploaded
        """
        metadata = LabelsetClientBase.create_metadata(label_table)
        if not metadata:
            return

        post_data = {
            "labelSetId": self._labelset_id,
            "meta": metadata,
        }
        self._client.labelset_post("updateLableSet", post_data)

    def list_segments(self) -> List[str]:
        """List all segment names in a labelset.

        :return: A list of segment names
        """
        post_data = {
            "labelSetId": self._labelset_id,
            "listType": _FULL_DATA_LIST_TYPE,
            "pageSize": -1,
        }
        segment_names = self._client.labelset_post("listSegments", post_data)["segments"]
        return [SegmentNameWrapper.get_local_segment_name(name) for name in segment_names]

    def publish(self) -> None:
        """Publish the labelset."""
        post_data = {"labelSetId": self._labelset_id, "status": LabelsetClientBase._PUBLISH_STATUS}
        self._client.labelset_post("updateLabelSetStatus", post_data)

    def is_published(self) -> bool:
        """Gheck whether the labelset is published.

        :return: Return `True` if the lableset is publish, otherwise return `False`
        """
        post_data = {"id": self._labelset_id, "projection": {"labelSet.status": 1}}
        data = self._client.labelset_post("listLabelSets", post_data)
        status = data["labelSets"][0]["labelSet"]["status"]
        return status == LabelsetClientBase._STATISTICSED_STATUS  # type: ignore[no-any-return]

    @staticmethod
    def create_metadata(label_tables: Union[LabelTable, List[LabelTable]]) -> Dict[str, Any]:
        """Return metadata created from the label table or the list of label tables.

        :return: A dict in metadata format
        """
        if isinstance(label_tables, LabelTable):
            label_tables = [label_tables]

        task_types = []
        for label_table in label_tables:
            if not label_table.categories:
                continue
            task_type = LabelsetClientBase._TASK_TYPE[label_table.label_type].copy()
            task_type["is_tracking"] = label_table.is_tracking
            task_type["categories"] = [{"nameEn": category} for category in label_table.categories]
            task_types.append(task_type)

        return {"taskTypes": task_types} if task_types else {}


class LabelsetClient(LabelsetClientBase):
    """A class which respresents a client of a labelset."""

    def get_segment(self, name: str) -> LabelSegmentClient:
        """Get a labelset segment according to the input name.

        :param name: The name of the requested labelset segment
        :return: The requested labelset segment client
        """

        return LabelSegmentClient(name, self._labelset_id, self._client)

    def upload_segment_object(self, segment: Segment, *, jobs: int = 1) -> LabelSegmentClient:
        """Upload a `Segment` to the labelset,
        This function will upload all labels contains in the input `Segment` object,

        :param segment: The `Segment` object contains the labels needs to be uploaded
        :param jobs: The number of the max workers in multithread upload
        :return: The client used for uploading the label in the `Segment`
        """
        segment_client = self.get_segment(segment.name)
        multithread_upload(segment_client.upload_data_object, segment, jobs=jobs)
        return segment_client


class FusionLabelsetClient(LabelsetClientBase):
    """A class which respresents a client of a fusion labelset."""

    def get_segment(self, name: str) -> FusionLabelSegmentClient:
        """Get a client of fusion labelset segment according to the input name.

        :param name: The name of the requested fusion labelset segment
        :return: The requested client of fusion labelset segment
        """
        return FusionLabelSegmentClient(name, self._labelset_id, self._client)

    def upload_segment_object(
        self,
        segment: FusionSegment,
        *,
        jobs: int = 1,
    ) -> FusionLabelSegmentClient:
        """Upload a `FusionSegment` to the labelset,
        This function will upload all labels contains in the input `FusionSegment` object,

        :param segment: The `FusionSegment` object contains the labels needs to be uploaded
        :param jobs: The number of the max workers in multithread upload
        :return: The client used for uploading the label in the `FusionSegment`
        """
        segment_client = self.get_segment(segment.name)
        multithread_upload(segment_client.upload_frame_object, segment, jobs=jobs)
        return segment_client
