# -*- coding: utf-8 -*-

from uuid import UUID

from arkindex_export import Classification
from arkindex_export.models import (
    DatasetElement,
    Element,
    Entity,
    EntityType,
    Image,
    Transcription,
    TranscriptionEntity,
)


def list_dataset_elements(dataset_id: UUID, set_name: str):
    return (
        Element.select()
        .join(Image)
        .switch(Element)
        .join(DatasetElement, on=DatasetElement.element)
        .where(
            DatasetElement.dataset == dataset_id, DatasetElement.set_name == set_name
        )
    )


def list_classifications(element_id: UUID):
    return Classification.select().where(Classification.element == element_id)


def list_transcriptions(element_id: UUID):
    return Transcription.select().where(Transcription.element == element_id)


def list_transcription_entities(transcription_id: UUID):
    return (
        TranscriptionEntity.select()
        .where(TranscriptionEntity.transcription == transcription_id)
        .join(Entity, on=TranscriptionEntity.entity)
        .join(EntityType, on=Entity.type)
    )