package dev.kaizen.app.net

import kotlinx.serialization.Serializable
import kotlinx.serialization.json.Json

private val USAGE = 0.toChar()
private val REASONING = 1.toChar()
private val SOURCES = 2.toChar()
private val TOOL = 3.toChar()
private val QUERY = 4.toChar()

private val lenientJson = Json { ignoreUnknownKeys = true }

@Serializable
data class ToolEvent(
    val type: String = "",
    val name: String = "",
    val args: kotlinx.serialization.json.JsonObject? = null,
    val result: String? = null,
    val error: String? = null,
    val elapsed: Long? = null,
)

@Serializable
data class SearchSource(
    val title: String = "",
    val url: String = "",
    val snippet: String = "",
)

@Serializable
private data class UsageJson(
    val i: Int = 0,
    val o: Int = 0,
)

data class StreamState(
    val content: String = "",
    val reasoning: String = "",
    val tools: List<ToolEvent> = emptyList(),
    val query: String = "",
    val sources: List<SearchSource> = emptyList(),
    val inputTokens: Int = 0,
    val outputTokens: Int = 0,
)

object StreamConsumer {

    fun visibleText(raw: String): String {
        val rest = stripLeadingToolEvents(raw) ?: return ""
        val segments = rest.split(REASONING)
        val sb = StringBuilder()
        for (i in segments.indices) {
            if (i % 2 == 0) sb.append(segments[i])
        }
        var text = sb.toString()
        val cut = listOf(text.indexOf(QUERY), text.indexOf(SOURCES), text.indexOf(USAGE))
            .filter { it != -1 }.minOrNull()
        if (cut != null) text = text.substring(0, cut)
        return text
    }

    private fun stripLeadingToolEvents(raw: String): String? {
        var rest = raw
        while (rest.isNotEmpty() && rest[0] == TOOL) {
            val nl = rest.indexOf('\n')
            if (nl == -1) return null
            rest = rest.substring(nl + 1)
        }
        return rest
    }

    class Incremental {
        private val content = StringBuilder(256)
        private val reasoningBuf = StringBuilder(256)
        private val toolBuffer = StringBuilder()
        private val trailingBuf = StringBuilder()
        private val tools = mutableListOf<ToolEvent>()
        private var pastToolEvents = false
        private var inReasoning = false
        private var inTrailing = false
        private var done = false
        private var trailingSentinel: Char? = null

        fun append(chars: CharArray, offset: Int, length: Int): StreamState {
            if (done) return buildState()

            if (!pastToolEvents) {
                toolBuffer.append(chars, offset, length)
                if (toolBuffer.isEmpty()) return buildState()
                val end = findToolEventsEnd(toolBuffer)
                if (end == -1) return buildState()
                parseToolEvents(toolBuffer, 0, end)
                pastToolEvents = true
                processRange(toolBuffer, end, toolBuffer.length)
                toolBuffer.setLength(0)
                toolBuffer.trimToSize()
            } else {
                processChars(chars, offset, length)
            }

            return buildState()
        }

        fun append(chunk: String): StreamState = append(chunk.toCharArray(), 0, chunk.length)

        fun currentContent(): String = content.toString()

        private fun buildState(): StreamState {
            val queryStr = parseTrailingField(QUERY)
            val sourcesList = parseTrailingSources()
            val usage = parseTrailingUsage()
            return StreamState(
                content = content.toString(),
                reasoning = reasoningBuf.toString(),
                tools = tools.toList(),
                query = queryStr,
                sources = sourcesList,
                inputTokens = usage?.i ?: 0,
                outputTokens = usage?.o ?: 0,
            )
        }

        private fun parseTrailingField(sentinel: Char): String {
            val full = trailingBuf.toString()
            val idx = full.indexOf(sentinel)
            if (idx == -1) return ""
            val after = full.substring(idx + 1)
            val end = listOf(after.indexOf(SOURCES), after.indexOf(USAGE), after.indexOf(QUERY))
                .filter { it != -1 }.minOrNull() ?: after.length
            return after.substring(0, end).trim()
        }

        private fun parseTrailingSources(): List<SearchSource> {
            val full = trailingBuf.toString()
            val idx = full.indexOf(SOURCES)
            if (idx == -1) return emptyList()
            val after = full.substring(idx + 1)
            val end = listOf(after.indexOf(USAGE)).filter { it != -1 }.minOrNull() ?: after.length
            val json = after.substring(0, end).trim()
            if (json.isEmpty()) return emptyList()
            return try {
                lenientJson.decodeFromString<List<SearchSource>>(json)
            } catch (_: Exception) {
                emptyList()
            }
        }

        private fun parseTrailingUsage(): UsageJson? {
            val full = trailingBuf.toString()
            val idx = full.indexOf(USAGE)
            if (idx == -1) return null
            val json = full.substring(idx + 1).trim()
            if (json.isEmpty()) return null
            return try { lenientJson.decodeFromString<UsageJson>(json) } catch (_: Exception) { null }
        }

        private fun parseToolEvents(buf: CharSequence, from: Int, to: Int) {
            val s = buf.toString()
            var pos = from
            while (pos < to && s[pos] == TOOL) {
                val nl = s.indexOf('\n', pos)
                if (nl == -1 || nl >= to) break
                val jsonStr = s.substring(pos + 1, nl)
                try {
                    tools.add(lenientJson.decodeFromString<ToolEvent>(jsonStr))
                } catch (_: Exception) { }
                pos = nl + 1
            }
        }

        private fun processRange(buf: CharSequence, from: Int, to: Int) {
            for (i in from until to) {
                if (done) break
                processChar(buf[i])
            }
        }

        private fun processChars(chars: CharArray, offset: Int, length: Int) {
            val end = offset + length
            for (i in offset until end) {
                if (done) break
                processChar(chars[i])
            }
        }

        private fun processChar(c: Char) {
            when {
                inTrailing -> trailingBuf.append(c)
                c == REASONING -> inReasoning = !inReasoning
                c == QUERY || c == SOURCES || c == USAGE -> {
                    inTrailing = true
                    trailingBuf.append(c)
                }
                inReasoning -> reasoningBuf.append(c)
                else -> content.append(c)
            }
        }

        companion object {
            private fun findToolEventsEnd(buf: StringBuilder): Int {
                var pos = 0
                val s = buf.toString()
                while (pos < s.length && s[pos] == TOOL) {
                    val nl = s.indexOf('\n', pos)
                    if (nl == -1) return -1
                    pos = nl + 1
                }
                return pos
            }
        }
    }
}
