ViewModelSerializer.kt

package com.dmm.eikaiwa.viewmodels

import kotlinx.serialization.KSerializer
import kotlinx.serialization.MissingFieldException
import kotlinx.serialization.descriptors.SerialDescriptor
import kotlinx.serialization.encoding.Decoder
import kotlinx.serialization.encoding.Encoder
import kotlinx.serialization.json.JsonDecoder
import kotlinx.serialization.json.JsonElement
import kotlinx.serialization.json.JsonEncoder
import kotlinx.serialization.json.boolean
import kotlinx.serialization.json.buildJsonObject
import kotlinx.serialization.json.jsonObject
import kotlinx.serialization.json.jsonPrimitive
import kotlinx.serialization.json.put
import kotlinx.serialization.modules.SerializersModule

/**
 * Abstract base class for serializing and deserializing ViewModels to and from JSON.
 * It handles common ViewModel serialization concerns like:
 * - Adding metadata fields (`_type`, `_version`).
 * - Optionally nesting the ViewModel's content under a `rootKey`.
 * - Managing references to avoid redundant serialization of shared ViewModels.
 *
 * Subclasses need to provide a `KSerializer` for the specific ViewModel type `T`.
 *
 * @param T The type of the ViewModel to be serialized/deserialized.
 *          It must implement both `ViewModelRecord` (for identity and root status)
 *          and `ViewModelMetadata` (for type and version information).
 * @property serializer The `KSerializer` for the specific ViewModel type `T`.
 *                      This serializer will be used to handle the actual content
 *                      serialization/deserialization of the ViewModel.
 * @property rootKey An optional string. If provided, the serialized ViewModel's content
 *                   will be nested under this key in the JSON object. If `null`,
 *                   the ViewModel's properties will be merged directly into the top-level JSON object.
 */
abstract class JsonViewModelSerializer<T>(
    private val serializer: KSerializer<T>,
    private val rootKey: String? = null
) : KSerializer<T> where T : ViewModelRecord, T : ViewModelMetadata {

    override val descriptor: SerialDescriptor get() = serializer.descriptor

    companion object {
        private var referenceIdIndex = 1
    }

    final override fun serialize(encoder: Encoder, value: T) {
        require(encoder is JsonEncoder)

        val contextSerializer = context(encoder.json.serializersModule)
        val referenceCache = contextSerializer.referenceCache
        val rootIds = contextSerializer.rootIds
        val isRootModel = rootIds?.contains(value.id) ?: false

        val jsonObject = buildJsonObject {
            put("_type", value.typeName)
            put("_version", value.version)

            val contentJsonElement = encoder.json.encodeToJsonElement(serializer, value)

            if (rootKey != null) {
                // If there's a rootKey, nest the contentJsonElement under it.
                put(rootKey, contentJsonElement)
            } else {
                // If no rootKey, merge the content into the current builder.
                contentJsonElement.jsonObject.entries.forEach { entry ->
                    put(entry.key, entry.value)
                }
            }
        }
        // Create reference if needed.
        if (value.root && isRootModel.not()) {
            val id = "ref:${referenceIdIndex++}"
            val referenceObject = buildJsonObject {
                put("_ref", id)
            }
            referenceCache.encoded[id] = jsonObject
            encoder.encodeJsonElement(referenceObject)
        } else {
            encoder.encodeJsonElement(jsonObject)
        }
    }

    @Suppress("UNCHECKED_CAST")
    final override fun deserialize(decoder: Decoder): T {
        require(decoder is JsonDecoder)

        val contextSerializer = context(decoder.json.serializersModule)
        val referenceCache = contextSerializer.referenceCache

        val element = decoder.decodeJsonElement()
        val jsonObject = element.jsonObject
        if ("_ref" in jsonObject) {
            val key = jsonObject["_ref"]!!.jsonPrimitive.content
            // Check if there is an already parsed model for that key.
            referenceCache.parsed[key]?.let {
                return it as? T ?: throw ModelTypesDoNotMatchException(
                    it.toString(),
                    serializer.descriptor.serialName
                )
            }
            // If not, grab raw object and parse it.
            val modelObject = referenceCache.raw[key]
                ?: throw ReferenceNotFoundException(key, referenceCache)
            val parsedModel = decodeViewModel(decoder, modelObject)
            referenceCache.parsed[key] = parsedModel
            return parsedModel.notNew()
        } else {
            return decodeViewModel(decoder, element).notNew()
        }
    }

    private fun decodeViewModel(decoder: Decoder, jsonElement: JsonElement): T {
        require(decoder is JsonDecoder)

        if (rootKey != null) {
            val rootKeyObject = jsonElement.jsonObject[rootKey]
                ?: throw MissingFieldException(rootKey, serializer.descriptor.serialName)
            return decoder.json.decodeFromJsonElement(serializer, rootKeyObject)
        }

        val model = decoder.json.decodeFromJsonElement(serializer, jsonElement)

        // Check for migration warning.
        try {
            jsonElement.jsonObject["_migrated"]?.jsonPrimitive?.boolean?.let {
                if (it) {
                    ViewModelMigrationWarning.modelsThatNeedMigration[model.typeName] = model.version
                }
            }
        } catch (_: Throwable) { }

        return model
    }

    private fun context(serializersModule: SerializersModule): ViewModelContextSerializer {
        val contextSerializer = serializersModule.getContextual(
            ViewModelContext::class
        ) as? ViewModelContextSerializer ?: throw ReferenceCacheNotFoundException()
        return contextSerializer
    }
}