package soy.iko.opencode.data.model

import androidx.compose.runtime.Immutable
import kotlinx.serialization.SerialName
import kotlinx.serialization.Serializable
import kotlinx.serialization.json.JsonElement
import kotlinx.serialization.json.JsonClassDiscriminator

/**
 * The state machine of a tool call, discriminated on `status`.
 * pending -> running -> completed | error
 */
@OptIn(kotlinx.serialization.ExperimentalSerializationApi::class)
@Serializable
@JsonClassDiscriminator("status")
sealed interface ToolState

@Serializable
@SerialName("pending")
data object ToolPending : ToolState

@Immutable
@Serializable
@SerialName("running")
data class ToolRunning(
    val input: JsonElement? = null,
    val title: String? = null,
    val time: PartTimeInfo? = null,
) : ToolState

@Immutable
@Serializable
@SerialName("completed")
data class ToolCompleted(
    val input: JsonElement? = null,
    val output: String? = null,
    val title: String? = null,
    val time: PartTimeInfo? = null,
) : ToolState

@Immutable
@Serializable
@SerialName("error")
data class ToolError(
    val input: JsonElement? = null,
    val error: String? = null,
    val time: PartTimeInfo? = null,
) : ToolState

/** Fallback for an unrecognized `status`. */
@Serializable
@SerialName("__unknown")
data object ToolUnknown : ToolState

/** The raw input arguments carried by a tool call, present on the running/completed/error
 *  states (pending/unknown carry none). */
fun ToolState.inputElement(): JsonElement? = when (this) {
    is ToolRunning -> input
    is ToolCompleted -> input
    is ToolError -> input
    else -> null
}