import kotlinx.serialization.*
import kotlinx.serialization.internal.ArrayListSerializer
import kotlinx.serialization.internal.StringDescriptor
import kotlinx.serialization.json.Json
import kotlinx.serialization.modules.SerializersModule

@Serializable
sealed class Foo {

    @Serializable
    @SerialName("foo_a")
    data class A(val s: String = "hunter") : Foo()

    sealed class B {
        @Serializable
        @SerialName("foo_b0")
        data class B0(val i: Int = 42) : Foo()

        @Serializable
        @SerialName("foo_b1")
        data class B1(val b: Boolean = false) : Foo()
    }
}

@Serializable
data class FooWrapper(@Polymorphic val foo: Foo)

/**
 * Foo Wrapper is used when we want to serialize Foo as a direct property.
 */
@Serializer(forClass = Foo::class)
object FooSerializer : KSerializer<Foo> {
    private val json by lazy { Json(context = fooContext) }

    override val descriptor: SerialDescriptor = StringDescriptor.withName("foo")

    override fun deserialize(decoder: Decoder): Foo =
        decoder.decodeString()
            .let { str -> json.parse(FooWrapper.serializer(), str).foo }

    override fun serialize(encoder: Encoder, obj: Foo) =
        FooWrapper(obj)
            .let { fooWrapped -> json.stringify(FooWrapper.serializer(), fooWrapped) }
            .let(encoder::encodeString)
}

/**
 * However, if we want to serialize Foo as part of a list, we need to use another
 * custom serializer! In this case: [FooListSerializer].
 */
@Serializable
data class MultiFoos(@Serializable(with = FooListSerializer::class) val foos: List<Foo>)

val fooContext = SerializersModule {
    polymorphic<Foo> {
        Foo.A::class with Foo.A.serializer()
        Foo.B.B0::class with Foo.B.B0.serializer()
        Foo.B.B1::class with Foo.B.B1.serializer()
    }
}

object FooListSerializer :
    KSerializer<List<Foo>> by ArrayListSerializer<Foo>(PolymorphicSerializer(Foo::class) as KSerializer<Foo>)

fun main(args: Array<String>) {
    val fooA = Foo.A()
    val fooB0 = Foo.B.B0()
    val fooB1 = Foo.B.B1()
    val foos = MultiFoos(listOf(fooA, fooB0, fooB1))

    Json(context = fooContext).apply {
        stringify(MultiFoos.serializer(), foos)
            .also(::println)
            .let { str -> parse(MultiFoos.serializer(), str) }
            .also(::println)
    }

    println("---")

    Json.apply {
        stringify(FooSerializer, fooA)
            .also(::println)
            .let { str -> parse(FooSerializer, str) }
            .also(::println)
    }
}