package net.kentonc.si.agent

import kotlinx.serialization.SerialName
import kotlinx.serialization.Serializable
import kotlinx.serialization.json.JsonElement
import kotlinx.serialization.json.JsonPrimitive
import kotlinx.serialization.json.JsonArray
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.JsonNull
import kotlinx.serialization.json.add
import kotlinx.serialization.json.buildJsonArray
import kotlinx.serialization.json.buildJsonObject
import kotlinx.serialization.json.put

/**
 * OpenAI-compatible wire types, mirroring src/llm.rs in the Rust original:
 * messages carry optional content (plain text or a parts list) plus tool
 * calls, and tool results come back as separate `tool` role messages.
 */

@Serializable
data class WireMessage(
    val role: String,
    val content: JsonElement? = null,
    @SerialName("tool_calls") val toolCalls: List<ToolCall> = emptyList(),
    @SerialName("tool_call_id") val toolCallId: String? = null,
) {
    companion object {
        fun system(text: String) = WireMessage(role = "system", content = JsonPrimitive(text))
        fun user(text: String, images: List<AgentImage>): WireMessage {
            val content: JsonElement = if (images.isEmpty()) JsonPrimitive(text)
            else buildJsonArray {
                if (text.isNotEmpty()) add(buildJsonObject {
                    put("type", "text"); put("text", text)
                })
                images.forEach { img -> add(img.toPart()) }
            }
            return WireMessage(role = "user", content = content)
        }
        fun assistant(content: String?, toolCalls: List<ToolCall>) = WireMessage(
            role = "assistant",
            content = content?.takeIf { it.isNotEmpty() }?.let { JsonPrimitive(it) },
            toolCalls = toolCalls,
        )
        fun toolResult(id: String, text: String) = WireMessage(
            role = "tool",
            content = JsonPrimitive(text),
            toolCallId = id,
        )
    }
}

@Serializable
data class ToolCall(
    val id: String = "",
    @SerialName("type") val type: String = "function",
    val function: FunctionCall,
)

@Serializable
data class FunctionCall(
    val name: String,
    val arguments: String = "",
)

/** One-line rendering of a call, for showing the user what's happening (ToolCall::summary). */
fun ToolCall.summary(): String {
    val args: JsonObject = runCatching {
        json.parseToJsonElement(function.arguments.ifBlank { "{}" })
    }.getOrNull() as? JsonObject ?: JsonObject(emptyMap())
    val parts = args.entries.map { (key, value) ->
        val text = when (value) {
            is JsonPrimitive -> if (key == "command" || key == "content") value.content else truncateFor(key, value.content)
            else -> truncateFor(key, value.toString())
        }
        "$key=$text"
    }
    return "${function.name}(${parts.joinToString(", ")})"
}

private fun truncateFor(key: String, text: String) =
    if (text.length <= 80) text else text.take(80) + "…"

/** Streamed chunk shapes. `content` stays a raw JSON element because servers emit
 *  either a string or an array of typed parts — handle both in [flattenContent]. */
@Serializable
data class StreamChunk(val choices: List<StreamChoice> = emptyList())

@Serializable
data class StreamChoice(val delta: StreamDelta? = null)

@Serializable
data class StreamDelta(
    val content: JsonElement? = null,
    @SerialName("tool_calls") val toolCalls: List<ToolCallDelta> = emptyList(),
)

@Serializable
data class ToolCallDelta(
    val index: Int = 0,
    val id: String? = null,
    val function: DeltaFunction? = null,
)

@Serializable
data class DeltaFunction(val name: String? = null, val arguments: String? = null)

/** Flatten a delta `content` field, string or parts list, to the text inside. */
fun flattenContent(el: JsonElement?): String? = when (el) {
    null, JsonNull -> null
    is JsonPrimitive -> el.content
    is JsonArray -> el.joinToString("") { part ->
        ((part as? JsonObject)?.get("text") as? JsonPrimitive)?.content ?: ""
    }
    is JsonObject -> ((el["text"]) as? JsonPrimitive)?.content
}

/** Something surfaced to the UI mid-turn, one stream like the Rust Item enum. */
sealed interface Item {
    data class Text(val delta: String) : Item
    data class ToolStart(val name: String, val summary: String) : Item
    data class ToolDone(val outcome: String) : Item
}

/** One local image to attach to the prompt (mirrors Image in llm.rs). */
data class AgentImage(val mime: String, val data: ByteArray) {
    /** data:<mime>;base64,<data> for the image_url part. */
    fun dataUrl(): String =
        "data:$mime;base64," + android.util.Base64.encodeToString(data, android.util.Base64.NO_WRAP)

    fun toPart() = buildJsonObject {
        put("type", "image_url")
        put("image_url", buildJsonObject { put("url", dataUrl()) })
    }

    override fun equals(other: Any?) = other is AgentImage && other.mime == mime && other.data.contentEquals(data)
    override fun hashCode() = mime.hashCode() * 31 + data.contentHashCode()
}

internal val json = kotlinx.serialization.json.Json {
    ignoreUnknownKeys = true
    encodeDefaults = false
}