package wang.yaojia.webterm.wire

import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonNull
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.JsonPrimitive
import kotlinx.serialization.json.booleanOrNull
import kotlinx.serialization.json.doubleOrNull
import kotlinx.serialization.json.intOrNull
import kotlinx.serialization.json.longOrNull

/**
 * Pure codec for the WS text-frame protocol (frozen contract, plan §3.1). NEVER throws in
 * either direction — the Android analogue of iOS `WireProtocol.MessageCodec`.
 *
 *  - [encode]        client → server. HAND-ROLLED so the bytes are IDENTICAL to the web
 *                    client / server `JSON.stringify` (explicit `sessionId:null`, top-level
 *                    `approve.mode`, JS control-char escaping). Every [ClientMessage] has a
 *                    representation, so encode is total.
 *  - [decodeServer]  server → client (UNTRUSTED). Tolerant kotlinx.serialization decode into a
 *                    [ServerMessage]; a malformed / unknown / wrong-typed-required frame → null
 *                    (drop + count, never crash — mirrors the server's own resilience).
 *  - [decodeClient]  the boundary validator — the Android mirror of the server's
 *                    `parseClientMessage` (type whitelist, `resize` 1..1000, `input.data`
 *                    passthrough, `attach.sessionId` UUID-v4 / explicit-null). Invalid → null.
 *                    Exact inverse of [encode] for round-trip verification.
 */
public object MessageCodec {
    private val json = Json {
        ignoreUnknownKeys = true
        isLenient = true
    }

    // ─── encode (client → server) ────────────────────────────────────────────────

    /**
     * Encode a [ClientMessage] as a JSON text frame the server's `parseClientMessage` accepts,
     * BYTE-IDENTICAL to the web client. String escaping matches `JSON.stringify` (short escapes
     * for `\b \t \n \f \r`, `\u00xx` lowercase-hex for the other C0 control bytes); `attach`
     * always carries an explicit `sessionId` key (null when absent); UUIDs are lowercased;
     * `approve.mode` is a TOP-LEVEL key (the server re-reads the raw frame for it).
     */
    public fun encode(message: ClientMessage): String = when (message) {
        is ClientMessage.Attach -> encodeAttach(message.sessionId, message.cwd)
        is ClientMessage.Input -> """{"type":"input","data":${jsonString(message.data)}}"""
        is ClientMessage.Resize -> """{"type":"resize","cols":${message.cols},"rows":${message.rows}}"""
        is ClientMessage.Approve ->
            message.mode?.let { """{"type":"approve","mode":"${it.wire}"}""" } ?: """{"type":"approve"}"""
        ClientMessage.Reject -> """{"type":"reject"}"""
    }

    private fun encodeAttach(sessionId: String?, cwd: String?): String {
        // The sessionId is escaped via jsonString (like every other string field) so encode stays
        // TOTAL and byte-identical to JSON.stringify — never interpolate an id raw (JSON injection).
        val sid = if (sessionId != null) jsonString(sessionId.lowercase()) else "null"
        val cwdPart = if (cwd != null) ",\"cwd\":${jsonString(cwd)}" else ""
        return """{"type":"attach","sessionId":$sid$cwdPart}"""
    }

    /** JSON string literal with `JSON.stringify`-identical escaping (iterates UTF-16 units). */
    private fun jsonString(value: String): String {
        val sb = StringBuilder(value.length + 2)
        sb.append('"')
        var i = 0
        while (i < value.length) {
            val c = value[i]
            when (c) {
                '"' -> sb.append("\\\"")
                '\\' -> sb.append("\\\\")
                '\b' -> sb.append("\\b") // 0x08
                '\t' -> sb.append("\\t") // 0x09
                '\n' -> sb.append("\\n") // 0x0A
                '\u000C' -> sb.append("\\f") // 0x0C
                '\r' -> sb.append("\\r") // 0x0D
                else ->
                    when {
                        c.code < 0x20 ->
                            sb.append("\\u").append(c.code.toString(16).padStart(4, '0'))
                        // Valid high+low surrogate pair: emit both UTF-16 units verbatim
                        // (JSON.stringify keeps the astral character as-is).
                        c.isHighSurrogate() && i + 1 < value.length && value[i + 1].isLowSurrogate() -> {
                            sb.append(c).append(value[i + 1])
                            i++ // consume the low surrogate too
                        }
                        // Lone surrogate (a high without a low, or a bare low): ES2019
                        // well-formed JSON.stringify emits \udXXX lowercase-hex.
                        c.isSurrogate() ->
                            sb.append("\\u").append(c.code.toString(16).padStart(4, '0'))
                        else -> sb.append(c)
                    }
            }
            i++
        }
        sb.append('"')
        return sb.toString()
    }

    // ─── decodeServer (server → client, untrusted) ───────────────────────────────

    /**
     * Decode an untrusted server text frame into a [ServerMessage], or null to DROP it (bad
     * JSON, non-object, unknown `type`, or a missing/wrong-typed REQUIRED field). Wrong-typed
     * OPTIONAL fields are tolerated as absent (mirrors the server's tolerant parsing). Unknown
     * keys are ignored. Never throws.
     */
    public fun decodeServer(raw: String): ServerMessage? {
        val root = parseObject(raw) ?: return null
        val type = root.stringField("type") ?: return null
        return when (ServerMessageType.fromWire(type)) {
            ServerMessageType.ATTACHED -> decodeAttached(root)
            ServerMessageType.OUTPUT -> root.stringField("data")?.let { ServerMessage.Output(it) }
            ServerMessageType.EXIT -> decodeExit(root)
            ServerMessageType.STATUS -> decodeStatus(root)
            ServerMessageType.TELEMETRY -> decodeTelemetry(root)
            null -> null
        }
    }

    private fun decodeAttached(obj: JsonObject): ServerMessage? {
        val raw = obj.stringField("sessionId") ?: return null
        if (!Validation.isValidSessionId(raw)) return null
        return ServerMessage.Attached(raw.lowercase())
    }

    private fun decodeExit(obj: JsonObject): ServerMessage? {
        val code = obj.intField("code") ?: return null
        return ServerMessage.Exit(code, obj.stringField("reason"))
    }

    private fun decodeStatus(obj: JsonObject): ServerMessage? {
        val status = obj.stringField("status")?.let { ClaudeStatus.fromWire(it) } ?: return null
        val detail = obj.stringField("detail")
        val pending = obj.boolField("pending") ?: false
        val gate = obj.stringField("gate")?.let { GateKind.fromWire(it) }
        return ServerMessage.Status(status, detail, pending, gate)
    }

    private fun decodeTelemetry(obj: JsonObject): ServerMessage? {
        val tel = obj["telemetry"] as? JsonObject ?: return null
        val at = tel.longField("at") ?: return null // only required field
        return ServerMessage.Telemetry(
            StatusTelemetry(
                contextUsedPct = tel.doubleField("contextUsedPct"),
                costUsd = tel.doubleField("costUsd"),
                linesAdded = tel.intField("linesAdded"),
                linesRemoved = tel.intField("linesRemoved"),
                model = tel.stringField("model"),
                effort = tel.stringField("effort"),
                pr = tel.decodeNested("pr", PrInfo.serializer()),
                rate = tel.decodeNested("rate", RateInfo.serializer()),
                at = at,
            ),
        )
    }

    // ─── decodeClient (boundary validator — server-parseClientMessage mirror) ─────

    /**
     * Validate + decode a client text frame the way the server's `parseClientMessage` would
     * ACCEPT/REJECT it: `type` in [ClientMessageType]; `resize` cols/rows non-string integers in
     * 1..1000; `input.data` a string (passed through VERBATIM); `attach.sessionId` present as
     * explicit null or a UUID-v4; `attach.cwd` (when present) an absolute path. Invalid → null.
     * The exact inverse of [encode], so `decodeClient(encode(m)) == m` for canonical `m`.
     */
    public fun decodeClient(raw: String): ClientMessage? {
        val root = parseObject(raw) ?: return null
        val type = root.stringField("type") ?: return null
        return when (ClientMessageType.fromWire(type)) {
            ClientMessageType.ATTACH -> decodeAttach(root)
            ClientMessageType.INPUT -> root.stringField("data")?.let { ClientMessage.Input(it) }
            ClientMessageType.RESIZE -> decodeResize(root)
            ClientMessageType.APPROVE ->
                ClientMessage.Approve(root.stringField("mode")?.let { ApproveMode.fromWire(it) })
            ClientMessageType.REJECT -> ClientMessage.Reject
            null -> null
        }
    }

    private fun decodeResize(obj: JsonObject): ClientMessage? {
        val cols = obj.intField("cols")?.takeIf { it in WireConstants.RESIZE_RANGE } ?: return null
        val rows = obj.intField("rows")?.takeIf { it in WireConstants.RESIZE_RANGE } ?: return null
        return ClientMessage.Resize(cols, rows)
    }

    private fun decodeAttach(obj: JsonObject): ClientMessage? {
        if ("sessionId" !in obj) return null // the key is REQUIRED (present as null or uuid)

        val cwd: String? =
            obj["cwd"]?.let { el ->
                val c = (el as? JsonPrimitive)?.takeIf { it.isString }?.content ?: return null
                if (!Validation.isAbsoluteCwd(c)) return null
                c
            }

        val sidEl = obj.getValue("sessionId")
        if (sidEl is JsonNull) return ClientMessage.Attach(null, cwd)
        val sid = (sidEl as? JsonPrimitive)?.takeIf { it.isString }?.content ?: return null
        if (!Validation.isValidSessionId(sid)) return null
        return ClientMessage.Attach(sid.lowercase(), cwd)
    }

    // ─── shared JSON helpers (tolerant, never throw) ─────────────────────────────

    private fun parseObject(raw: String): JsonObject? =
        runCatching { json.parseToJsonElement(raw) }.getOrNull() as? JsonObject

    /** A JSON STRING value only (a number/bool/null/object → null → treated as absent). */
    private fun JsonObject.stringField(key: String): String? =
        (this[key] as? JsonPrimitive)?.takeIf { it.isString }?.content

    /** A JSON non-string INTEGER only (`"80"`, `80.5`, null → null). */
    private fun JsonObject.intField(key: String): Int? =
        (this[key] as? JsonPrimitive)?.takeIf { !it.isString }?.intOrNull

    /** A JSON non-string integral LONG only (for the millisecond `at` timestamp). */
    private fun JsonObject.longField(key: String): Long? =
        (this[key] as? JsonPrimitive)?.takeIf { !it.isString }?.longOrNull

    /** A JSON non-string NUMBER as Double (a quoted number/other → null → treated as absent). */
    private fun JsonObject.doubleField(key: String): Double? =
        (this[key] as? JsonPrimitive)?.takeIf { !it.isString }?.doubleOrNull

    /** A JSON BOOLEAN only (a quoted "true" / other → null → caller defaults to false). */
    private fun JsonObject.boolField(key: String): Boolean? =
        (this[key] as? JsonPrimitive)?.takeIf { !it.isString }?.booleanOrNull

    /** Decode a nested @Serializable object tolerantly: any decode failure → null (absent). */
    private fun <T> JsonObject.decodeNested(
        key: String,
        deserializer: kotlinx.serialization.DeserializationStrategy<T>,
    ): T? = this[key]?.let { runCatching { json.decodeFromJsonElement(deserializer, it) }.getOrNull() }
}
