from django_elasticsearch_dsl import Document, fields
from django_elasticsearch_dsl.registries import registry

from .models import StaffCity, StaffGroup


@registry.register_document
class DepartmentDocument(Document):
    class Index:
        name = "departments"

    class Django:
        model = StaffGroup
        fields = {
            "name",
            "level",
        }


@registry.register_document
class CityDocument(Document):
    name = fields.TextField()

    def prepare_name(self, instance: StaffCity):
        return instance.name

    class Index:
        name = "cities"

    class Django:
        model = StaffCity
