import kotlinx.serialization.*
import kotlinx.serialization.descriptors.*
import kotlinx.serialization.encoding.*
import kotlinx.serialization.json.*
import kotlinx.serialization.protobuf.*

public object JsonObjectProtobufSerializer : KSerializer<JsonObject> {
    override val descriptor: SerialDescriptor =
        SerialDescriptor("kotlinx.serialization.json.JsonObject", ProtobufStruct.serializer().descriptor)

    override fun serialize(encoder: Encoder, value: JsonObject): Unit = if (encoder is JsonEncoder) {
        encoder.encodeJsonElement(value)
    } else {
        encoder.encodeSerializableValue(ProtobufStruct.serializer(), value.toStruct())
    }

    override fun deserialize(decoder: Decoder): JsonObject = if (decoder is JsonDecoder) {
        decoder.decodeJsonElement().jsonObject
    } else {
        decoder.decodeSerializableValue(ProtobufStruct.serializer()).toJsonObject()
    }
}

public object JsonElementProtobufSerializer : KSerializer<JsonElement> {
    override val descriptor: SerialDescriptor =
        SerialDescriptor("kotlinx.serialization.json.JsonElement", ProtobufValue.serializer().descriptor)

    override fun serialize(encoder: Encoder, value: JsonElement): Unit = if (encoder is JsonEncoder) {
        encoder.encodeJsonElement(value)
    } else {
        encoder.encodeSerializableValue(ProtobufValue.serializer(), value.toValue())
    }

    override fun deserialize(decoder: Decoder): JsonElement = if (decoder is JsonDecoder) {
        decoder.decodeJsonElement()
    } else {
        decoder.decodeSerializableValue(ProtobufValue.serializer()).toJsonElement()
    }
}

public object JsonArrayProtobufSerializer : KSerializer<JsonArray> {
    override val descriptor: SerialDescriptor =
        SerialDescriptor("kotlinx.serialization.json.JsonArray", ProtobufListValue.serializer().descriptor)

    override fun serialize(encoder: Encoder, value: JsonArray): Unit = if (encoder is JsonEncoder) {
        encoder.encodeJsonElement(value)
    } else {
        encoder.encodeSerializableValue(ProtobufListValue.serializer(), value.toListValue())
    }

    override fun deserialize(decoder: Decoder): JsonArray = if (decoder is JsonDecoder) {
        decoder.decodeJsonElement().jsonArray
    } else {
        decoder.decodeSerializableValue(ProtobufListValue.serializer()).toJsonArray()
    }
}

@Serializable
private data class ProtobufStruct(@ProtoNumber(1) val fields: Map<String, ProtobufValue>)
@Serializable
private data class ProtobufValue(@ProtoOneOf val kind: ProtobufValueKind)
@Serializable
private sealed interface ProtobufValueKind
@Serializable
private data class ProtobufNullKind(@ProtoNumber(1) val value: ProtobufNullValue) : ProtobufValueKind
@Serializable
private data class ProtobufNumberKind(@ProtoNumber(2) val value: Double) : ProtobufValueKind
@Serializable
private data class ProtobufStringKind(@ProtoNumber(3) val value: String) : ProtobufValueKind
@Serializable
private data class ProtobufBoolKind(@ProtoNumber(4) val value: Boolean) : ProtobufValueKind
@Serializable
private data class ProtobufStructKind(@ProtoNumber(5) val value: ProtobufStruct) : ProtobufValueKind
@Serializable
private data class ProtobufListKind(@ProtoNumber(6) val value: ProtobufListValue) : ProtobufValueKind
@Serializable
private enum class ProtobufNullValue { INSTANCE }
@Serializable
private data class ProtobufListValue(@ProtoNumber(1) val values: List<ProtobufValue>)

private fun JsonObject.toStruct(): ProtobufStruct = ProtobufStruct(mapValues { it.value.toValue() })
private fun JsonElement.toValue(): ProtobufValue = ProtobufValue(
    when (this) {
        is JsonObject -> ProtobufStructKind(toStruct())
        is JsonArray -> ProtobufListKind(toListValue())
        is JsonPrimitive -> if (isString) {
            ProtobufStringKind(content)
        } else {
            booleanOrNull?.let { ProtobufBoolKind(it) }
                ?: doubleOrNull?.let { ProtobufNumberKind(it) }
                ?: ProtobufNullKind(ProtobufNullValue.INSTANCE)
        }
    }
)
private fun JsonArray.toListValue(): ProtobufListValue = ProtobufListValue(map { it.toValue() })

private fun ProtobufStruct.toJsonObject(): JsonObject = JsonObject(fields.mapValues { it.value.toJsonElement() })
private fun ProtobufValue.toJsonElement(): JsonElement = when (kind) {
    is ProtobufNullKind -> JsonNull
    is ProtobufNumberKind -> JsonPrimitive(kind.value)
    is ProtobufStringKind -> JsonPrimitive(kind.value)
    is ProtobufBoolKind -> JsonPrimitive(kind.value)
    is ProtobufStructKind -> kind.value.toJsonObject()
    is ProtobufListKind -> kind.value.toJsonArray()
}
private fun ProtobufListValue.toJsonArray(): JsonArray = JsonArray(values.map { it.toJsonElement() })