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

from uuid import UUID

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


def retrieve_element(element_id: UUID) -> Element:
    return Element.get_by_id(element_id)


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)
    )