package dev.kaizen.app.live

import kotlinx.serialization.Serializable
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.boolean
import kotlinx.serialization.json.booleanOrNull
import kotlinx.serialization.json.double
import kotlinx.serialization.json.doubleOrNull
import kotlinx.serialization.json.int
import kotlinx.serialization.json.intOrNull
import kotlinx.serialization.json.jsonObject
import kotlinx.serialization.json.jsonPrimitive
import kotlinx.serialization.json.long
import kotlinx.serialization.json.longOrNull

private val json = Json { ignoreUnknownKeys = true; encodeDefaults = false }

// ─── Client → Server ──────────────────────────────────────────────────────

@Serializable
data class SetupMessage(
    val type: String = "setup",
    val model: String,
    val voice: String,
    val conversationId: String? = null,
    val tools: Boolean? = null,
)

@Serializable
data class ResumeMessage(
    val type: String = "resume",
    val sessionId: String,
    val lastSeqId: Int? = null,
)

@Serializable
data class AudioMessage(val type: String = "audio", val data: String)

@Serializable
data class TextMessage(val type: String = "text", val text: String)

@Serializable
data class SimpleMessage(val type: String)

fun encodeSetup(model: String, voice: String, conversationId: String?, tools: Boolean? = true): String =
    json.encodeToString(SetupMessage.serializer(), SetupMessage(model = model, voice = voice, conversationId = conversationId, tools = tools))

fun encodeResume(sessionId: String, lastSeqId: Int?): String =
    json.encodeToString(ResumeMessage.serializer(), ResumeMessage(sessionId = sessionId, lastSeqId = lastSeqId))

fun encodeAudio(base64Pcm: String): String =
    json.encodeToString(AudioMessage.serializer(), AudioMessage(data = base64Pcm))

fun encodeText(text: String): String =
    json.encodeToString(TextMessage.serializer(), TextMessage(text = text))

fun encodeInterrupt(): String = """{"type":"interrupt"}"""
fun encodeEnd(): String = """{"type":"end"}"""

// ─── Server → Client ──────────────────────────────────────────────────────

data class LiveUsage(
    val inputTokens: Int = 0,
    val outputTokens: Int = 0,
    val cost: Double? = null,
    val durationMs: Long = 0,
)

sealed interface ServerEvent {
    data class Ready(val sessionId: String) : ServerEvent
    data class Resumed(val sessionId: String) : ServerEvent
    data class Audio(val data: String, val seq: Int) : ServerEvent
    data class Text(val delta: String, val seq: Int) : ServerEvent
    data class Transcript(val text: String, val final: Boolean) : ServerEvent
    data class Vad(val speaking: Boolean) : ServerEvent
    data object Interrupted : ServerEvent
    data class TurnEnd(
        val usage: LiveUsage,
        val userText: String,
        val assistantText: String,
        val userMessageId: String? = null,
        val assistantMessageId: String? = null,
    ) : ServerEvent
    data class ToolCall(val name: String, val status: String) : ServerEvent
    data class Backpressure(val throttleMs: Int) : ServerEvent
    data class GoAway(val remainingMs: Int, val reason: String) : ServerEvent
    data class Error(val code: String, val message: String?) : ServerEvent
}

fun parseServerEvent(raw: String): ServerEvent? {
    val obj: JsonObject
    try {
        obj = json.parseToJsonElement(raw).jsonObject
    } catch (_: Exception) {
        return null
    }

    val type = obj["type"]?.jsonPrimitive?.content ?: return null

    return when (type) {
        "ready" -> ServerEvent.Ready(
            sessionId = obj["sessionId"]?.jsonPrimitive?.content ?: return null,
        )
        "resumed" -> ServerEvent.Resumed(
            sessionId = obj["sessionId"]?.jsonPrimitive?.content ?: return null,
        )
        "audio" -> ServerEvent.Audio(
            data = obj["data"]?.jsonPrimitive?.content ?: return null,
            seq = obj["seq"]?.jsonPrimitive?.intOrNull ?: 0,
        )
        "text" -> ServerEvent.Text(
            delta = obj["delta"]?.jsonPrimitive?.content ?: return null,
            seq = obj["seq"]?.jsonPrimitive?.intOrNull ?: 0,
        )
        "transcript" -> ServerEvent.Transcript(
            text = obj["text"]?.jsonPrimitive?.content ?: "",
            final = obj["final"]?.jsonPrimitive?.booleanOrNull ?: false,
        )
        "vad" -> ServerEvent.Vad(
            speaking = obj["speaking"]?.jsonPrimitive?.booleanOrNull ?: false,
        )
        "interrupted" -> ServerEvent.Interrupted
        "turn_end" -> {
            val u = obj["usage"]?.jsonObject
            ServerEvent.TurnEnd(
                usage = LiveUsage(
                    inputTokens = u?.get("inputTokens")?.jsonPrimitive?.intOrNull ?: 0,
                    outputTokens = u?.get("outputTokens")?.jsonPrimitive?.intOrNull ?: 0,
                    cost = u?.get("cost")?.jsonPrimitive?.doubleOrNull,
                    durationMs = u?.get("durationMs")?.jsonPrimitive?.longOrNull ?: 0,
                ),
                userText = obj["userText"]?.jsonPrimitive?.content ?: "",
                assistantText = obj["assistantText"]?.jsonPrimitive?.content ?: "",
                userMessageId = obj["userMessageId"]?.jsonPrimitive?.content,
                assistantMessageId = obj["assistantMessageId"]?.jsonPrimitive?.content,
            )
        }
        "tool_call" -> ServerEvent.ToolCall(
            name = obj["name"]?.jsonPrimitive?.content ?: "",
            status = obj["status"]?.jsonPrimitive?.content ?: "",
        )
        "backpressure" -> ServerEvent.Backpressure(
            throttleMs = obj["throttleMs"]?.jsonPrimitive?.intOrNull ?: 0,
        )
        "go_away" -> ServerEvent.GoAway(
            remainingMs = obj["remainingMs"]?.jsonPrimitive?.intOrNull ?: 0,
            reason = obj["reason"]?.jsonPrimitive?.content ?: "",
        )
        "error" -> ServerEvent.Error(
            code = obj["code"]?.jsonPrimitive?.content ?: "unknown",
            message = obj["message"]?.jsonPrimitive?.content,
        )
        else -> null
    }
}
