@file:Suppress(
  "CANNOT_OVERRIDE_INVISIBLE_MEMBER",
  "INVISIBLE_MEMBER",
  "INVISIBLE_REFERENCE",
)

import androidx.collection.ObjectIntMap
import androidx.compose.runtime.DerivedState
import androidx.compose.runtime.DerivedStateObserver
import androidx.compose.runtime.SnapshotMutationPolicy
import androidx.compose.runtime.derivedStateObservers
import androidx.compose.runtime.derivedStateOf
import androidx.compose.runtime.snapshots.StateObject
import androidx.compose.runtime.structuralEqualityPolicy

/**
 * Helpers for working with [derivedStateOf] state objects. All these APIs are internal to the Compose
 * runtime, so we can't access them normally. We cheat by disabling internal checks on this file, but
 * that's really dangerous since it allows us to access tons of stuff, so these helpers are in their
 * own file to isolate from the higher-level code.
 */
internal object DerivedStateHelper {

  fun StateObject.asMaybeDerivedState(): DerivedStateProxy<*>? =
    if (this is DerivedState<*>) DerivedStateProxy(this)
    else null

  fun StateObject.asDerivedState(): DerivedStateProxy<*> =
    DerivedStateProxy(this as DerivedState<*>)

  inline fun <R> observeDerivedStateRecalculations(
    crossinline onStart: () -> Unit,
    crossinline onDone: () -> Unit,
    block: () -> R
  ) {
    val observers = derivedStateObservers()
    try {
      observers.add(object : DerivedStateObserver {
        override fun start(derivedState: DerivedState<*>) {
          onStart()
        }

        override fun done(derivedState: DerivedState<*>) {
          onDone()
        }
      })
      block()
    } finally {
      observers.removeAt(observers.lastIndex)
    }
  }

  /**
   * No-overhead wrapper of a [DerivedState] that allows other code in this module to access properties
   * of a [DerivedState] without enabling unsafe access.
   */
  @JvmInline
  value class DerivedStateProxy<T>(private val derivedState: DerivedState<T>) {
    val policy: SnapshotMutationPolicy<T>? get() = derivedState.policy
    val currentValue: T get() = derivedState.currentRecord.currentValue
    val dependencies: ObjectIntMap<StateObject> get() = derivedState.currentRecord.dependencies

    /**
     * Compares [value] to this state's current value using its [SnapshotMutationPolicy].
     */
    fun isCurrentValueEquivalentTo(value: Any?): Boolean {
      val policy = policy ?: structuralEqualityPolicy()
      @Suppress("UNCHECKED_CAST")
      return policy.equivalent(value as T, currentValue)
    }
  }
}