from bson.codec_options import CodecOptions, TypeCodec, TypeRegistry

from maps_adv.geosmb.harmonist.server.lib.enums import (
    ColumnType,
    InputDataType,
    PipelineStep,
    StepStatus,
)

__all__ = ["codec_options"]


class BaseEnumCodec(TypeCodec):
    python_type = None
    bson_type = "string"

    def transform_python(self, value):
        return value.name

    def transform_bson(self, value):
        return self.python_type[value]


class ColumnTypeCodec(BaseEnumCodec):
    python_type = ColumnType


class PipelineStepCodec(BaseEnumCodec):
    python_type = PipelineStep


class StepStatusCodec(BaseEnumCodec):
    python_type = StepStatus


class InputDataTypeCodec(BaseEnumCodec):
    python_type = InputDataType


column_type_codec = ColumnTypeCodec()
input_data_type_codec = InputDataTypeCodec()
pipeline_step_codec = PipelineStepCodec()
pipeline_step_status_codec = StepStatusCodec()


codec_options = CodecOptions(
    type_registry=TypeRegistry(
        [
            column_type_codec,
            input_data_type_codec,
            pipeline_step_codec,
            pipeline_step_status_codec,
        ]
    )
)
