"""
This script defines object-relational mapping helpers to ease
 the manipulation of Kili data structures.
"""
from dataclasses import dataclass

class DictClass(dict):
    """
    A python class that acts like dict
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.__dict__ = self

@dataclass
class AnnotationFormat:
    # pylint: disable=invalid-name
    """
    List of annotation formats
    """

    Latest = 'latest'
    Raw = 'raw'
    Simple = 'simple'
    YoloV4Pytorch = 'yolo_v4_pytorch'


@dataclass
class JobMLTask:
    # pylint: disable=invalid-name
    """
    List of ML Tasks
    """
    AssetAnnotation = 'ASSET_ANNOTATION'
    Classification = 'CLASSIFICATION'
    NamedEntitiesRecognition = 'NAMED_ENTITIES_RECOGNITION'
    NamedEntitiesRelation = 'NAMED_ENTITIES_RELATION'
    ObjectRelation = 'OBJECT_RELATION'
    ObjectDetection = 'OBJECT_DETECTION'
    PoseEstimation = 'POSE_ESTIMATION'
    RegionDetection = 'REGION_DETECTION'
    SpeechToText = 'SPEECH_TO_TEXT'
    Transcription = 'TRANSCRIPTION'

@dataclass
class JobTool:
    # pylint: disable=invalid-name
    """
    List of tools
    """
    Marker = 'marker'
    Polygon = 'polygon'
    Polyline = 'polyline'
    Range = 'range'
    Rectangle = 'rectangle'
    Semantic = 'semantic'
    Vector = 'vector'


def get_polygon(annotation):
    """
    Extracts a polygon from an annotation

    Parameters
    ----------
    - annotation : Kili annotation
    """
    try:
        return annotation['boundingPoly'][0]['normalizedVertices']
    except KeyError:
        return None


def get_category(annotation):
    """
    Extracts a category from an annotation

    Parameters
    ----------
    - annotation : Kili annotation
    """
    try:
        return annotation['categories'][0]['name']
    except KeyError:
        return None


def get_named_entity(annotation):
    """
    Extracts a named entity from an annotation

    Parameters
    ----------
    - annotation : Kili annotation
    """
    try:
        return {
            'beginId': annotation['beginId'],
            'beginOffset': annotation['beginOffset'],
            'content': annotation['content'],
            'endId': annotation['endId'],
            'endOffset': annotation['endOffset']
        }
    except KeyError:
        return None


def format_image_annotation(annotation):
    """
    Extracts a category, polygon, named entity from an annotation
    depending of the annotation type

    Parameters
    ----------
    - annotation : Kili annotation
    """
    category = get_category(annotation)
    polygon = get_polygon(annotation)
    named_entity = get_named_entity(annotation)
    if polygon is not None:
        return (category, polygon)
    if named_entity is not None:
        return (category, named_entity)
    if category is not None:
        return category
    return None


class Label(DictClass):
    """
    Label class
    """

    jsonResponse = {}

    def json_response(self, _format=AnnotationFormat.Raw):
        """
        Format a json response

        Parameters
        ----------
        - _format: expected format
        """
        if 'jsonResponse' not in self:
            raise Exception(
                'You did not fetch jsonResponse for' \
                f' label "{self["id"] if "id" in self else self}"')
        if _format == AnnotationFormat.Raw:
            return self.jsonResponse
        if _format == AnnotationFormat.Simple:
            job_names = self.jsonResponse.keys()
            if len(job_names) > 1:
                return {
                    'error': 'Simple format is not adapted' \
                    ' when there is more than one job.' \
                    ' Please choose another annotation format.'}
            for job_name in job_names:
                job_response = self.jsonResponse[job_name]
                category = get_category(job_response)
                if category is not None:
                    return category
                if 'annotations' not in job_response:
                    return None
                return [format_image_annotation(annotation) \
                    for annotation in job_response['annotations']]
            return None
        raise Exception(f'format "{_format}" is not a valid annotation format.')


class Asset(DictClass):
    """
    Asset class
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if 'labels' in self:
            labels = []
            for label in self['labels']:
                labels.append(Label(label))
            self.labels = labels
        if 'latestLabel' in self:
            latest_label = self['latestLabel']
            if latest_label is not None:
                self.latestLabel = Label(latest_label) # pylint: disable=invalid-name
