diff --git a/treehouse/treehouse-compose/src/commonMain/kotlin/app/cash/treehouse/compose/applier.kt b/treehouse/treehouse-compose/src/commonMain/kotlin/app/cash/treehouse/compose/applier.kt
index 77520f3d4..4383be9a1 100644
--- a/treehouse/treehouse-compose/src/commonMain/kotlin/app/cash/treehouse/compose/applier.kt
+++ b/treehouse/treehouse-compose/src/commonMain/kotlin/app/cash/treehouse/compose/applier.kt
@@ -180,13 +180,19 @@ internal class ProtocolApplier(
   override fun remove(index: Int, count: Int) {
     // Children instances are never removed from their parents.
     val current = current as ProtocolChildrenNode
-
     val children = current.children
+
+    // TODO We should not have to track this and send it as part of the protocol.
+    //  Ideally this would be entirely encapsulated on the display-side with additional bookkeeping.
+    //  For now, we track it here and send it in the protocol as a simple solution.
+    val removedIds = ArrayList<Long>(count)
     for (i in index until index + count) {
-      nodes.remove(children[i].id)
+      removedIds.add(children[i].id)
     }
+
+    nodes.keys.removeAll(removedIds)
     children.remove(index, count)
-    childrenDiffs.add(ChildrenDiff.Remove(current.id, current.tag, index, count))
+    childrenDiffs.add(ChildrenDiff.Remove(current.id, current.tag, index, count, removedIds))
   }

   override fun move(from: Int, to: Int, count: Int) {
diff --git a/treehouse/treehouse-protocol/src/commonMain/kotlin/app/cash/treehouse/protocol/protocol.kt b/treehouse/treehouse-protocol/src/commonMain/kotlin/app/cash/treehouse/protocol/protocol.kt
index 42f316c5d..e16bb4a0f 100644
--- a/treehouse/treehouse-protocol/src/commonMain/kotlin/app/cash/treehouse/protocol/protocol.kt
+++ b/treehouse/treehouse-protocol/src/commonMain/kotlin/app/cash/treehouse/protocol/protocol.kt
@@ -84,7 +84,14 @@ public sealed class ChildrenDiff {
     override val tag: Int,
     val index: Int,
     val count: Int,
-  ) : ChildrenDiff()
+    val removedIds: List<Long>,
+  ) : ChildrenDiff() {
+    init {
+      require(count == removedIds.size) {
+        "Count $count != Removed ID list size ${removedIds.size}"
+      }
+    }
+  }

   public companion object {
     public const val RootId: Long = 0L
diff --git a/treehouse/treehouse-protocol/src/commonTest/kotlin/app/cash/treehouse/protocol/ProtocolTest.kt b/treehouse/treehouse-protocol/src/commonTest/kotlin/app/cash/treehouse/protocol/ProtocolTest.kt
index b6adaebf1..31a20c77a 100644
--- a/treehouse/treehouse-protocol/src/commonTest/kotlin/app/cash/treehouse/protocol/ProtocolTest.kt
+++ b/treehouse/treehouse-protocol/src/commonTest/kotlin/app/cash/treehouse/protocol/ProtocolTest.kt
@@ -22,6 +22,7 @@ import kotlinx.serialization.modules.SerializersModule
 import kotlinx.serialization.modules.polymorphic
 import kotlin.test.Test
 import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith

 class ProtocolTest {
   private val format = Json {
@@ -51,7 +52,7 @@ class ProtocolTest {
         ChildrenDiff.Clear,
         ChildrenDiff.Insert(1, 2, 3, 4, 5),
         ChildrenDiff.Move(1, 2, 3, 4, 5),
-        ChildrenDiff.Remove(1, 2, 3, 4),
+        ChildrenDiff.Remove(1, 2, 3, 4, listOf(5, 6, 7, 8)),
       ),
       propertyDiffs = listOf(
         PropertyDiff(1, 2, "Hello"),
@@ -63,7 +64,7 @@ class ProtocolTest {
       """["clear",{}],""" +
       """["insert",{"id":1,"tag":2,"childId":3,"kind":4,"index":5}],""" +
       """["move",{"id":1,"tag":2,"fromIndex":3,"toIndex":4,"count":5}],""" +
-      """["remove",{"id":1,"tag":2,"index":3,"count":4}]""" +
+      """["remove",{"id":1,"tag":2,"index":3,"count":4,"removedIds":[5,6,7,8]}]""" +
       """],"propertyDiffs":[""" +
       """{"id":1,"tag":2,"value":["kotlin.String","Hello"]},""" +
       """{"id":1,"tag":2}""" +
@@ -80,6 +81,12 @@ class ProtocolTest {
     assertJsonRoundtrip(Diff.serializer(), model, json)
   }

+  @Test fun removeCountMustMatchListSize() {
+    assertFailsWith<IllegalArgumentException>("Count 4 != Removed ID list size 3") {
+      ChildrenDiff.Remove(1, 2, 3, 4, listOf(5, 6, 7))
+    }
+  }
+
   private fun <T> assertJsonRoundtrip(serializer: KSerializer<T>, model: T, json: String) {
     assertEquals(json, format.encodeToString(serializer, model))
     assertEquals(model, format.decodeFromString(serializer, json))
diff --git a/treehouse/treehouse-widget/src/commonMain/kotlin/app/cash/treehouse/widget/WidgetDisplay.kt b/treehouse/treehouse-widget/src/commonMain/kotlin/app/cash/treehouse/widget/WidgetDisplay.kt
index 1c04a3088..c3fdca8ae 100644
--- a/treehouse/treehouse-widget/src/commonMain/kotlin/app/cash/treehouse/widget/WidgetDisplay.kt
+++ b/treehouse/treehouse-widget/src/commonMain/kotlin/app/cash/treehouse/widget/WidgetDisplay.kt
@@ -55,7 +55,7 @@ public class WidgetDisplay<T : Any>(
         }
         is ChildrenDiff.Remove -> {
           children.remove(childrenDiff.index, childrenDiff.count)
-          // TODO we need to remove widgets from our map!
+          widgets.keys.removeAll(childrenDiff.removedIds)
         }
         ChildrenDiff.Clear -> {
           children.clear()

we need to remove widgets from our map!
we need to remove widgets from our map!

https://github.com/square/treehouse/blob/1ad7b7242bb7c68e31465a3e65ecfd7bded817ae/treehouse-widget/src/commonMain/kotlin/app/cash/treehouse/widget/WidgetDisplay.kt#L30

```text

package app.cash.treehouse.widget

import app.cash.treehouse.protocol.Diff
import app.cash.treehouse.protocol.WidgetDiff

class WidgetDisplay<T : Any>(
  private val root: Widget<T>,
  private val factory: Widget.Factory<T>,
  private val events: EventSink,
) {
  private val widgets = mutableMapOf(Diff.RootId to root)

  fun apply(diff: Diff) {
    for (widgetDiff in diff.widgetDiffs) {
      val widget = checkNotNull(widgets[widgetDiff.id]) {
        "Unknown widget ID ${widgetDiff.id}"
      }

      when (widgetDiff) {
        is WidgetDiff.Insert -> {
          val childWidget = factory.create(widget.value, widgetDiff.kind, widgetDiff.childId, events)
          widgets[widgetDiff.childId] = childWidget
          widget.children(widgetDiff.childrenIndex).insert(widgetDiff.index, childWidget.value)
        }
        is WidgetDiff.Move -> {
          widget.children(widgetDiff.childrenIndex).move(widgetDiff.fromIndex, widgetDiff.toIndex, widgetDiff.count)
        }
        is WidgetDiff.Remove -> {
          widget.children(widgetDiff.childrenIndex).remove(widgetDiff.index, widgetDiff.count)
          // TODO we need to remove widgets from our map!
        }
        WidgetDiff.Clear -> {
          widget.children(Diff.RootChildrenIndex).clear()
          widgets.clear()
          widgets[Diff.RootId] = root
        }
      }
    }

    for (propertyDiff in diff.propertyDiffs) {
      val widget = checkNotNull(widgets[propertyDiff.id]) {
        "Unknown widget iD ${propertyDiff.id}"
      }

      widget.apply(propertyDiff)
    }
  }
}

```

edf3f9a0d1576c577b8785b13ce44e6b2bcae7c4

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.widget

import kotlinx.cinterop.convert
import platform.UIKit.UIView
import platform.UIKit.insertSubview
import platform.UIKit.removeFromSuperview
import platform.UIKit.subviews
import platform.darwin.NSInteger

public class UIViewChildren(
  private val root: UIView,
) : Widget.Children<UIView> {
  override fun insert(index: Int, widget: UIView) {
    root.insertSubview(widget, index.convert<NSInteger>())
  }

  override fun move(fromIndex: Int, toIndex: Int, count: Int) {
    val views = Array(count) {
      val subview = root.subviews[fromIndex] as UIView
      subview.removeFromSuperview()
      subview
    }

    val newIndex = if (toIndex > fromIndex) {
      toIndex - count
    } else {
      toIndex
    }
    views.forEachIndexed { offset, view ->
      root.insertSubview(view, (newIndex + offset).convert<NSInteger>())
    }
  }

  override fun remove(index: Int, count: Int) {
    repeat(count) {
      (root.subviews[index] as UIView).removeFromSuperview()
    }
  }

  override fun clear() {
    for (subview in root.subviews) {
      (subview as UIView).removeFromSuperview()
    }
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package example.treehouse

import app.cash.treehouse.schema.Children
import app.cash.treehouse.schema.Property
import app.cash.treehouse.schema.Schema
import app.cash.treehouse.schema.Widget

@Schema(
  [
    Box::class,
    Text::class,
    Button::class,
    TextInput::class,
  ]
)
public interface ExampleSchema

@Widget(1)
public data class Box(
  @Children(1) val children: List<Any>,
)

@Widget(2)
public data class Text(
  @Property(1) val text: String?,
)

@Widget(3)
public data class Button(
  @Property(1) val text: String?,
  @Property(2) val onClick: () -> Unit,
)

@Widget(4)
public data class TextInput(
  @Property(1) val text: String?,
  @Property(2) val onChange: (String) -> Unit,
)

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package example.android

import android.app.Activity
import android.os.Bundle
import android.util.Log
import android.view.ViewGroup.LayoutParams.MATCH_PARENT
import android.widget.LinearLayout
import android.widget.LinearLayout.LayoutParams
import android.widget.LinearLayout.VERTICAL
import app.cash.treehouse.compose.AndroidUiDispatcher
import app.cash.treehouse.compose.TreehouseComposition
import app.cash.treehouse.widget.WidgetDisplay
import example.android.sunspot.AndroidSunspotBox
import example.android.sunspot.AndroidSunspotWidgetFactory
import example.shared.Counter
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.cancel

class MainActivity : Activity() {
  private val scope = CoroutineScope(AndroidUiDispatcher.Main)

  override fun onCreate(savedInstanceState: Bundle?) {
    super.onCreate(savedInstanceState)

    val root = LinearLayout(this).apply {
      orientation = VERTICAL
      layoutParams = LayoutParams(MATCH_PARENT, MATCH_PARENT)
    }
    setContentView(root)

    val display = WidgetDisplay(
      root = AndroidSunspotBox(root),
      factory = AndroidSunspotWidgetFactory(this),
    )

    val composition = TreehouseComposition(
      scope = scope,
      display = display::apply,
      onDiff = { Log.d("TreehouseDiff", it.toString()) },
      onEvent = { Log.d("TreehouseEvent", it.toString()) },
    )

    composition.setContent {
      Counter()
    }
  }

  override fun onDestroy() {
    scope.cancel()
    super.onDestroy()
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package example.browser

import app.cash.treehouse.compose.TreehouseComposition
import app.cash.treehouse.compose.WindowAnimationFrameClock
import app.cash.treehouse.widget.WidgetDisplay
import example.browser.sunspot.HtmlSunspotBox
import example.browser.sunspot.HtmlSunspotNodeFactory
import example.shared.Counter
import kotlinx.browser.document
import kotlinx.coroutines.GlobalScope
import kotlinx.coroutines.plus
import org.w3c.dom.HTMLElement

fun main() {
  val content = document.getElementById("content")!! as HTMLElement
  val display = WidgetDisplay(
    root = HtmlSunspotBox(content),
    factory = HtmlSunspotNodeFactory(document),
  )

  val composition = TreehouseComposition(
    scope = GlobalScope + WindowAnimationFrameClock,
    display = display::apply,
    onDiff = { console.log("TreehouseDiff", it.toString()) },
    onEvent = { console.log("TreehouseEvent", it.toString()) },
  )

  composition.setContent {
    Counter()
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package example.shared

import androidx.compose.runtime.Composable
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.setValue
import example.sunspot.compose.SunspotButton
import example.sunspot.compose.SunspotText

@Composable
fun Counter(value: Int = 0) {
  var count by remember { mutableStateOf(value) }

  SunspotButton("-1", onClick = { count-- })
  SunspotText(count.toString())
  SunspotButton("+1", onClick = { count++ })
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package example.sunspot

import app.cash.treehouse.schema.Children
import app.cash.treehouse.schema.Default
import app.cash.treehouse.schema.Property
import app.cash.treehouse.schema.Schema
import app.cash.treehouse.schema.Widget

@Schema(
  [
    SunspotBox::class,
    SunspotText::class,
    SunspotButton::class,
  ]
)
interface Sunspot

@Widget(1)
data class SunspotBox(
  @Children(1) val children: List<Any>,
)

@Widget(2)
data class SunspotText(
  @Property(1) val text: String?,
  @Property(2) @Default("\"black\"") val color: String = "black",
)

@Widget(3)
data class SunspotButton(
  @Property(1) val text: String?,
  @Property(2) @Default("true") val enabled: Boolean = true,
  @Property(3) val onClick: (() -> Unit)? = null,
)

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package example.android.sunspot

import android.view.View
import android.widget.LinearLayout
import app.cash.treehouse.widget.ViewGroupChildren
import example.sunspot.widget.SunspotBox

class AndroidSunspotBox(
  override val value: LinearLayout,
) : SunspotBox<View> {
  override val children = ViewGroupChildren(value)
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package example.android.sunspot

import android.view.View
import android.widget.Button
import example.sunspot.widget.SunspotButton

class AndroidSunspotButton(
  override val value: Button,
) : SunspotButton<View> {
  override fun text(text: String?) {
    value.text = text
  }

  override fun enabled(enabled: Boolean) {
    value.isEnabled = enabled
  }

  override fun onClick(onClick: (() -> Unit)?) {
    value.setOnClickListener(
      if (onClick != null) {
        { onClick() }
      } else {
        null
      }
    )
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package example.android.sunspot

import android.graphics.Color
import android.view.View
import android.widget.TextView
import example.sunspot.widget.SunspotText

class AndroidSunspotText(
  override val value: TextView,
) : SunspotText<View> {
  override fun text(text: String?) {
    value.text = text
  }

  override fun color(color: String) {
    value.setTextColor(Color.parseColor(color))
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package example.android.sunspot

import android.content.Context
import android.view.View
import android.widget.Button
import android.widget.LinearLayout
import android.widget.TextView
import example.sunspot.widget.SunspotBox
import example.sunspot.widget.SunspotButton
import example.sunspot.widget.SunspotText
import example.sunspot.widget.SunspotWidgetFactory

class AndroidSunspotWidgetFactory(
  private val context: Context,
) : SunspotWidgetFactory<View> {
  override fun SunspotBox(): SunspotBox<View> {
    val view = LinearLayout(context)
    return AndroidSunspotBox(view)
  }

  override fun SunspotText(): SunspotText<View> {
    val view = TextView(context)
    return AndroidSunspotText(view)
  }

  override fun SunspotButton(): SunspotButton<View> {
    val view = Button(context)
    return AndroidSunspotButton(view)
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package example.browser.sunspot

import example.sunspot.widget.SunspotBox
import example.sunspot.widget.SunspotButton
import example.sunspot.widget.SunspotText
import example.sunspot.widget.SunspotWidgetFactory
import org.w3c.dom.Document
import org.w3c.dom.HTMLButtonElement
import org.w3c.dom.HTMLDivElement
import org.w3c.dom.HTMLElement
import org.w3c.dom.HTMLSpanElement

class HtmlSunspotNodeFactory(
  private val document: Document,
) : SunspotWidgetFactory<HTMLElement> {
  override fun SunspotBox(): SunspotBox<HTMLElement> {
    val div = document.createElement("div") as HTMLDivElement
    return HtmlSunspotBox(div)
  }

  override fun SunspotText(): SunspotText<HTMLElement> {
    val span = document.createElement("span") as HTMLSpanElement
    return HtmlSunspotText(span)
  }

  override fun SunspotButton(): SunspotButton<HTMLElement> {
    val button = document.createElement("button") as HTMLButtonElement
    return HtmlSunspotButton(button)
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package example.browser.sunspot

import app.cash.treehouse.widget.HTMLElementChildren
import example.sunspot.widget.SunspotBox
import example.sunspot.widget.SunspotButton
import example.sunspot.widget.SunspotText
import org.w3c.dom.HTMLButtonElement
import org.w3c.dom.HTMLElement
import org.w3c.dom.HTMLSpanElement

class HtmlSunspotBox(
  override val value: HTMLElement,
) : SunspotBox<HTMLElement> {
  override val children = HTMLElementChildren(value)
}

class HtmlSunspotText(
  override val value: HTMLSpanElement,
) : SunspotText<HTMLElement> {
  override fun text(text: String?) {
    value.textContent = text
  }

  override fun color(color: String) {
    value.style.color = color
  }
}

class HtmlSunspotButton(
  override val value: HTMLButtonElement,
) : SunspotButton<HTMLElement> {
  override fun text(text: String?) {
    value.textContent = text
  }

  override fun enabled(enabled: Boolean) {
    value.disabled = !enabled
  }

  override fun onClick(onClick: (() -> Unit)?) {
    value.onclick = if (onClick != null) {
      { onClick() }
    } else {
      null
    }
  }
}

/*
 * Copyright 2020 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package androidx.compose.runtime

import androidx.compose.runtime.snapshots.SnapshotMutableState
import kotlinx.browser.window
import kotlin.coroutines.resume
import kotlin.coroutines.suspendCoroutine
import kotlin.time.DurationUnit
import kotlin.time.ExperimentalTime
import kotlin.time.toDuration

internal actual open class ThreadLocal<T> actual constructor(
  initialValue: () -> T
) {
  private var value: T = initialValue()

  actual fun get(): T = value

  actual fun set(value: T) {
    this.value = value
  }
}

internal actual class SnapshotThreadLocal<T> actual constructor() {
  private var value: T? = null

  actual fun get() = value

  actual fun set(value: T?) {
    this.value = value
  }
}

actual class AtomicReference<V> actual constructor(private var value: V) {
  actual fun get(): V = value

  actual fun set(value: V) {
    this.value = value
  }

  actual fun getAndSet(value: V): V {
    val oldValue = this.value
    this.value = value
    return oldValue
  }

  actual fun compareAndSet(expect: V, newValue: V): Boolean =
    if (expect == value) {
      value = newValue
      true
    } else {
      false
    }
}

internal actual fun identityHashCode(instance: Any?): Int {
  if (instance == null) {
    return 0
  }

  val hashCode = instance.asDynamic()[IDENTITY_HASHCODE_FIELD]
  if (hashCode != null) {
    return hashCode
  }

  return when (jsTypeOf(instance)) {
    "object", "function" -> memoizeIdentityHashCode(instance)
    else -> throw IllegalArgumentException(
      "identity hash code for ${jsTypeOf(instance)} is not supported"
    )
  }
}

private var nextHash = 1
private const val IDENTITY_HASHCODE_FIELD = "kotlinIdentityHashcodeValue$"

@NoLiveLiterals
private fun memoizeIdentityHashCode(instance: Any?): Int {
  val value = nextHash++

  val descriptor = js("new Object()")
  descriptor.value = value
  descriptor.writable = false
  descriptor.configurable = false
  descriptor.enumerable = false

  js("Object").defineProperty(instance, IDENTITY_HASHCODE_FIELD, descriptor)

  return value
}

actual annotation class TestOnly

actual inline fun <R> synchronized(lock: Any, block: () -> R): R =
  block()

actual val DefaultMonotonicFrameClock: MonotonicFrameClock = MonotonicClockImpl()

@OptIn(ExperimentalTime::class)
private class MonotonicClockImpl : MonotonicFrameClock {
  override suspend fun <R> withFrameNanos(
    onFrame: (Long) -> R
  ): R = suspendCoroutine { continuation ->
    window.requestAnimationFrame {
      val duration = it.toDuration(DurationUnit.MILLISECONDS)
      val result = onFrame(duration.toLongNanoseconds())
      continuation.resume(result)
    }
  }
}

internal actual fun <T> createSnapshotMutableState(
  value: T,
  policy: SnapshotMutationPolicy<T>
): SnapshotMutableState<T> = SnapshotMutableStateImpl(value, policy)

/*
 * Copyright 2019 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package androidx.compose.runtime

// todo: trace?
internal actual object Trace {
  actual fun beginSection(name: String): Any? {
    return null
  }

  actual fun endSection(token: Any?) {
  }
}

actual annotation class CheckResult actual constructor(actual val suggest: String)

/*
 * Copyright 2020 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package androidx.compose.runtime

import androidx.compose.runtime.snapshots.SnapshotMutableState
import kotlinx.atomicfu.atomic
import kotlinx.coroutines.yield
import kotlin.native.identityHashCode
import kotlin.system.getTimeNanos
import kotlin.time.ExperimentalTime

// TODO actual thread local lol
internal actual open class ThreadLocal<T> actual constructor(
  initialValue: () -> T
) {
  private var value: T = initialValue()

  actual fun get(): T = value

  actual fun set(value: T) {
    this.value = value
  }
}

actual class AtomicReference<V> actual constructor(value: V) {
  private val delegate = atomic(value)

  actual fun get(): V = delegate.value

  actual fun set(value: V) {
    delegate.value = value
  }

  actual fun getAndSet(value: V): V =
    delegate.getAndSet(value)

  actual fun compareAndSet(expect: V, newValue: V): Boolean =
    delegate.compareAndSet(expect, newValue)
}

internal actual fun identityHashCode(instance: Any?): Int =
  instance.identityHashCode()

actual annotation class TestOnly

actual inline fun <R> synchronized(lock: Any, block: () -> R): R =
  block()

actual val DefaultMonotonicFrameClock: MonotonicFrameClock = MonotonicClockImpl()

@OptIn(ExperimentalTime::class)
private class MonotonicClockImpl : MonotonicFrameClock {
  override suspend fun <R> withFrameNanos(
    onFrame: (Long) -> R
  ): R {
    yield()
    return onFrame(getTimeNanos())
  }
}

internal actual object Trace {
  actual fun beginSection(name: String): Any? {
    return null
  }

  actual fun endSection(token: Any?) {
  }
}

actual annotation class CheckResult actual constructor(actual val suggest: String)

internal actual fun <T> createSnapshotMutableState(
  value: T,
  policy: SnapshotMutationPolicy<T>
): SnapshotMutableState<T> =
  SnapshotMutableStateImpl(value, policy)

//fixme: not actually thread local
internal actual class SnapshotThreadLocal<T> actual constructor() {
  private var value: T? = null

  actual fun get(): T? = value
  actual fun set(value: T?) {
    this.value = value
  }
}

/*
 * Copyright 2021 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// This file copied from Compose since it lives inside Compose UI but is general-purpose.
package app.cash.treehouse.compose

import android.os.Looper
import android.view.Choreographer
import androidx.compose.runtime.MonotonicFrameClock
import androidx.core.os.HandlerCompat
import kotlinx.coroutines.CoroutineDispatcher
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.runBlocking
import kotlin.coroutines.CoroutineContext

/**
 * A [CoroutineDispatcher] that will perform dispatch during a [handler] callback or
 * [choreographer]'s animation frame stage, whichever comes first. Use [Main] to obtain
 * a dispatcher for the process's main thread (i.e. the activity thread).
 */
// Implementation note: the constructor is private to direct users toward the companion object
// accessors for the main/current threads. A choreographer must be obtained from its current
// thread as per the only public API surface for obtaining one as of this writing, and the
// choreographer and handler must match. Constructing an AndroidUiDispatcher with a handler
// not marked as async will adversely affect dispatch behavior but not to the point of
// incorrectness; more operations would be deferred to the choreographer frame as racing handler
// messages would wait behind a frame barrier.
public class AndroidUiDispatcher private constructor(
  public val choreographer: Choreographer,
  private val handler: android.os.Handler
) : CoroutineDispatcher() {

  // Guards all properties in this class
  private val lock = Any()

  private val toRunTrampolined = ArrayDeque<Runnable>()
  private var toRunOnFrame = mutableListOf<Choreographer.FrameCallback>()
  private var spareToRunOnFrame = mutableListOf<Choreographer.FrameCallback>()
  private var scheduledTrampolineDispatch = false
  private var scheduledFrameDispatch = false

  private val dispatchCallback = object : Choreographer.FrameCallback, Runnable {
    override fun run() {
      performTrampolineDispatch()
      synchronized(lock) {
        if (toRunOnFrame.isEmpty()) {
          choreographer.removeFrameCallback(this)
          scheduledFrameDispatch = false
        }
      }
    }

    override fun doFrame(frameTimeNanos: Long) {
      handler.removeCallbacks(this)
      performTrampolineDispatch()
      performFrameDispatch(frameTimeNanos)
    }
  }

  private fun nextTask(): Runnable? = synchronized(lock) {
    toRunTrampolined.removeFirstOrNull()
  }

  private fun performTrampolineDispatch() {
    do {
      var task = nextTask()
      while (task != null) {
        task.run()
        task = nextTask()
      }
    } while (
    // We don't dispatch holding the lock so that other tasks can get in on our
    // trampolining time slice, but once we're done, make sure nothing added a new task
    // before we set scheduledDispatch = false, which would prevent the next dispatch
    // from being correctly scheduled. Loop to run these stragglers now.
      synchronized(lock) {
        if (toRunTrampolined.isEmpty()) {
          scheduledTrampolineDispatch = false
          false
        } else true
      }
    )
  }

  private fun performFrameDispatch(frameTimeNanos: Long) {
    val toRun = synchronized(lock) {
      if (!scheduledFrameDispatch) return
      scheduledFrameDispatch = false
      val result = toRunOnFrame
      toRunOnFrame = spareToRunOnFrame
      spareToRunOnFrame = result
      result
    }
    for (i in 0 until toRun.size) {
      // This callback will not and must not throw, see AndroidUiFrameClock
      toRun[i].doFrame(frameTimeNanos)
    }
    toRun.clear()
  }

  internal fun postFrameCallback(callback: Choreographer.FrameCallback) {
    synchronized(lock) {
      toRunOnFrame.add(callback)
      if (!scheduledFrameDispatch) {
        scheduledFrameDispatch = true
        choreographer.postFrameCallback(dispatchCallback)
      }
    }
  }

  internal fun removeFrameCallback(callback: Choreographer.FrameCallback) {
    synchronized(lock) {
      toRunOnFrame.remove(callback)
    }
  }

  /**
   * A [MonotonicFrameClock] associated with this [AndroidUiDispatcher]'s [choreographer]
   * that may be used to await [Choreographer] frame dispatch.
   */
  public val frameClock: MonotonicFrameClock = AndroidUiFrameClock(choreographer)

  override fun dispatch(context: CoroutineContext, block: Runnable) {
    synchronized(lock) {
      toRunTrampolined.addLast(block)
      if (!scheduledTrampolineDispatch) {
        scheduledTrampolineDispatch = true
        handler.post(dispatchCallback)
        if (!scheduledFrameDispatch) {
          scheduledFrameDispatch = true
          choreographer.postFrameCallback(dispatchCallback)
        }
      }
    }
  }

  public companion object {
    /**
     * The [CoroutineContext] containing the [AndroidUiDispatcher] and its [frameClock] for the
     * process's main thread.
     */
    public val Main: CoroutineContext by lazy {
      val dispatcher = AndroidUiDispatcher(
        if (isMainThread()) Choreographer.getInstance()
        else runBlocking(Dispatchers.Main) { Choreographer.getInstance() },
        HandlerCompat.createAsync(Looper.getMainLooper())
      )

      dispatcher + dispatcher.frameClock
    }
  }
}

private fun isMainThread() = Looper.myLooper() === Looper.getMainLooper()

/*
 * Copyright 2021 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// This file copied from Compose since it lives inside Compose UI but is general-purpose.
package app.cash.treehouse.compose

import android.view.Choreographer
import androidx.compose.runtime.MonotonicFrameClock
import kotlinx.coroutines.suspendCancellableCoroutine
import kotlin.coroutines.ContinuationInterceptor
import kotlin.coroutines.coroutineContext

public class AndroidUiFrameClock(
  private val choreographer: Choreographer
) : MonotonicFrameClock {
  override suspend fun <R> withFrameNanos(
    onFrame: (Long) -> R
  ): R {
    val uiDispatcher = coroutineContext[ContinuationInterceptor] as? AndroidUiDispatcher
    return suspendCancellableCoroutine { co ->
      // Important: this callback won't throw, and AndroidUiDispatcher counts on it.
      val callback = Choreographer.FrameCallback { frameTimeNanos ->
        co.resumeWith(runCatching { onFrame(frameTimeNanos) })
      }

      // If we're on an AndroidUiDispatcher then we post callback to happen *after*
      // the greedy trampoline dispatch is complete.
      // This means that onFrame will run on the current choreographer frame if one is
      // already in progress, but withFrameNanos will *not* resume until the frame
      // is complete. This prevents multiple calls to withFrameNanos immediately dispatching
      // on the same frame.

      if (uiDispatcher != null && uiDispatcher.choreographer == choreographer) {
        uiDispatcher.postFrameCallback(callback)
        co.invokeOnCancellation { uiDispatcher.removeFrameCallback(callback) }
      } else {
        choreographer.postFrameCallback(callback)
        co.invokeOnCancellation { choreographer.removeFrameCallback(callback) }
      }
    }
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.compose

import androidx.compose.runtime.AbstractApplier
import androidx.compose.runtime.Applier
import androidx.compose.runtime.Composable
import androidx.compose.runtime.ComposeNode
import app.cash.treehouse.protocol.ChildrenDiff
import app.cash.treehouse.protocol.ChildrenDiff.Companion.RootChildrenTag
import app.cash.treehouse.protocol.ChildrenDiff.Companion.RootId
import app.cash.treehouse.protocol.Diff
import app.cash.treehouse.protocol.Event
import app.cash.treehouse.protocol.PropertyDiff

/**
 * A synthetic node which allows the applier to differentiate between multiple groups of children.
 *
 * Compose's tree assumes each node only has single list of children. Or, put another way, even if
 * you apply multiple children Compose treats them as a single list of child nodes. In order to
 * differentiate between these children lists we introduce synthetic nodes. Every real node which
 * supports one or more groups of children will have one or more of these synthetic nodes as its
 * direct descendants. The nodes which are produced by each group of children will then become the
 * descendants of those synthetic nodes.
 *
 * This function is named weirdly to prevent normal usage since bad things will happen.
 *
 * @see ProtocolApplier
 * @suppress
 */
@Composable
public fun `$SyntheticChildren`(tag: Int, content: @Composable () -> Unit) {
  ComposeNode<ProtocolChildrenNode.Intermediate, Applier<ProtocolNode>>(
    factory = {
      ProtocolChildrenNode.Intermediate(tag)
    },
    update = {
    },
    content = content,
  )
}

/**
 * A node which exists in the tree to emulate supporting multiple children sets but which does not
 * appear directly in the protocol. The ID of these nodes mirrors that of its parent to simplify
 * creation of the protocol diffs. This is safe because these types never appears in the node map.
 */
private sealed class ProtocolChildrenNode(
  val tag: Int,
) : ProtocolNode(-1) {
  class Intermediate(tag: Int) : ProtocolChildrenNode(tag)
  class Root : ProtocolChildrenNode(RootChildrenTag) {
    init {
      id = RootId
    }
  }
}

/**
 * @suppress
 */
public open class ProtocolNode(
  public val type: Int,
) {
  public var id: Long = -1
    internal set

  internal lateinit var applier: ProtocolApplier

  internal val children = mutableListOf<ProtocolNode>()

  public fun appendDiff(diff: PropertyDiff) {
    applier.appendDiff(diff)
  }

  public open fun sendEvent(event: Event) {
    throw IllegalStateException("Node ID $id of type $type does not handle events")
  }
}

/**
 * An [Applier] which records operations on the tree as models which can then be separately applied
 * by the display layer. Additionally, it has special handling for emulating nodes which contain
 * multiple children.
 *
 * Nodes in the tree are required to alternate between [ProtocolChildrenNode] instances and
 * non-[ProtocolChildrenNode] [ProtocolNode] subtypes starting from the root. This invariant is
 * maintained by virtue of the fact that all of the input `@Composables` should be generated code.
 *
 * For example, a node tree may look like this:
 * ```
 *                    Children(tag=1)
 *                     /          \
 *                    /            \
 *            ToolbarNode        ListNode
 *             /     \                 \
 *            /       \                 \
 * Children(tag=1)  Children(tag=2)   Children(tag=1)
 *        |              |               /       \
 *        |              |              /         \
 *   ButtonNode     ButtonNode     TextNode     TextNode
 * ```
 * But the protocol diff output would only record non-[ProtocolChildrenNode] nodes using their
 * [ProtocolChildrenNode.tag] value:
 * ```
 * Insert(id=<root-id>, tag=1, type=<toolbar-type>, childId=<toolbar-id>)
 * Insert(id=<toolbar-id>, tag=1, type=<button-type>, childId=..)
 * Insert(id=<toolbar-id>, tag=2, type=<button-type>, childId=..)
 * Insert(id=<root-id>, tag=1, type=<list-type>, childId=<list-id>)
 * Insert(id=<list-id>, tag=1, type=<text-type>, childId=..)
 * Insert(id=<list-id>, tag=1, type=<text-type>, childId=..)
 * ```
 */
internal class ProtocolApplier(
  private val onDiff: (Diff) -> Unit,
) : AbstractApplier<ProtocolNode>(ProtocolChildrenNode.Root()) {
  private var nextId = RootId + 1
  val nodes = mutableMapOf(root.id to root)
  private var childrenDiffs = mutableListOf<ChildrenDiff>()
  private var propertyDiffs = mutableListOf<PropertyDiff>()

  fun appendDiff(diff: PropertyDiff) {
    propertyDiffs.add(diff)
  }

  override fun onEndChanges() {
    val existingChildrenDiffs = childrenDiffs
    val existingPropertyDiffs = propertyDiffs
    if (existingPropertyDiffs.isNotEmpty() || existingChildrenDiffs.isNotEmpty()) {
      childrenDiffs = mutableListOf()
      propertyDiffs = mutableListOf()

      val diff = Diff(
        childrenDiffs = existingChildrenDiffs,
        propertyDiffs = existingPropertyDiffs,
      )
      onDiff(diff)
    }
  }

  override fun insertTopDown(index: Int, instance: ProtocolNode) {
    current.children.add(index, instance)

    if (instance is ProtocolChildrenNode) {
      // Inherit the ID from the current node such that changes to the children can be reported
      // as if they occurred directly on the parent.
      instance.id = current.id
      // We do not add children instances to the map (they have no unique IDs and are only
      // available through indexing on the parent) and we do not send them over the wire to the
      // display (they are always implied by the display interfaces).
    } else {
      val current = current as ProtocolChildrenNode

      val id = nextId++
      instance.id = id
      instance.applier = this

      nodes[id] = instance
      childrenDiffs.add(ChildrenDiff.Insert(current.id, current.tag, id, instance.type, index))
    }
  }

  override fun insertBottomUp(index: Int, instance: ProtocolNode) {
    // Ignored, we insert top-down for now.
  }

  override fun remove(index: Int, count: Int) {
    // Children instances are never removed from their parents.
    val current = current as ProtocolChildrenNode

    val children = current.children
    for (i in index until index + count) {
      nodes.remove(children[i].id)
    }
    children.remove(index, count)
    childrenDiffs.add(ChildrenDiff.Remove(current.id, current.tag, index, count))
  }

  override fun move(from: Int, to: Int, count: Int) {
    // Children instances are never moved within their parents.
    val current = current as ProtocolChildrenNode

    current.children.move(from, to, count)
    childrenDiffs.add(ChildrenDiff.Move(current.id, current.tag, from, to, count))
  }

  override fun onClear() {
    current.children.clear()
    nodes.clear()
    nodes[current.id] = current // Restore root node into map.
    childrenDiffs.add(ChildrenDiff.Clear)
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.compose

import androidx.compose.runtime.Composable
import androidx.compose.runtime.Composition
import androidx.compose.runtime.MonotonicFrameClock
import androidx.compose.runtime.Recomposer
import androidx.compose.runtime.snapshots.ObserverHandle
import androidx.compose.runtime.snapshots.Snapshot
import app.cash.treehouse.protocol.Diff
import app.cash.treehouse.protocol.Event
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.CoroutineStart.UNDISPATCHED
import kotlinx.coroutines.Job
import kotlinx.coroutines.launch

public interface TreehouseComposition {
  public fun sendEvent(event: Event)
  public fun setContent(content: @Composable () -> Unit)
  public fun cancel()
}

/**
 * @param scope A [CoroutineScope] whose [coroutineContext][kotlin.coroutines.CoroutineContext]
 * must have a [MonotonicFrameClock] key which is being ticked.
 */
public fun TreehouseComposition(
  scope: CoroutineScope,
  display: (diff: Diff, events: (Event) -> Unit) -> Unit,
  onDiff: (Diff) -> Unit = {},
  onEvent: (Event) -> Unit = {},
): TreehouseComposition {
  val server = RealTreehouseComposition(scope, display, onDiff, onEvent)
  server.launch()
  return server
}

private class RealTreehouseComposition(
  private val scope: CoroutineScope,
  private val display: (Diff, (Event) -> Unit) -> Unit,
  private val onDiff: (Diff) -> Unit,
  private val onEvent: (Event) -> Unit,
) : TreehouseComposition {
  // Caching type conversion of function reference to lambda to avoid future allocations.
  private val eventSink: (Event) -> Unit = ::sendEvent
  private val applier = ProtocolApplier { diff ->
    onDiff(diff)
    display(diff, eventSink)
  }

  private val recomposer = Recomposer(scope.coroutineContext)
  private val composition = Composition(applier, recomposer)

  private lateinit var snapshotHandle: ObserverHandle
  private var snapshotJob: Job? = null
  private lateinit var recomposeJob: Job

  fun launch() {
    // Set up a trigger to apply changes on the next frame if a global write was observed.
    // TODO where should this live?
    var applyScheduled = false
    snapshotHandle = Snapshot.registerGlobalWriteObserver {
      if (!applyScheduled) {
        applyScheduled = true
        snapshotJob = scope.launch {
          applyScheduled = false
          Snapshot.sendApplyNotifications()
        }
      }
    }

    // These launch undispatched so that they reach their first suspension points before returning
    // control to the caller.
    recomposeJob = scope.launch(start = UNDISPATCHED) {
      recomposer.runRecomposeAndApplyChanges()
    }
  }

  override fun sendEvent(event: Event) {
    onEvent(event)

    val node = applier.nodes[event.id]
    if (node == null) {
      // TODO how to handle race where an incoming event targets this removed node?
      throw IllegalArgumentException("Unknown node ${event.id} for event with tag ${event.tag}")
    }
    node.sendEvent(event)
  }

  override fun setContent(content: @Composable () -> Unit) {
    composition.setContent(content)
  }

  override fun cancel() {
    snapshotHandle.dispose()
    snapshotJob?.cancel()
    recomposeJob.cancel()
    recomposer.cancel()
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.compose

import androidx.compose.runtime.MonotonicFrameClock
import kotlinx.browser.window
import kotlin.coroutines.resume
import kotlin.coroutines.suspendCoroutine

public object WindowAnimationFrameClock : MonotonicFrameClock {
  override suspend fun <R> withFrameNanos(
    onFrame: (Long) -> R
  ): R = suspendCoroutine { continuation ->
    window.requestAnimationFrame {
      val durationMillis = it.toLong()
      val durationNanos = durationMillis * 1000000
      val result = onFrame(durationNanos)
      continuation.resume(result)
    }
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.gradle

import app.cash.exhaustive.Exhaustive
import org.gradle.api.Project
import org.gradle.api.provider.Provider
import org.jetbrains.kotlin.gradle.dsl.KotlinJvmOptions
import org.jetbrains.kotlin.gradle.plugin.KotlinCompilation
import org.jetbrains.kotlin.gradle.plugin.KotlinCompilerPluginSupportPlugin
import org.jetbrains.kotlin.gradle.plugin.KotlinPlatformType.androidJvm
import org.jetbrains.kotlin.gradle.plugin.KotlinPlatformType.common
import org.jetbrains.kotlin.gradle.plugin.KotlinPlatformType.js
import org.jetbrains.kotlin.gradle.plugin.KotlinPlatformType.jvm
import org.jetbrains.kotlin.gradle.plugin.KotlinPlatformType.native
import org.jetbrains.kotlin.gradle.plugin.SubpluginArtifact
import org.jetbrains.kotlin.gradle.plugin.SubpluginOption

public class TreehousePlugin : KotlinCompilerPluginSupportPlugin {
  override fun apply(project: Project) {
    super.apply(project)

    project.configurations.all { configuration ->
      configuration.resolutionStrategy.dependencySubstitution { it ->
        val composeRuntime = it.module("androidx.compose.runtime:runtime:$composeVersion")
        val justification =
          "Our Android jar includes the same types in the same package as AndroidX. " +
            "This prevents using other Compose-based libraries such as Compose UI in an app. " +
            "Prefer the canonical version so that dependency resolution will de-duplicate them."

        it.substitute(it.module("app.cash.treehouse:compose-runtime-android"))
          .using(composeRuntime)
          .because(justification)
        it.substitute(it.module("app.cash.treehouse:compose-runtime-android-debug"))
          .using(composeRuntime)
          .because(justification)
      }
    }
  }

  override fun isApplicable(kotlinCompilation: KotlinCompilation<*>): Boolean = true

  override fun getCompilerPluginId(): String = "app.cash.treehouse"

  override fun getPluginArtifact(): SubpluginArtifact = SubpluginArtifact(
    "androidx.compose.compiler",
    "compiler",
    composeVersion,
  )

  override fun getPluginArtifactForNative(): SubpluginArtifact = SubpluginArtifact(
    "app.cash.treehouse",
    "compose-compiler-hosted",
    treehouseVersion,
  )

  override fun applyToCompilation(kotlinCompilation: KotlinCompilation<*>): Provider<List<SubpluginOption>> {
    kotlinCompilation.dependencies {
      implementation("app.cash.treehouse:treehouse-compose:$treehouseVersion")
    }

    @Exhaustive when (kotlinCompilation.platformType) {
      androidJvm, jvm -> {
        if ((kotlinCompilation.kotlinOptions as KotlinJvmOptions).useOldBackend) {
          throw IllegalStateException("Treehouse only works with the default IR-based backend")
        }
      }
      js -> {
        // This enables a workaround for Compose lambda generation to function correctly in JS.
        kotlinCompilation.kotlinOptions.freeCompilerArgs +=
          listOf("-P", "plugin:androidx.compose.compiler.plugins.kotlin:generateDecoys=true")
      }
      common, native -> {
        // Nothing to do!
      }
    }

    return kotlinCompilation.target.project.provider { emptyList() }
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.gradle

import org.gradle.api.Project

public interface TreehouseSchemaExtension {
  public var source: Project?
  public var type: String?
}

// Gradle requires this type to be open since it runtime extends it.
internal open class TreehouseSchemaExtensionImpl @JvmOverloads constructor(
  override var source: Project? = null,
  override var type: String? = null,
) : TreehouseSchemaExtension

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.gradle

import app.cash.treehouse.gradle.TreehouseSchemaGeneratorPlugin.Strategy.Compose
import app.cash.treehouse.gradle.TreehouseSchemaGeneratorPlugin.Strategy.Test
import app.cash.treehouse.gradle.TreehouseSchemaGeneratorPlugin.Strategy.Widget
import org.gradle.api.Plugin
import org.gradle.api.Project
import org.gradle.api.attributes.Usage
import org.gradle.api.attributes.Usage.JAVA_RUNTIME
import org.gradle.api.attributes.Usage.USAGE_ATTRIBUTE
import org.gradle.api.tasks.JavaExec
import org.jetbrains.kotlin.gradle.dsl.KotlinMultiplatformExtension
import org.jetbrains.kotlin.gradle.plugin.KotlinPlatformType
import java.io.File

@Suppress("unused") // Invoked reflectively by Gradle.
public class TreehouseSchemaComposePlugin : TreehouseSchemaGeneratorPlugin(Compose)

@Suppress("unused") // Invoked reflectively by Gradle.
public class TreehouseSchemaTestPlugin : TreehouseSchemaGeneratorPlugin(Test)

@Suppress("unused") // Invoked reflectively by Gradle.
public class TreehouseSchemaWidgetPlugin : TreehouseSchemaGeneratorPlugin(Widget)

public abstract class TreehouseSchemaGeneratorPlugin(
  private val strategy: Strategy,
) : Plugin<Project> {
  public enum class Strategy(
    internal val generatorFlag: String,
    internal val dependencyCoordinate: String,
  ) {
    Compose("--compose", "app.cash.treehouse:treehouse-compose:$treehouseVersion"),
    Test("--test", "app.cash.treehouse:treehouse-widget:$treehouseVersion"),
    Widget("--widget", "app.cash.treehouse:treehouse-widget:$treehouseVersion"),
  }

  override fun apply(project: Project) {
    var applied = false

    if (strategy == Compose) {
      project.plugins.apply(TreehousePlugin::class.java)
    }

    val extension = project.extensions.create(
      TreehouseSchemaExtension::class.java,
      "treehouseSchema",
      TreehouseSchemaExtensionImpl::class.java,
    )

    val configuration = project.configurations.create("treehouseSchema") {
      // Ensure we get JVM artifacts from any multiplatform dependencies for use with JavaExec.
      val runtimeUsage = project.objects.named(Usage::class.java, JAVA_RUNTIME)
      it.attributes.attribute(USAGE_ATTRIBUTE, runtimeUsage)
      it.attributes.attribute(KotlinPlatformType.attribute, KotlinPlatformType.jvm)
    }
    project.dependencies.add(
      configuration.name,
      "app.cash.treehouse:treehouse-schema-generator:$treehouseVersion",
    )

    val generatedDir = File(project.buildDir, "generated/treehouse")
    val generate = project.tasks.register("treehouseGenerate", JavaExec::class.java) { exec ->
      exec.outputs.dir(generatedDir)

      exec.classpath(configuration)
      exec.main = "app.cash.treehouse.schema.generator.Main"

      exec.doFirst {
        generatedDir.deleteRecursively()
      }
    }

    project.plugins.withId("org.jetbrains.kotlin.multiplatform") {
      applied = true
    }

    project.afterEvaluate {
      check(applied) {
        "Treehouse schema plugin requires the Kotlin multiplatform plugin to be applied."
      }

      val schemaType = requireNotNull(extension.type) {
        "Treehouse schema type name must be specified!"
      }
      val schemaProject = requireNotNull(extension.source) {
        "Treehouse schema project must be specified!"
      }

      project.dependencies.add(configuration.name, schemaProject)

      generate.configure {
        it.args = listOf(
          strategy.generatorFlag,
          "--out", generatedDir.toString(),
          schemaType
        )
      }

      val kotlin = project.extensions.getByType(KotlinMultiplatformExtension::class.java)

      kotlin.sourceSets.getByName("commonMain") { sourceSet ->
        sourceSet.kotlin.srcDir(generatedDir)
        sourceSet.dependencies {
          api(strategy.dependencyCoordinate)

          if (strategy == Test) {
            api(schemaProject)
          }
        }
      }

      kotlin.targets.all { target ->
        target.compilations.all { compilation ->
          compilation.compileKotlinTaskProvider.configure {
            it.dependsOn(generate.get())
          }
        }
      }
    }
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.gradle

import org.gradle.api.Plugin
import org.gradle.api.Project
import org.jetbrains.kotlin.gradle.dsl.KotlinJvmProjectExtension
import org.jetbrains.kotlin.gradle.dsl.KotlinMultiplatformExtension
import org.jetbrains.kotlin.gradle.plugin.KotlinPlatformType
import org.jetbrains.kotlin.gradle.plugin.KotlinTarget

@Suppress("unused") // Invoked reflectively by Gradle.
public class TreehouseSchemaPlugin : Plugin<Project> {
  override fun apply(project: Project) {
    var applied = false

    project.plugins.withId("org.jetbrains.kotlin.jvm") {
      applied = true

      val kotlin = project.extensions.getByType(KotlinJvmProjectExtension::class.java)
      kotlin.target.applySchemaAnnotationDependency()
    }
    project.plugins.withId("org.jetbrains.kotlin.multiplatform") {
      applied = true

      val kotlin = project.extensions.getByType(KotlinMultiplatformExtension::class.java)
      kotlin.targets.all { it.applySchemaAnnotationDependency() }

      project.afterEvaluate {
        check(kotlin.targets.any { it.platformType == KotlinPlatformType.jvm }) {
          "Treehouse schema plugin requires a jvm() target when used with Kotlin multiplatform"
        }
      }
    }

    project.afterEvaluate {
      check(applied) {
        "Treehouse schema plugin requires the Kotlin JVM or multiplatform plugin to be applied."
      }
    }
  }

  private fun KotlinTarget.applySchemaAnnotationDependency() {
    compilations.getByName("main") { compilation ->
      compilation.dependencies {
        api("app.cash.treehouse:treehouse-schema-annotations:$treehouseVersion")
      }
    }
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.protocol

import kotlinx.serialization.Polymorphic
import kotlinx.serialization.SerialName
import kotlinx.serialization.Serializable

@Serializable
public data class Event(
  /** Identifier for the widget from which this event originated. */
  val id: Long,
  /** Identifies which event occurred on the widget with [id]. */
  val tag: Int,
  @Polymorphic val value: Any? = null,
)

@Serializable
public data class Diff(
  val childrenDiffs: List<ChildrenDiff> = emptyList(),
  val propertyDiffs: List<PropertyDiff> = emptyList()
)

@Serializable
public data class PropertyDiff(
  /** Identifier for the widget whose property has changed. */
  val id: Long,
  /** Identifies which property changed on the widget with [id]. */
  val tag: Int,
  @Polymorphic val value: Any? = null,
)

@Serializable
public sealed class ChildrenDiff {
  /** Identifier for the widget whose children have changed. */
  public abstract val id: Long
  /** Identifies which group of children changed on the widget with [id]. */
  public abstract val tag: Int

  @Serializable
  @SerialName("clear")
  public object Clear : ChildrenDiff() {
    override val id: Long get() = RootId
    override val tag: Int get() = RootChildrenTag
  }

  @Serializable
  @SerialName("insert")
  public data class Insert(
    override val id: Long,
    override val tag: Int,
    val childId: Long,
    val kind: Int,
    val index: Int,
  ) : ChildrenDiff()

  @Serializable
  @SerialName("move")
  public data class Move(
    override val id: Long,
    override val tag: Int,
    val fromIndex: Int,
    val toIndex: Int,
    val count: Int,
  ) : ChildrenDiff()

  @Serializable
  @SerialName("remove")
  public data class Remove(
    override val id: Long,
    override val tag: Int,
    val index: Int,
    val count: Int,
  ) : ChildrenDiff()

  public companion object {
    public const val RootId: Long = 0L
    public const val RootChildrenTag: Int = 1
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.schema

import kotlin.reflect.KClass

/**
 * Annotates an otherwise unused type with a set of the [Widget]-annotated classes which are all
 * part of this schema.
 *
 * ```
 * @Schema([
 *   Box::class,
 *   Button::class,
 *   Text::class,
 * ])
 * interface Name
 * ```
 */
public annotation class Schema(val widgets: Array<KClass<*>>)

/**
 * Annotates a data class which represents a widget in a UI tree. Each widget in a [Schema] must
 * have a unique [tag] among all [@Widget][Widget] annotations in the [Schema].
 *
 * All of the properties in the class must be annotated with either [Property] or [Children].
 *
 * ```
 * @Widget(1)
 * data class Box(
 *   @Property(1) val orientation: Orientation,
 *   @Children val children: List<Any>,
 * )
 * ```
 */
public annotation class Widget(val tag: Int)

/**
 * Annotates a [Widget] property which represents a property on the associated UI widget. Properties
 * in a [Widget] class must have a unique [tag] among all [@Property][Property] annotations in
 * the class.
 *
 * ```
 * @Widget(1)
 * data class Text(
 *   @Property(1) val text: String,
 * )
 * ```
 */
public annotation class Property(val tag: Int)

/**
 * Annotates a [Widget] property as representing child widgets which are contained within the
 * enclosing widget. Children in a [Widget] class must have a unique [tag] among all
 * [@Children][Children] annotations in the class. The type of the property must be `List<Any>`.
 *
 * ```
 * @Widget(1)
 * data class Box(
 *   @Children(1) val children: List<Any>,
 * )
 * ```
 */
public annotation class Children(val tag: Int)

/**
 * Annotates a [Property] with an associated default expression. The [expression] is not
 * type-checked and will be used verbatim in the generated code.
 *
 * ```
 * @Widget(1)
 * data class Button(
 *   @Property(1) val text: String,
 *   @Property(2) @Default("true") val enabled: Boolean,
 * )
 * ```
 */
public annotation class Default(val expression: String)

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.widget

import android.view.View
import android.view.ViewGroup

public class ViewGroupChildren(private val parent: ViewGroup) : Widget.Children<View> {
  override fun insert(index: Int, widget: View) {
    parent.addView(widget, index)
  }

  override fun move(fromIndex: Int, toIndex: Int, count: Int) {
    val views = Array(count) { offset ->
      parent.getChildAt(fromIndex + offset)
    }
    parent.removeViews(fromIndex, count)

    val newIndex = if (toIndex > fromIndex) {
      toIndex - count
    } else {
      toIndex
    }
    views.forEachIndexed { offset, view ->
      parent.addView(view, newIndex + offset)
    }
  }

  override fun remove(index: Int, count: Int) {
    parent.removeViews(index, count)
  }

  override fun clear() {
    parent.removeAllViews()
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.widget

import kotlin.jvm.JvmOverloads

public class MutableListChildren<T : Any>
@JvmOverloads constructor(
  public val list: MutableList<T> = mutableListOf(),
) : Widget.Children<T>, Iterable<T> {
  override fun insert(index: Int, widget: T) {
    list.add(index, widget)
  }

  override fun move(fromIndex: Int, toIndex: Int, count: Int) {
    /*
     * Copyright 2019 The Android Open Source Project
     *
     * Licensed under the Apache License, Version 2.0 (the "License");
     * you may not use this file except in compliance with the License.
     * You may obtain a copy of the License at
     *
     *      http://www.apache.org/licenses/LICENSE-2.0
     *
     * Unless required by applicable law or agreed to in writing, software
     * distributed under the License is distributed on an "AS IS" BASIS,
     * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     * See the License for the specific language governing permissions and
     * limitations under the License.
     */
    val dest = if (fromIndex > toIndex) toIndex else toIndex - count
    if (count == 1) {
      if (fromIndex == toIndex + 1 || fromIndex == toIndex - 1) {
        // Adjacent elements, perform swap to avoid backing array manipulations.
        val fromEl = list.get(fromIndex)
        val toEl = list.set(toIndex, fromEl)
        list.set(fromIndex, toEl)
      } else {
        val fromEl = list.removeAt(fromIndex)
        list.add(dest, fromEl)
      }
    } else {
      val subView = list.subList(fromIndex, fromIndex + count)
      val subCopy = subView.toMutableList()
      subView.clear()
      list.addAll(dest, subCopy)
    }
  }

  override fun remove(index: Int, count: Int) {
    /*
     * Copyright 2019 The Android Open Source Project
     *
     * Licensed under the Apache License, Version 2.0 (the "License");
     * you may not use this file except in compliance with the License.
     * You may obtain a copy of the License at
     *
     *      http://www.apache.org/licenses/LICENSE-2.0
     *
     * Unless required by applicable law or agreed to in writing, software
     * distributed under the License is distributed on an "AS IS" BASIS,
     * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     * See the License for the specific language governing permissions and
     * limitations under the License.
     */
    if (count == 1) {
      list.removeAt(index)
    } else {
      list.subList(index, index + count).clear()
    }
  }

  override fun clear() {
    list.clear()
  }

  override fun iterator(): Iterator<T> = list.iterator()
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.widget

import app.cash.treehouse.protocol.Event
import app.cash.treehouse.protocol.PropertyDiff

public interface Widget<T : Any> {
  public val value: T

  public fun apply(diff: PropertyDiff, events: (Event) -> Unit)

  public fun children(tag: Int): Children<T> {
    throw IllegalArgumentException("Widget does not support children")
  }

  /**
   * An interface for manipulating a widget's list of child widgets.
   *
   * Arguments to these methods can be assumed to be validated against the current state of the
   * list. No additional validation needs to be performed (for example, checking index bounds).
   */
  public interface Children<T : Any> {
    /** Insert child [widget] at [index]. */
    public fun insert(index: Int, widget: T)
    /**
     * Move [count] child widgets from [fromIndex] to [toIndex].
     *
     * Both indices are relative to the position before the change. For example, to move the
     * widget at position 1 to after the widget at position 2, [fromIndex] should be 1 and
     * [toIndex] should be 3. If the widgets were `A B C D E`, calling `move(1, 3, 1)` would
     * result in the widgets being reordered to `A C B D E`.
     */
    public fun move(fromIndex: Int, toIndex: Int, count: Int)
    /** Remove [count] child widgets starting from [index]. */
    public fun remove(index: Int, count: Int)
    /** Remove all child widgets. */
    public fun clear()
  }

  public interface Factory<T : Any> {
    public fun create(kind: Int, id: Long): Widget<T>
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.widget

import app.cash.treehouse.protocol.ChildrenDiff
import app.cash.treehouse.protocol.ChildrenDiff.Companion.RootChildrenTag
import app.cash.treehouse.protocol.ChildrenDiff.Companion.RootId
import app.cash.treehouse.protocol.Diff
import app.cash.treehouse.protocol.Event

public interface Display {
  public fun apply(diff: Diff, events: (Event) -> Unit)
}

public class WidgetDisplay<T : Any>(
  private val root: Widget<T>,
  private val factory: Widget.Factory<T>,
) : Display {
  init {
    // Check that the root widget has a group of children with the shared root tag. This call
    // will throw if that invariant does not hold.
    root.children(RootChildrenTag)
  }

  private val widgets = mutableMapOf(RootId to root)

  override fun apply(diff: Diff, events: (Event) -> Unit) {
    for (childrenDiff in diff.childrenDiffs) {
      val widget = checkNotNull(widgets[childrenDiff.id]) {
        "Unknown widget ID ${childrenDiff.id}"
      }
      val children = widget.children(childrenDiff.tag)

      when (childrenDiff) {
        is ChildrenDiff.Insert -> {
          val childWidget = factory.create(childrenDiff.kind, childrenDiff.childId)
          widgets[childrenDiff.childId] = childWidget
          children.insert(childrenDiff.index, childWidget.value)
        }
        is ChildrenDiff.Move -> {
          children.move(childrenDiff.fromIndex, childrenDiff.toIndex, childrenDiff.count)
        }
        is ChildrenDiff.Remove -> {
          children.remove(childrenDiff.index, childrenDiff.count)
          // TODO we need to remove widgets from our map!
        }
        ChildrenDiff.Clear -> {
          children.clear()
          widgets.clear()
          widgets[RootId] = root
        }
      }
    }

    for (propertyDiff in diff.propertyDiffs) {
      val widget = checkNotNull(widgets[propertyDiff.id]) {
        "Unknown widget ID ${propertyDiff.id}"
      }

      widget.apply(propertyDiff, events)
    }
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.widget

import kotlinx.dom.clear
import org.w3c.dom.HTMLElement
import org.w3c.dom.get

public class HTMLElementChildren(
  private val parent: HTMLElement,
) : Widget.Children<HTMLElement> {
  override fun insert(index: Int, widget: HTMLElement) {
    // Null element returned when index == childCount causes insertion at end.
    val current = parent.children[index]
    parent.insertBefore(widget, current)
  }

  override fun move(fromIndex: Int, toIndex: Int, count: Int) {
    val elements = Array(count) {
      val element = parent.children[fromIndex] as HTMLElement
      parent.removeChild(element)
      element
    }

    val newIndex = if (toIndex > fromIndex) {
      toIndex - count
    } else {
      toIndex
    }
    elements.forEachIndexed { offset, element ->
      // Null element returned when newIndex + offset == childCount causes insertion at end.
      val current = parent.children[newIndex + offset]
      parent.insertBefore(element, current)
    }
  }

  override fun remove(index: Int, count: Int) {
    repeat(count) {
      parent.removeChild(parent.children[index]!!)
    }
  }

  override fun clear() {
    parent.clear()
  }
}

/*
 * Copyright 2020 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

@file:OptIn(InternalComposeApi::class)
package androidx.compose.runtime.internal

import androidx.compose.runtime.ComposeCompilerApi
import androidx.compose.runtime.Composer
import androidx.compose.runtime.InternalComposeApi
import androidx.compose.runtime.RecomposeScope
import androidx.compose.runtime.Stable

/**
 * A Restart is created to hold composable lambdas to track when they are invoked allowing
 * the invocations to be invalidated when a new composable lambda is created during composition.
 *
 * This allows much of the call-graph to be skipped when a composable function is passed through
 * multiple levels of composable functions.
 */
@Suppress("NAME_SHADOWING")
@Stable
@OptIn(ComposeCompilerApi::class)
/* ktlint-disable parameter-list-wrapping */ // TODO(https://github.com/pinterest/ktlint/issues/921): reenable
internal actual class ComposableLambdaImpl actual constructor(
  val key: Int,
  private val tracked: Boolean
) : ComposableLambda {
  private var _block: Any? = null
  private var scope: RecomposeScope? = null
  private var scopes: MutableList<RecomposeScope>? = null

  private fun trackWrite() {
    if (tracked) {
      val scope = this.scope
      if (scope != null) {
        scope.invalidate()
        this.scope = null
      }
      val scopes = this.scopes
      if (scopes != null) {
        for (index in 0 until scopes.size) {
          val item = scopes[index]
          item.invalidate()
        }
        scopes.clear()
      }
    }
  }

  private fun trackRead(composer: Composer) {
    if (tracked) {
      val scope = composer.recomposeScope
      if (scope != null) {
        // Find the first invalid scope and replace it or record it if no scopes are invalid
        composer.recordUsed(scope)
        val lastScope = this.scope
        if (lastScope.replacableWith(scope)) {
          this.scope = scope
        } else {
          val lastScopes = scopes
          if (lastScopes == null) {
            val newScopes = mutableListOf<RecomposeScope>()
            scopes = newScopes
            newScopes.add(scope)
          } else {
            for (index in 0 until lastScopes.size) {
              val scopeAtIndex = lastScopes[index]
              if (scopeAtIndex.replacableWith(scope)) {
                lastScopes[index] = scope
                return
              }
            }
            lastScopes.add(scope)
          }
        }
      }
    }
  }

  actual fun update(block: Any) {
    if (_block != block) {
      val oldBlockNull = _block == null
      _block = block
      if (!oldBlockNull) {
        trackWrite()
      }
    }
  }

  override operator fun invoke(c: Composer, changed: Int): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(0) else sameBits(0)
    val result = (_block as (c: Composer, changed: Int) -> Any?)(c, dirty)
    c.endRestartGroup()?.updateScope(this::invoke)
    return result
  }

  override operator fun invoke(p1: Any?, c: Composer, changed: Int): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(1) else sameBits(1)
    val result = (
      _block as (
        p1: Any?,
        c: Composer,
        changed: Int
      ) -> Any?
      )(
      p1,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ -> this(p1, nc, changed or 0b1) }
    return result
  }

  override operator fun invoke(p1: Any?, p2: Any?, c: Composer, changed: Int): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(2) else sameBits(2)
    val result = (_block as (p1: Any?, p2: Any?, c: Composer, changed: Int) -> Any?)(
      p1,
      p2,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ -> this(p1, p2, nc, changed or 0b1) }
    return result
  }

  override operator fun invoke(p1: Any?, p2: Any?, p3: Any?, c: Composer, changed: Int): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(3) else sameBits(3)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        c: Composer,
        changed: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ -> this(p1, p2, p3, nc, changed or 0b1) }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    c: Composer,
    changed: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(4) else sameBits(4)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        c: Composer,
        changed: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, nc, changed or 0b1)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    c: Composer,
    changed: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(5) else sameBits(5)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        c: Composer,
        changed: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, p5, nc, changed or 0b1)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    c: Composer,
    changed: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(6) else sameBits(6)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        c: Composer,
        changed: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, p5, p6, nc, changed or 0b1)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    c: Composer,
    changed: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(7) else sameBits(7)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        c: Composer,
        changed: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, p5, p6, p7, nc, changed or 0b1)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    c: Composer,
    changed: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(8) else sameBits(8)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        c: Composer,
        changed: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, p5, p6, p7, p8, nc, changed or 0b1)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    c: Composer,
    changed: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(9) else sameBits(9)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        c: Composer,
        changed: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, p5, p6, p7, p8, p9, nc, changed or 0b1)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(10) else sameBits(10)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, p5, p6, p7, p8, p9, p10, nc, changed or 0b1, changed)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(11) else sameBits(11)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        p11: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      p11,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, p5, p6, p7, p8, p9, p10, p11, nc, changed or 0b1, changed1)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(12) else sameBits(12)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        p11: Any?,
        p12: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      p11,
      p12,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, p5, p6, p7, p8, p9, p10, p11, p12, nc, changed or 0b1, changed1)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(13) else sameBits(13)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        p11: Any?,
        p12: Any?,
        p13: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      p11,
      p12,
      p13,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(
        p1,
        p2,
        p3,
        p4,
        p5,
        p6,
        p7,
        p8,
        p9,
        p10,
        p11,
        p12,
        p13,
        nc,
        changed or 0b1,
        changed1
      )
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    p14: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(14) else sameBits(14)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        p11: Any?,
        p12: Any?,
        p13: Any?,
        p14: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      p11,
      p12,
      p13,
      p14,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(
        p1,
        p2,
        p3,
        p4,
        p5,
        p6,
        p7,
        p8,
        p9,
        p10,
        p11,
        p12,
        p13,
        p14,
        nc,
        changed or 0b1,
        changed1
      )
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    p14: Any?,
    p15: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(15) else sameBits(15)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        p11: Any?,
        p12: Any?,
        p13: Any?,
        p14: Any?,
        p15: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      p11,
      p12,
      p13,
      p14,
      p15,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(
        p1,
        p2,
        p3,
        p4,
        p5,
        p6,
        p7,
        p8,
        p9,
        p10,
        p11,
        p12,
        p13,
        p14,
        p15,
        nc,
        changed or 0b1,
        changed1
      )
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    p14: Any?,
    p15: Any?,
    p16: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(16) else sameBits(16)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        p11: Any?,
        p12: Any?,
        p13: Any?,
        p14: Any?,
        p15: Any?,
        p16: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      p11,
      p12,
      p13,
      p14,
      p15,
      p16,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(
        p1,
        p2,
        p3,
        p4,
        p5,
        p6,
        p7,
        p8,
        p9,
        p10,
        p11,
        p12,
        p13,
        p14,
        p15,
        p16,
        nc,
        changed or 0b1,
        changed1
      )
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    p14: Any?,
    p15: Any?,
    p16: Any?,
    p17: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(17) else sameBits(17)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        p11: Any?,
        p12: Any?,
        p13: Any?,
        p14: Any?,
        p15: Any?,
        p16: Any?,
        p17: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      p11,
      p12,
      p13,
      p14,
      p15,
      p16,
      p17,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(
        p1,
        p2,
        p3,
        p4,
        p5,
        p6,
        p7,
        p8,
        p9,
        p10,
        p11,
        p12,
        p13,
        p14,
        p15,
        p16,
        p17,
        nc,
        changed or 0b1,
        changed1
      )
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    p14: Any?,
    p15: Any?,
    p16: Any?,
    p17: Any?,
    p18: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(18) else sameBits(18)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        p11: Any?,
        p12: Any?,
        p13: Any?,
        p14: Any?,
        p15: Any?,
        p16: Any?,
        p17: Any?,
        p18: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      p11,
      p12,
      p13,
      p14,
      p15,
      p16,
      p17,
      p18,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(
        p1,
        p2,
        p3,
        p4,
        p5,
        p6,
        p7,
        p8,
        p9,
        p10,
        p11,
        p12,
        p13,
        p14,
        p15,
        p16,
        p17,
        p18,
        nc,
        changed or 0b1,
        changed1
      )
    }
    return result
  }
}

@ComposeCompilerApi
@Stable
actual interface ComposableLambda {
  actual operator fun invoke(c: Composer, changed: Int): Any?

  actual operator fun invoke(p1: Any?, c: Composer, changed: Int): Any?

  actual operator fun invoke(p1: Any?, p2: Any?, c: Composer, changed: Int): Any?

  actual operator fun invoke(p1: Any?, p2: Any?, p3: Any?, c: Composer, changed: Int): Any?

  actual operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    c: Composer,
    changed: Int
  ): Any?

  actual operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    c: Composer,
    changed: Int
  ): Any?

  actual operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    c: Composer,
    changed: Int
  ): Any?

  actual operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    c: Composer,
    changed: Int
  ): Any?

  actual operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    c: Composer,
    changed: Int
  ): Any?

  actual operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    c: Composer,
    changed: Int
  ): Any?

  actual operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any?

  actual operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any?

  actual operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any?

  actual operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any?

  actual operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    p14: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any?

  actual operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    p14: Any?,
    p15: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any?

  actual operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    p14: Any?,
    p15: Any?,
    p16: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any?

  actual operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    p14: Any?,
    p15: Any?,
    p16: Any?,
    p17: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any?

  actual operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    p14: Any?,
    p15: Any?,
    p16: Any?,
    p17: Any?,
    p18: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any?
}

/*
 * Copyright 2020 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

@file:OptIn(InternalComposeApi::class)
package androidx.compose.runtime.internal

import androidx.compose.runtime.ComposeCompilerApi
import androidx.compose.runtime.Composer
import androidx.compose.runtime.InternalComposeApi
import androidx.compose.runtime.RecomposeScope
import androidx.compose.runtime.Stable

/**
 * A Restart is created to hold composable lambdas to track when they are invoked allowing
 * the invocations to be invalidated when a new composable lambda is created during composition.
 *
 * This allows much of the call-graph to be skipped when a composable function is passed through
 * multiple levels of composable functions.
 */
@Suppress("NAME_SHADOWING")
@Stable
@OptIn(ComposeCompilerApi::class)
/* ktlint-disable parameter-list-wrapping */ // TODO(https://github.com/pinterest/ktlint/issues/921): reenable
internal actual class ComposableLambdaImpl actual constructor(
  val key: Int,
  private val tracked: Boolean
) : ComposableLambda {
  private var _block: Any? = null
  private var scope: RecomposeScope? = null
  private var scopes: MutableList<RecomposeScope>? = null

  private fun trackWrite() {
    if (tracked) {
      val scope = this.scope
      if (scope != null) {
        scope.invalidate()
        this.scope = null
      }
      val scopes = this.scopes
      if (scopes != null) {
        for (index in 0 until scopes.size) {
          val item = scopes[index]
          item.invalidate()
        }
        scopes.clear()
      }
    }
  }

  private fun trackRead(composer: Composer) {
    if (tracked) {
      val scope = composer.recomposeScope
      if (scope != null) {
        // Find the first invalid scope and replace it or record it if no scopes are invalid
        composer.recordUsed(scope)
        val lastScope = this.scope
        if (lastScope.replacableWith(scope)) {
          this.scope = scope
        } else {
          val lastScopes = scopes
          if (lastScopes == null) {
            val newScopes = mutableListOf<RecomposeScope>()
            scopes = newScopes
            newScopes.add(scope)
          } else {
            for (index in 0 until lastScopes.size) {
              val scopeAtIndex = lastScopes[index]
              if (scopeAtIndex.replacableWith(scope)) {
                lastScopes[index] = scope
                return
              }
            }
            lastScopes.add(scope)
          }
        }
      }
    }
  }

  actual fun update(block: Any) {
    if (_block != block) {
      val oldBlockNull = _block == null
      _block = block
      if (!oldBlockNull) {
        trackWrite()
      }
    }
  }

  override operator fun invoke(c: Composer, changed: Int): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(0) else sameBits(0)
    val result = (_block as (c: Composer, changed: Int) -> Any?)(c, dirty)
    c.endRestartGroup()?.updateScope(this as (Composer, Int) -> Unit)
    return result
  }

  override operator fun invoke(p1: Any?, c: Composer, changed: Int): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(1) else sameBits(1)
    val result = (
      _block as (
        p1: Any?,
        c: Composer,
        changed: Int
      ) -> Any?
      )(
      p1,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ -> this(p1, nc, changed or 0b1) }
    return result
  }

  override operator fun invoke(p1: Any?, p2: Any?, c: Composer, changed: Int): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(2) else sameBits(2)
    val result = (_block as (p1: Any?, p2: Any?, c: Composer, changed: Int) -> Any?)(
      p1,
      p2,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ -> this(p1, p2, nc, changed or 0b1) }
    return result
  }

  override operator fun invoke(p1: Any?, p2: Any?, p3: Any?, c: Composer, changed: Int): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(3) else sameBits(3)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        c: Composer,
        changed: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ -> this(p1, p2, p3, nc, changed or 0b1) }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    c: Composer,
    changed: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(4) else sameBits(4)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        c: Composer,
        changed: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, nc, changed or 0b1)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    c: Composer,
    changed: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(5) else sameBits(5)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        c: Composer,
        changed: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, p5, nc, changed or 0b1)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    c: Composer,
    changed: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(6) else sameBits(6)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        c: Composer,
        changed: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, p5, p6, nc, changed or 0b1)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    c: Composer,
    changed: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(7) else sameBits(7)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        c: Composer,
        changed: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, p5, p6, p7, nc, changed or 0b1)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    c: Composer,
    changed: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(8) else sameBits(8)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        c: Composer,
        changed: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, p5, p6, p7, p8, nc, changed or 0b1)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    c: Composer,
    changed: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed or if (c.changed(this)) differentBits(9) else sameBits(9)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        c: Composer,
        changed: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      c,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, p5, p6, p7, p8, p9, nc, changed or 0b1)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(10) else sameBits(10)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, p5, p6, p7, p8, p9, p10, nc, changed or 0b1, changed)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(11) else sameBits(11)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        p11: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      p11,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, p5, p6, p7, p8, p9, p10, p11, nc, changed or 0b1, changed1)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(12) else sameBits(12)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        p11: Any?,
        p12: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      p11,
      p12,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(p1, p2, p3, p4, p5, p6, p7, p8, p9, p10, p11, p12, nc, changed or 0b1, changed1)
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(13) else sameBits(13)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        p11: Any?,
        p12: Any?,
        p13: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      p11,
      p12,
      p13,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(
        p1,
        p2,
        p3,
        p4,
        p5,
        p6,
        p7,
        p8,
        p9,
        p10,
        p11,
        p12,
        p13,
        nc,
        changed or 0b1,
        changed1
      )
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    p14: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(14) else sameBits(14)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        p11: Any?,
        p12: Any?,
        p13: Any?,
        p14: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      p11,
      p12,
      p13,
      p14,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(
        p1,
        p2,
        p3,
        p4,
        p5,
        p6,
        p7,
        p8,
        p9,
        p10,
        p11,
        p12,
        p13,
        p14,
        nc,
        changed or 0b1,
        changed1
      )
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    p14: Any?,
    p15: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(15) else sameBits(15)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        p11: Any?,
        p12: Any?,
        p13: Any?,
        p14: Any?,
        p15: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      p11,
      p12,
      p13,
      p14,
      p15,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(
        p1,
        p2,
        p3,
        p4,
        p5,
        p6,
        p7,
        p8,
        p9,
        p10,
        p11,
        p12,
        p13,
        p14,
        p15,
        nc,
        changed or 0b1,
        changed1
      )
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    p14: Any?,
    p15: Any?,
    p16: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(16) else sameBits(16)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        p11: Any?,
        p12: Any?,
        p13: Any?,
        p14: Any?,
        p15: Any?,
        p16: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      p11,
      p12,
      p13,
      p14,
      p15,
      p16,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(
        p1,
        p2,
        p3,
        p4,
        p5,
        p6,
        p7,
        p8,
        p9,
        p10,
        p11,
        p12,
        p13,
        p14,
        p15,
        p16,
        nc,
        changed or 0b1,
        changed1
      )
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    p14: Any?,
    p15: Any?,
    p16: Any?,
    p17: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(17) else sameBits(17)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        p11: Any?,
        p12: Any?,
        p13: Any?,
        p14: Any?,
        p15: Any?,
        p16: Any?,
        p17: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      p11,
      p12,
      p13,
      p14,
      p15,
      p16,
      p17,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(
        p1,
        p2,
        p3,
        p4,
        p5,
        p6,
        p7,
        p8,
        p9,
        p10,
        p11,
        p12,
        p13,
        p14,
        p15,
        p16,
        p17,
        nc,
        changed or 0b1,
        changed1
      )
    }
    return result
  }

  override operator fun invoke(
    p1: Any?,
    p2: Any?,
    p3: Any?,
    p4: Any?,
    p5: Any?,
    p6: Any?,
    p7: Any?,
    p8: Any?,
    p9: Any?,
    p10: Any?,
    p11: Any?,
    p12: Any?,
    p13: Any?,
    p14: Any?,
    p15: Any?,
    p16: Any?,
    p17: Any?,
    p18: Any?,
    c: Composer,
    changed: Int,
    changed1: Int
  ): Any? {
    val c = c.startRestartGroup(key)
    trackRead(c)
    val dirty = changed1 or if (c.changed(this)) differentBits(18) else sameBits(18)
    val result = (
      _block as (
        p1: Any?,
        p2: Any?,
        p3: Any?,
        p4: Any?,
        p5: Any?,
        p6: Any?,
        p7: Any?,
        p8: Any?,
        p9: Any?,
        p10: Any?,
        p11: Any?,
        p12: Any?,
        p13: Any?,
        p14: Any?,
        p15: Any?,
        p16: Any?,
        p17: Any?,
        p18: Any?,
        c: Composer,
        changed: Int,
        changed1: Int
      ) -> Any?
      )(
      p1,
      p2,
      p3,
      p4,
      p5,
      p6,
      p7,
      p8,
      p9,
      p10,
      p11,
      p12,
      p13,
      p14,
      p15,
      p16,
      p17,
      p18,
      c,
      changed,
      dirty
    )
    c.endRestartGroup()?.updateScope { nc, _ ->
      this(
        p1,
        p2,
        p3,
        p4,
        p5,
        p6,
        p7,
        p8,
        p9,
        p10,
        p11,
        p12,
        p13,
        p14,
        p15,
        p16,
        p17,
        p18,
        nc,
        changed or 0b1,
        changed1
      )
    }
    return result
  }
}

@ComposeCompilerApi
@Stable
actual interface ComposableLambda :
  Function2<Composer, Int, Any?>,
  Function3<Any?, Composer, Int, Any?>,
  Function4<Any?, Any?, Composer, Int, Any?>,
  Function5<Any?, Any?, Any?, Composer, Int, Any?>,
  Function6<Any?, Any?, Any?, Any?, Composer, Int, Any?>,
  Function7<Any?, Any?, Any?, Any?, Any?, Composer, Int, Any?>,
  Function8<Any?, Any?, Any?, Any?, Any?, Any?, Composer, Int, Any?>,
  Function9<Any?, Any?, Any?, Any?, Any?, Any?, Any?, Composer, Int, Any?>,
  Function10<Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Composer, Int, Any?>,
  Function11<Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Composer, Int, Any?>,
  Function13<Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Composer, Int, Int,
    Any?>,
  Function14<Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Composer, Int, Int,
    Any?>,
  Function15<Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Composer,
    Int, Int, Any?>,
  Function16<Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?,
    Composer, Int, Int, Any?>,
  Function17<Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?,
    Composer, Int,
    Int, Any?>,
  Function18<Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?,
    Any?, Composer, Int, Int, Any?>,
  Function19<Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?,
    Any?, Any?, Composer, Int, Int, Any?>,
  Function20<Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?,
    Any?, Any?, Any?, Composer, Int, Int, Any?>,
  Function21<Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?, Any?,
    Any?, Any?, Any?, Any?, Composer, Int, Int, Any?>

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.schema.generator

import app.cash.exhaustive.Exhaustive
import app.cash.treehouse.schema.parser.Children
import app.cash.treehouse.schema.parser.Event
import app.cash.treehouse.schema.parser.Property
import app.cash.treehouse.schema.parser.Schema
import app.cash.treehouse.schema.parser.Widget
import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.FileSpec
import com.squareup.kotlinpoet.FunSpec
import com.squareup.kotlinpoet.KModifier.OVERRIDE
import com.squareup.kotlinpoet.KModifier.PRIVATE
import com.squareup.kotlinpoet.KModifier.PUBLIC
import com.squareup.kotlinpoet.ParameterSpec
import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
import com.squareup.kotlinpoet.PropertySpec
import com.squareup.kotlinpoet.TypeSpec
import com.squareup.kotlinpoet.asTypeName
import com.squareup.kotlinpoet.joinToCode
import com.squareup.kotlinpoet.jvm.jvmField

/*
@Composable
fun Button(
  text: String,
  enabled: Boolean = true,
  onClick: (() -> Unit)? = null,
) {
  ComposeNode<ButtonComposeNode, ProtocolApplier>(
    factory = ::ButtonNode,
    update = {
      set(text) {
        appendDiff(PropertyDiff(id, 1, text))
      }
      set(enabled) {
        appendDiff(PropertyDiff(id, 2, enabled))
      }
      set(onClick) {
        val onClickSet = onClick != null
        if (onClickSet != (this.onClick != null)) {
          appendDiff(PropertyDiff(id, 3, onClickSet))
        }
        this.onClick = onClick
      }
    }
  )
}

private class ButtonComposeNode(id: Long) : Node(id, 2) {
  var onClick: (() -> Unit)? = null

  override fun sendEvent(event: Event) {
    when (event.eventId) {
      1L -> onClick?.invoke()
      else -> throw IllegalArgumentException("Unknown event ID ${event.eventId}")
    }
  }
}
*/
internal fun generateComposeNode(schema: Schema, widget: Widget): FileSpec {
  val events = widget.traits.filterIsInstance<Event>()
  val nodeType = if (events.isEmpty()) {
    protocolNode
  } else {
    schema.composeNodeType(widget)
  }
  val applierOfServerNode = applier.parameterizedBy(protocolNode)
  return FileSpec.builder(schema.composePackage, widget.flatName)
    .addFunction(
      FunSpec.builder(widget.flatName)
        .addModifiers(PUBLIC)
        .addAnnotation(composable)
        .apply {
          for (trait in widget.traits) {
            addParameter(
              when (trait) {
                is Property -> {
                  ParameterSpec.builder(trait.name, trait.type.asTypeName())
                    .apply {
                      trait.defaultExpression?.let { defaultValue(it) }
                    }
                    .build()
                }
                is Event -> {
                  ParameterSpec.builder(trait.name, trait.lambdaType)
                    .defaultValue("null")
                    .build()
                }
                is Children -> {
                  ParameterSpec.builder(trait.name, composableLambda)
                    .build()
                }
              }
            )
          }

          val arguments = mutableListOf<CodeBlock>()

          arguments += CodeBlock.builder()
            .add("factory = ")
            .apply {
              if (events.isEmpty()) {
                add("{\n⇥%T(%L)\n⇤}", nodeType, widget.tag)
              } else {
                add("::%T", nodeType)
              }
            }
            .build()

          val updateLambda = CodeBlock.builder()
          val childrenLambda = CodeBlock.builder()
          for (trait in widget.traits) {
            @Exhaustive when (trait) {
              is Property -> {
                updateLambda.apply {
                  add("set(%N) {\n", trait.name)
                  indent()
                  add("appendDiff(%T(this.id, %L, %N))\n", propertyDiff, trait.tag, trait.name)
                  unindent()
                  add("}\n")
                }
              }
              is Event -> {
                updateLambda.apply {
                  add("set(%N) {\n", trait.name)
                  indent()
                  add("val %1NSet = %1N != null\n", trait.name)
                  add("if (%1NSet != (this.%1N != null)) {\n", trait.name)
                  indent()
                  add("appendDiff(%T(this.id, %L, %NSet))\n", propertyDiff, trait.tag, trait.name)
                  unindent()
                  add("}\n")
                  add("this.%1N = %1N\n", trait.name)
                  unindent()
                  add("}\n")
                }
              }
              is Children -> {
                childrenLambda.apply {
                  add("%M(%L) {\n", syntheticChildren, trait.tag)
                  indent()
                  add("%N()\n", trait.name)
                  unindent()
                  add("}\n")
                }
              }
            }
          }

          arguments += CodeBlock.builder()
            .add("update = {\n")
            .indent()
            .add(updateLambda.build())
            .unindent()
            .add("}")
            .build()

          if (childrenLambda.isNotEmpty()) {
            arguments += CodeBlock.builder()
              .add("content = {\n")
              .indent()
              .add(childrenLambda.build())
              .unindent()
              .add("}")
              .build()
          }

          addStatement(
            "%M<%T, %T>(%L)", composeNode, nodeType, applierOfServerNode,
            arguments.joinToCode(",\n", "\n", ",\n")
          )
        }
        .build()
    )
    .apply {
      if (events.isNotEmpty()) {
        addType(
          TypeSpec.classBuilder(nodeType)
            .addModifiers(PRIVATE)
            .superclass(protocolNode)
            .addSuperclassConstructorParameter("%L", widget.tag)
            .apply {
              for (event in events) {
                addProperty(
                  PropertySpec.builder(event.name, event.lambdaType)
                    .mutable(true)
                    .initializer("null")
                    .jvmField() // Method count optimization as this is implementation detail.
                    .build()
                )
              }
            }
            .addFunction(
              FunSpec.builder("sendEvent")
                .addModifiers(PUBLIC, OVERRIDE)
                .addParameter("event", eventType)
                .beginControlFlow("when (val tag = event.tag)")
                .apply {
                  for (event in events) {
                    if (event.parameterType != null) {
                      addStatement(
                        "%L -> %N?.invoke(event.value as %T)", event.tag, event.name,
                        event.parameterType!!.asTypeName()
                      )
                    } else {
                      addStatement("%L -> %N?.invoke()", event.tag, event.name)
                    }
                  }
                }
                .addStatement("else -> throw %T(\"Unknown tag \$tag\")", iae)
                .endControlFlow()
                .build()
            )
            .build()
        )
      }
    }
    .build()
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
@file:JvmName("Main")

package app.cash.treehouse.schema.generator

import app.cash.exhaustive.Exhaustive
import app.cash.treehouse.schema.generator.TreehouseGenerator.Type.Compose
import app.cash.treehouse.schema.generator.TreehouseGenerator.Type.Test
import app.cash.treehouse.schema.generator.TreehouseGenerator.Type.Widget
import app.cash.treehouse.schema.parser.parseSchema
import com.github.ajalt.clikt.core.CliktCommand
import com.github.ajalt.clikt.core.CliktError
import com.github.ajalt.clikt.parameters.arguments.argument
import com.github.ajalt.clikt.parameters.arguments.convert
import com.github.ajalt.clikt.parameters.arguments.help
import com.github.ajalt.clikt.parameters.options.help
import com.github.ajalt.clikt.parameters.options.option
import com.github.ajalt.clikt.parameters.options.required
import com.github.ajalt.clikt.parameters.options.switch
import com.github.ajalt.clikt.parameters.types.path

public fun main(vararg args: String) {
  TreehouseGenerator().main(args)
}

private class TreehouseGenerator : CliktCommand() {
  enum class Type {
    Compose,
    Test,
    Widget,
  }

  private val type by option()
    .switch(
      "--compose" to Compose,
      "--test" to Test,
      "--widget" to Widget,
    )
    .help("Type of code to generate")
    .required()

  private val out by option().path().required()
    .help("Directory into which generated files are written")

  private val schemaType by argument("schema")
    .help("Fully-qualified class name for the @Schema-annotated interface")
    .convert {
      // Replace with https://youtrack.jetbrains.com/issue/KT-10440 once it ships.
      try {
        Class.forName(it).kotlin
      } catch (e: ClassNotFoundException) {
        throw CliktError("Unable to load class: $it", e)
      }
    }

  override fun run() {
    val schema = parseSchema(schemaType)
    @Exhaustive when (type) {
      Widget -> {
        generateWidgetFactory(schema).writeTo(out)
        for (widget in schema.widgets) {
          generateWidget(schema, widget).writeTo(out)
        }
      }
      Compose -> {
        for (widget in schema.widgets) {
          generateComposeNode(schema, widget).writeTo(out)
        }
      }
      Test -> {
        generateSchemaWidgetFactory(schema).writeTo(out)
        for (widget in schema.widgets) {
          generateSchemaWidget(schema, widget).writeTo(out)
        }
      }
    }
  }
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.schema.generator

import app.cash.treehouse.schema.parser.Event
import app.cash.treehouse.schema.parser.Schema
import app.cash.treehouse.schema.parser.Widget
import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.LambdaTypeName
import com.squareup.kotlinpoet.TypeName
import com.squareup.kotlinpoet.UNIT
import com.squareup.kotlinpoet.asClassName
import com.squareup.kotlinpoet.asTypeName

/**
 * Returns a single string that is likely to be unique within a schema, like `MapEntry` or
 * `NavigationBarButton`.
 */
internal val Widget.flatName: String
  get() = type.asClassName().simpleNames.joinToString(separator = "")

internal val Event.lambdaType: TypeName
  get() {
    parameterType?.let { parameterType ->
      return LambdaTypeName.get(null, parameterType.asTypeName(), returnType = UNIT)
        .copy(nullable = true)
    }
    return noArgumentEventLambda
  }

private val noArgumentEventLambda = LambdaTypeName.get(returnType = UNIT).copy(nullable = true)

internal fun Schema.composeNodeType(widget: Widget): ClassName {
  return ClassName(composePackage, widget.flatName + "ComposeNode")
}

internal val Schema.composePackage get() = "$`package`.compose"

internal fun Schema.widgetType(widget: Widget): ClassName {
  return ClassName(displayPackage, widget.flatName)
}

internal fun Schema.getWidgetFactoryType(): ClassName {
  return ClassName(displayPackage, "${name}WidgetFactory")
}

internal val Schema.displayPackage get() = "$`package`.widget"

internal fun Schema.testType(widget: Widget): ClassName {
  return ClassName(testPackage, "Schema${widget.flatName}")
}

internal val Schema.testPackage get() = "$`package`.test"

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.schema.generator

import app.cash.exhaustive.Exhaustive
import app.cash.treehouse.schema.parser.Children
import app.cash.treehouse.schema.parser.Event
import app.cash.treehouse.schema.parser.Property
import app.cash.treehouse.schema.parser.Schema
import app.cash.treehouse.schema.parser.Widget
import com.squareup.kotlinpoet.ANY
import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.FileSpec
import com.squareup.kotlinpoet.FunSpec
import com.squareup.kotlinpoet.KModifier.ABSTRACT
import com.squareup.kotlinpoet.KModifier.OVERRIDE
import com.squareup.kotlinpoet.KModifier.PRIVATE
import com.squareup.kotlinpoet.ParameterSpec
import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
import com.squareup.kotlinpoet.PropertySpec
import com.squareup.kotlinpoet.STRING
import com.squareup.kotlinpoet.TypeSpec
import com.squareup.kotlinpoet.asClassName
import com.squareup.kotlinpoet.asTypeName
import com.squareup.kotlinpoet.joinToCode

/*
class SchemaSunspotWidgetFactory : SunspotWidgetFactory<Nothing> {
  override fun SunspotBox(): SunspotBox<Nothing> = SchemaSunspotBox(scope)
  // ...
}

interface SchemaWidget : Widget<SchemaWidget> {
  val schema: Any
  fun assertSchema(expected: Any, path: String = "$")
}
*/
internal fun generateSchemaWidgetFactory(schema: Schema): FileSpec {
  val widgetFactoryType = ClassName(schema.testPackage, "Schema${schema.name}WidgetFactory")
  val schemaWidgetType = ClassName(schema.testPackage, "SchemaWidget")

  return FileSpec.builder(schema.testPackage, widgetFactoryType.simpleName)
    .addType(
      TypeSpec.objectBuilder(widgetFactoryType)
        .addSuperinterface(schema.getWidgetFactoryType().parameterizedBy(schemaWidgetType))
        .apply {
          for (widget in schema.widgets) {
            addFunction(
              FunSpec.builder(widget.flatName)
                .addModifiers(OVERRIDE)
                .returns(schema.widgetType(widget).parameterizedBy(schemaWidgetType))
                .addStatement("return %T()", schema.testType(widget))
                .build()
            )
          }
        }
        .build()
    )
    .addType(
      TypeSpec.interfaceBuilder(schemaWidgetType.simpleName)
        .addSuperinterface(widget.parameterizedBy(schemaWidgetType))
        .addProperty("schema", ANY, ABSTRACT)
        .addFunction(
          FunSpec.builder("assertSchema")
            .addModifiers(ABSTRACT)
            .addParameter("expected", ANY)
            .addParameter(
              ParameterSpec.builder("path", STRING)
                .defaultValue(""""$"""")
                .build()
            )
            .build()
        )
        .build()
    )
    .build()
}

/*
private class SchemaSunspotBox : SunspotBox<SchemaWidget>, SchemaWidget {
  // ...
}
 */
internal fun generateSchemaWidget(schema: Schema, widget: Widget): FileSpec {
  val className = schema.testType(widget)
  val widgetType = widget.type.asClassName()
  val schemaWidgetType = ClassName(schema.testPackage, "SchemaWidget")
  return FileSpec.builder(schema.testPackage, className.simpleName)
    .addType(
      TypeSpec.classBuilder(className)
        .addSuperinterface(schema.widgetType(widget).parameterizedBy(schemaWidgetType))
        .addSuperinterface(schemaWidgetType)
        .addProperty(
          PropertySpec.builder("value", schemaWidgetType, OVERRIDE)
            .getter(
              FunSpec.getterBuilder()
                .addStatement("return this")
                .build()
            )
            .build()
        )
        .apply {
          val schemaParameters = mutableListOf<CodeBlock>()
          for (trait in widget.traits) {
            @Exhaustive when (trait) {
              is Property -> {
                addProperty(
                  PropertySpec.builder(
                    trait.name, trait.type.asTypeName().copy(nullable = true),
                    PRIVATE
                  )
                    .mutable(true)
                    .initializer("null")
                    .build()
                )
                addFunction(
                  FunSpec.builder(trait.name)
                    .addModifiers(OVERRIDE)
                    .addParameter(trait.name, trait.type.asTypeName())
                    .addStatement("this.%1N = %1N", trait.name)
                    .build()
                )
                schemaParameters += if (trait.type.isMarkedNullable) {
                  CodeBlock.of("%1N = %1N", trait.name)
                } else {
                  CodeBlock.of("%1N = %1N ?: error(%2S)", trait.name, "Required property '${trait.name}' not set")
                }
              }
              is Event -> {
                addProperty(
                  PropertySpec.builder(trait.name, trait.lambdaType, PRIVATE)
                    .mutable(true)
                    .initializer("null")
                    .build()
                )
                addFunction(
                  FunSpec.builder(trait.name)
                    .addModifiers(OVERRIDE)
                    .addParameter(trait.name, trait.lambdaType)
                    .addStatement("this.%1N = %1N", trait.name)
                    .build()
                )
                schemaParameters += CodeBlock.of(
                  "%1N = %1N ?: { error(%2S) }", trait.name, "No '${trait.name}' lambda set"
                )
              }
              is Children -> {
                addProperty(
                  PropertySpec.builder(trait.name, mutableListChildren.parameterizedBy(schemaWidgetType))
                    .addModifiers(OVERRIDE)
                    .initializer("%T()", mutableListChildren)
                    .build()
                )
                schemaParameters += CodeBlock.of("%1N = %1N.map(%2T::schema)", trait.name, schemaWidgetType)
              }
            }
          }
          addProperty(
            PropertySpec.builder("schema", widgetType, OVERRIDE)
              .getter(
                FunSpec.getterBuilder()
                  .apply {
                    if (schemaParameters.isEmpty()) {
                      addStatement("return %T", widgetType)
                    } else {
                      addStatement(
                        "return %T(%L)", widgetType,
                        schemaParameters.joinToCode(separator = ",\n", prefix = "\n", suffix = "\n")
                      )
                    }
                  }
                  .build()
              )
              .build()
          )
        }
        .addFunction(
          FunSpec.builder("assertSchema")
            .addModifiers(OVERRIDE)
            .addParameter("expected", ANY)
            .addParameter("path", STRING)
            .beginControlFlow("if (expected !is %T)", widgetType)
            .addStatement(
              "throw %T(\"Expected: \${expected::class} but was: ${widgetType.simpleNames.joinToString(".")}\\n  at \$path\")",
              ae
            )
            .endControlFlow()
            .apply {
              for (trait in widget.traits) {
                @Exhaustive when (trait) {
                  is Property -> {
                    beginControlFlow("if (expected.%1N != %1N)", trait.name)
                    addStatement(
                      "throw %1T(\"Expected: \${expected.%2N} but was: \$%2N\\n  at \$path.%2N\")",
                      ae, trait.name
                    )
                    endControlFlow()
                  }
                  is Children -> {
                    addStatement("%N.forEachIndexed { index, child ->", trait.name)
                    addStatement(
                      "⇥child.assertSchema(expected.%1N[index], \"\$path.%1N[\$index]\")⇤",
                      trait.name
                    )
                    addStatement("}")
                  }
                  is Event -> {
                    // Ignored for schema assertion.
                  }
                }
              }
            }
            .build()
        )
        .build()
    )
    .build()
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.schema.generator

import com.squareup.kotlinpoet.AnnotationSpec
import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.LambdaTypeName
import com.squareup.kotlinpoet.MemberName
import com.squareup.kotlinpoet.UNIT

internal val eventType = ClassName("app.cash.treehouse.protocol", "Event")
internal val propertyDiff = ClassName("app.cash.treehouse.protocol", "PropertyDiff")

internal val eventSink = LambdaTypeName.get(receiver = null, eventType, returnType = UNIT)

internal val widget = ClassName("app.cash.treehouse.widget", "Widget")
internal val widgetChildren = widget.nestedClass("Children")
internal val widgetFactory = widget.nestedClass("Factory")
internal val mutableListChildren = ClassName("app.cash.treehouse.widget", "MutableListChildren")

internal val protocolApplier = ClassName("app.cash.treehouse.compose", "ProtocolApplier")
internal val protocolNode = ClassName("app.cash.treehouse.compose", "ProtocolNode")
internal val syntheticChildren = MemberName("app.cash.treehouse.compose", "\$SyntheticChildren")

internal val applier = ClassName("androidx.compose.runtime", "Applier")
internal val composable = ClassName("androidx.compose.runtime", "Composable")
internal val composeNode = MemberName("androidx.compose.runtime", "ComposeNode")

internal val composableLambda = LambdaTypeName.get(returnType = UNIT)
  .copy(
    annotations = listOf(
      AnnotationSpec.builder(composable).build(),
    )
  )

internal val ae = ClassName("kotlin", "AssertionError")
internal val iae = ClassName("kotlin", "IllegalArgumentException")

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.schema.generator

import app.cash.exhaustive.Exhaustive
import app.cash.treehouse.schema.parser.Children
import app.cash.treehouse.schema.parser.Event
import app.cash.treehouse.schema.parser.Property
import app.cash.treehouse.schema.parser.Schema
import app.cash.treehouse.schema.parser.Widget
import com.squareup.kotlinpoet.ANY
import com.squareup.kotlinpoet.BOOLEAN
import com.squareup.kotlinpoet.FileSpec
import com.squareup.kotlinpoet.FunSpec
import com.squareup.kotlinpoet.INT
import com.squareup.kotlinpoet.KModifier.ABSTRACT
import com.squareup.kotlinpoet.KModifier.OVERRIDE
import com.squareup.kotlinpoet.KModifier.PUBLIC
import com.squareup.kotlinpoet.LONG
import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
import com.squareup.kotlinpoet.PropertySpec
import com.squareup.kotlinpoet.TypeSpec
import com.squareup.kotlinpoet.TypeVariableName
import com.squareup.kotlinpoet.asTypeName

/*
interface SunspotWidgetFactory<T : Any> : Widget.Factory<T> {
  fun SunspotText(): SunspotText<T>
  fun SunspotButton(): SunspotButton<T>

  override fun create(, kind: Int, id: Long): TreeNode<T> {
    return when (kind) {
      1 -> SunspotText()
      2 -> SunspotButton()
      else -> throw IllegalArgumentException("Unknown kind $kind")
    }
  }
}
*/
internal fun generateWidgetFactory(schema: Schema): FileSpec {
  return FileSpec.builder(schema.displayPackage, schema.getWidgetFactoryType().simpleName)
    .addType(
      TypeSpec.interfaceBuilder(schema.getWidgetFactoryType())
        .addModifiers(PUBLIC)
        .addTypeVariable(typeVariableT)
        .addSuperinterface(factoryOfT)
        .apply {
          for (node in schema.widgets) {
            addFunction(
              FunSpec.builder(node.flatName)
                .addModifiers(PUBLIC, ABSTRACT)
                .returns(schema.widgetType(node).parameterizedBy(typeVariableT))
                .build()
            )
          }
        }
        .addFunction(
          FunSpec.builder("create")
            .addModifiers(PUBLIC, OVERRIDE)
            .addParameter("kind", INT)
            .addParameter("id", LONG)
            .returns(widgetOfT)
            .beginControlFlow("return when (kind)")
            .apply {
              for (node in schema.widgets.sortedBy { it.tag }) {
                addStatement("%L -> %N()", node.tag, node.flatName)
              }
            }
            .addStatement("else -> throw %T(\"Unknown kind \$kind\")", iae)
            .endControlFlow()
            .build()
        )
        .build()
    )
    .build()
}

/*
interface SunspotButton<out T: Any> : SunspotNode<T> {
  fun text(text: String?)
  fun enabled(enabled: Boolean)
  fun onClick(onClick: (() -> Unit)?)

  override fun apply(diff: PropertyDiff, events: (Event) -> Unit) {
    when (val tag = diff.tag) {
      1 -> text(diff.value as String?)
      2 -> enabled(diff.value as Boolean)
      3 -> {
        val onClick: (() -> Unit)? = if (diff.value as Boolean) {
          val event = Event(diff.id, 3, null);
          { events(event) }
        } else {
          null
        }
        onClick(onClick)
      }
      else -> throw IllegalArgumentException("Unknown tag $tag")
    }
  }
}
*/
internal fun generateWidget(schema: Schema, widget: Widget): FileSpec {
  return FileSpec.builder(schema.displayPackage, widget.flatName)
    .addType(
      TypeSpec.interfaceBuilder(widget.flatName)
        .addModifiers(PUBLIC)
        .addTypeVariable(typeVariableT)
        .addSuperinterface(widgetOfT)
        .apply {
          for (trait in widget.traits) {
            when (trait) {
              is Property -> {
                addFunction(
                  FunSpec.builder(trait.name)
                    .addModifiers(PUBLIC, ABSTRACT)
                    .addParameter(trait.name, trait.type.asTypeName())
                    .build()
                )
              }
              is Event -> {
                addFunction(
                  FunSpec.builder(trait.name)
                    .addModifiers(PUBLIC, ABSTRACT)
                    .addParameter(trait.name, trait.lambdaType)
                    .build()
                )
              }
              is Children -> {
                addProperty(
                  PropertySpec.builder(trait.name, childrenOfT)
                    .addModifiers(PUBLIC, ABSTRACT)
                    .build()
                )
              }
            }
          }
          val childrens = widget.traits.filterIsInstance<Children>()
          if (childrens.isNotEmpty()) {
            addFunction(
              FunSpec.builder("children")
                .addModifiers(PUBLIC, OVERRIDE)
                .addParameter("tag", INT)
                .returns(childrenOfT)
                .beginControlFlow("return when (tag)")
                .apply {
                  for (children in childrens) {
                    addStatement("%L -> %N", children.tag, children.name)
                  }
                }
                .addStatement("else -> throw %T(\"Unknown tag \$tag\")", iae)
                .endControlFlow()
                .build()
            )
          }
        }
        .addFunction(
          FunSpec.builder("apply")
            .addModifiers(PUBLIC, OVERRIDE)
            .addParameter("diff", propertyDiff)
            .addParameter("events", eventSink)
            .beginControlFlow("when (val tag = diff.tag)")
            .apply {
              for (trait in widget.traits) {
                @Exhaustive when (trait) {
                  is Property -> {
                    addStatement(
                      "%L -> %N(diff.value as %T)", trait.tag, trait.name,
                      trait.type.asTypeName()
                    )
                  }
                  is Event -> {
                    beginControlFlow("%L ->", trait.tag)
                    beginControlFlow("val %N: %T = if (diff.value as %T)", trait.name, trait.lambdaType, BOOLEAN)
                    addStatement(
                      "{ events(%T(diff.id, %L, %L)) }", eventType, trait.tag,
                      if (trait.parameterType != null) "it" else "null"
                    )
                    nextControlFlow("else")
                    addStatement("null")
                    endControlFlow()
                    addStatement("%1N(%1N)", trait.name)
                    endControlFlow()
                  }
                  is Children -> {}
                }
              }
            }
            .addStatement("else -> throw %T(\"Unknown tag \$tag\")", iae)
            .endControlFlow()
            .build()
        )
        .build()
    )
    .build()
}

private val typeVariableT = TypeVariableName("T", listOf(ANY))
private val widgetOfT = widget.parameterizedBy(typeVariableT)
private val childrenOfT = widgetChildren.parameterizedBy(typeVariableT)
private val factoryOfT = widgetFactory.parameterizedBy(typeVariableT)

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.schema.parser

import kotlin.reflect.KClass
import kotlin.reflect.KType

public data class Schema(
  val name: String,
  val `package`: String,
  val widgets: List<Widget>,
)

public data class Widget(
  val tag: Int,
  /** Either a 'data class' or 'object'. */
  val type: KClass<*>,
  /** Non-empty list for a 'data class' [type] or empty list for 'object' [type]. */
  val traits: List<Trait>,
)

public sealed class Trait {
  public abstract val tag: Int
  public abstract val name: String
  public abstract val defaultExpression: String?
}

public data class Property(
  override val tag: Int,
  override val name: String,
  val type: KType,
  override val defaultExpression: String?,
) : Trait()

public data class Event(
  override val tag: Int,
  override val name: String,
  override val defaultExpression: String?,
  val parameterType: KType?,
) : Trait()

public data class Children(
  override val tag: Int,
  override val name: String,
) : Trait() {
  override val defaultExpression: String? get() = null
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.schema.parser

import kotlin.reflect.KClass
import kotlin.reflect.KTypeProjection.Companion.invariant
import kotlin.reflect.full.createType
import kotlin.reflect.full.findAnnotation
import kotlin.reflect.full.isSubtypeOf
import kotlin.reflect.full.primaryConstructor
import kotlin.reflect.full.starProjectedType
import kotlin.reflect.full.withNullability
import app.cash.treehouse.schema.Children as ChildrenAnnotation
import app.cash.treehouse.schema.Default as DefaultAnnotation
import app.cash.treehouse.schema.Property as PropertyAnnotation
import app.cash.treehouse.schema.Schema as SchemaAnnotation
import app.cash.treehouse.schema.Widget as WidgetAnnotation

private val childrenType = List::class.createType(
  arguments = listOf(invariant(Any::class.createType()))
)
private val eventType = Function::class.starProjectedType
private val optionalEventType = eventType.withNullability(true)

public fun parseSchema(schemaType: KClass<*>): Schema {
  val widgets = mutableListOf<Widget>()

  val widgetTypes = requireNotNull(schemaType.findAnnotation<SchemaAnnotation>()) {
    "Schema ${schemaType.qualifiedName} missing @Schema annotation"
  }.widgets

  val duplicatedWidgets = widgetTypes.groupBy { it }.filterValues { it.size > 1 }.keys
  if (duplicatedWidgets.isNotEmpty()) {
    throw IllegalArgumentException(
      buildString {
        append("Schema contains repeated widget")
        if (duplicatedWidgets.size > 1) {
          append('s')
        }
        duplicatedWidgets.joinTo(this, prefix = "\n\n- ", separator = "\n- ") { it.qualifiedName!! }
      }
    )
  }

  for (widgetType in widgetTypes) {
    val widgetAnnotation = requireNotNull(widgetType.findAnnotation<WidgetAnnotation>()) {
      "${widgetType.qualifiedName} missing @Widget annotation"
    }

    val traits = if (widgetType.isData) {
      widgetType.primaryConstructor!!.parameters.map {
        val property = it.findAnnotation<PropertyAnnotation>()
        val children = it.findAnnotation<ChildrenAnnotation>()
        val defaultExpression = it.findAnnotation<DefaultAnnotation>()?.expression

        if (property != null) {
          if (it.type.isSubtypeOf(eventType) || it.type.isSubtypeOf(optionalEventType)) {
            val arguments = it.type.arguments.dropLast(1) // Drop return type.
            require(arguments.size <= 1) {
              "@Property ${widgetType.qualifiedName}#${it.name} lambda type can only have zero or one arguments. Found: $arguments"
            }
            Event(property.tag, it.name!!, defaultExpression, arguments.singleOrNull()?.type)
          } else {
            Property(property.tag, it.name!!, it.type, defaultExpression)
          }
        } else if (children != null) {
          require(it.type == childrenType) {
            "@Children ${widgetType.qualifiedName}#${it.name} must be of type 'List<Any>'"
          }
          Children(children.tag, it.name!!)
        } else {
          throw IllegalArgumentException("Unannotated parameter \"${it.name}\" on ${widgetType.qualifiedName}")
        }
      }
    } else if (widgetType.objectInstance != null) {
      emptyList()
    } else {
      throw IllegalArgumentException(
        "@Widget ${widgetType.qualifiedName} must be 'data' class or 'object'"
      )
    }

    val badChildren =
      traits.filterIsInstance<Children>().groupBy(Children::tag).filterValues { it.size > 1 }
    if (badChildren.isNotEmpty()) {
      throw IllegalArgumentException(
        buildString {
          appendLine("Widget ${widgetType.qualifiedName}'s @Children tags must be unique")
          for ((tag, group) in badChildren) {
            append("\n- @Children($tag): ")
            group.joinTo(this) { it.name }
          }
        }
      )
    }

    val badProperties =
      traits.filterIsInstance<Property>().groupBy(Property::tag).filterValues { it.size > 1 }
    if (badProperties.isNotEmpty()) {
      throw IllegalArgumentException(
        buildString {
          appendLine("Widget ${widgetType.qualifiedName}'s @Property tags must be unique")
          for ((tag, group) in badProperties) {
            append("\n- @Property($tag): ")
            group.joinTo(this) { it.name }
          }
        }
      )
    }

    widgets += Widget(widgetAnnotation.tag, widgetType, traits)
  }

  val badWidgets = widgets.groupBy(Widget::tag).filterValues { it.size > 1 }
  if (badWidgets.isNotEmpty()) {
    throw IllegalArgumentException(
      buildString {
        appendLine("Schema @Widget tags must be unique")
        for ((tag, group) in badWidgets) {
          append("\n- @Widget($tag): ")
          group.joinTo(this) { it.type.qualifiedName!! }
        }
      }
    )
  }

  return Schema(schemaType.simpleName!!, schemaType.packageName, widgets)
}

/*
 * Copyright (C) 2021 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package app.cash.treehouse.schema.parser

import kotlin.reflect.KClass

internal val KClass<*>.packageName: String
  get() {
    // Replace with https://youtrack.jetbrains.com/issue/KT-18104 once it ships.
    // Note: Class.packageName isn't available until Java 9.
    val javaClass = java
    require(!javaClass.isPrimitive && !javaClass.isArray)
    return javaClass.name.substringBeforeLast(".", missingDelimiterValue = "")
  }