enum class PartitionColumnType(val value: String) {
    EVENT_TIME("event_time"),
    PROCESSING_TIME("processing_time")
}


class PartitionTypeBucketAssigner(private val partitionType: PartitionColumnType) : BucketAssigner<MyProtobufModel, String> {

    private val formatter = DateTimeFormatter.ofPattern("'year='yyyy/'month='M/'day='d")
    val ZONE_UTC: ZoneId = ZoneId.of("UTC")

    override fun getBucketId(element: MyProtobufModel, context: BucketAssigner.Context): String {
        return when (partitionType) {
            PartitionColumnType.EVENT_TIME -> {
                val eventTimestamp = Instant.ofEpochMilli(element.timestamp.getSeconds() * 1000)
                val eventTime = LocalDateTime.ofInstant(eventTimestamp, ZONE_UTC)

                formatter.format(eventTime)
            }
            PartitionColumnType.PROCESSING_TIME -> {
                val contextProcessingTime = context.currentProcessingTime()
                val instant = Instant.ofEpochMilli(contextProcessingTime)
                val processingTime = LocalDateTime.ofInstant(instant, ZONE_UTC)

                formatter.format(processingTime)
            }
        }
    }

    override fun getSerializer(): SimpleVersionedStringSerializer {
        return SimpleVersionedStringSerializer.INSTANCE
    }
}