package com.garfiec.librechat.core.network.sse

import com.garfiec.librechat.core.model.Conversation
import com.garfiec.librechat.core.model.Message
import com.garfiec.librechat.core.model.StreamEvent
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonElement
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.JsonPrimitive
import kotlinx.serialization.json.booleanOrNull
import kotlinx.serialization.json.contentOrNull
import kotlinx.serialization.json.intOrNull
import kotlinx.serialization.json.jsonArray
import kotlinx.serialization.json.jsonObject
import kotlinx.serialization.json.jsonPrimitive
import co.touchlab.kermit.Logger

/**
 * Maps raw SSE events into domain [StreamEvent]s.
 *
 * Supports both the LangGraph nested event format (used by agents endpoint)
 * and the legacy flat format. The backend sends events in this structure:
 *
 * **LangGraph events** have an `"event"` key:
 * ```json
 * {"event":"on_message_delta","data":{"id":"step_xxx","delta":{"content":[{"type":"text","text":"Hello"}]}}}
 * ```
 *
 * **Control events** have top-level flag keys:
 * ```json
 * {"final":true,"conversation":{...},"requestMessage":{...},"responseMessage":{...}}
 * {"created":{"message":{...}}}
 * {"sync":true,"resumeState":{...}}
 * ```
 *
 * NOTE: This class holds mutable state that is reset per-connection via [resetState].
 * It is NOT thread-safe. Only use from a single coroutine (as done in [SseClient.connect]).
 */
class SseEventMapper(private val json: Json) {

    // Track agent context from on_run_step for cross-event correlation.
    // on_message_delta events don't carry agentId/groupId; the server only
    // includes them on on_run_step events. We store the last values and
    // apply them to subsequent content events.
    private var activeAgentId: String? = null
    private var activeGroupId: Int? = null

    /** Resets tracked state. Call when starting a new SSE stream. */
    fun resetState() {
        activeAgentId = null
        activeGroupId = null
    }

    /**
     * Safely extracts a string from a [JsonElement] that may be a [JsonPrimitive] (string)
     * or a [JsonObject]/[JsonArray]. Tool call args and outputs from agent SSE streams
     * (e.g. image generation) arrive as JSON objects, not string primitives.
     * Returns [JsonPrimitive.contentOrNull] for primitives, or [JsonElement.toString]
     * (the raw JSON text) for objects and arrays.
     */
    private fun JsonElement?.toStringValue(): String? = when (this) {
        null -> null
        is JsonPrimitive -> contentOrNull
        else -> toString()
    }

    fun map(event: SseEvent): StreamEvent? {
        if (event.data.isBlank() || event.data == "[DONE]") return null

        return try {
            val root = json.parseToJsonElement(event.data).jsonObject
            // Check SSE event type for attachment events sent as
            // `event: attachment\ndata: {...}` (the data object won't have
            // an "event" key — it's the raw attachment metadata).
            if (event.event == "attachment" || event.event == "librechat:attachment") {
                return mapAttachment(root)
            }
            mapJsonObject(root)
        } catch (e: Exception) {
            Logger.w("SSE", e) { "SSE parse error: event=${event.event}" }
            StreamEvent.Error(message = "Parse error: ${e.message}")
        }
    }

    private fun mapJsonObject(root: JsonObject): StreamEvent? {
        // 1. Check for "final" control event (highest priority)
        if (root["final"]?.jsonPrimitive?.booleanOrNull == true) {
            return mapFinalEvent(root)
        }

        // 2. Check for "created" control event
        if (root.containsKey("created")) {
            return mapCreatedEvent(root)
        }

        // 3. Check for "sync" control event
        if (root["sync"]?.jsonPrimitive?.booleanOrNull == true) {
            return mapSyncEvent(root)
        }

        // 4. Check for "error" field (may be a string or an object)
        val errorText = root["error"]?.toStringValue()
        if (errorText != null) {
            return StreamEvent.Error(message = errorText)
        }

        // 5. Check for LangGraph nested event (has "event" key)
        val eventType = root["event"]?.jsonPrimitive?.contentOrNull
        if (eventType != null) {
            return mapLangGraphEvent(eventType, root)
        }

        // 6. Legacy flat format fallback
        return mapLegacyEvent(root)
    }

    // --- Control events ---

    private fun mapFinalEvent(root: JsonObject): StreamEvent {
        val parseErrors = mutableListOf<String>()

        val conversation = root["conversation"]?.let {
            try { json.decodeFromJsonElement(Conversation.serializer(), it) } catch (e: Exception) {
                val msg = "Failed to parse final conversation: ${e.message}"
                Logger.w("SSE", e) { msg }
                parseErrors.add(msg)
                null
            }
        }
        val requestMessage = root["requestMessage"]?.let {
            try { json.decodeFromJsonElement(Message.serializer(), it) } catch (e: Exception) {
                val msg = "Failed to parse final requestMessage: ${e.message}"
                Logger.w("SSE", e) { msg }
                parseErrors.add(msg)
                null
            }
        }
        val responseMessage = root["responseMessage"]?.let {
            try { json.decodeFromJsonElement(Message.serializer(), it) } catch (e: Exception) {
                val msg = "Failed to parse final responseMessage: ${e.message}"
                Logger.w("SSE", e) { msg }
                parseErrors.add(msg)
                null
            }
        }
        // Legacy: some events use "message" instead of "responseMessage"
        val legacyMessage = root["message"]?.let {
            try { json.decodeFromJsonElement(Message.serializer(), it) } catch (e: Exception) {
                val msg = "Failed to parse final legacy message: ${e.message}"
                Logger.w("SSE", e) { msg }
                parseErrors.add(msg)
                null
            }
        }

        // If ALL critical fields failed to parse, emit an Error instead
        val allFieldsNull = conversation == null && requestMessage == null &&
            responseMessage == null && legacyMessage == null
        if (allFieldsNull && parseErrors.isNotEmpty()) {
            Logger.e("SSE") { "Final event: all fields failed to parse -- $parseErrors" }
            return StreamEvent.Error(
                message = "Failed to parse final event: ${parseErrors.joinToString("; ")}",
            )
        }

        return StreamEvent.Final(
            message = legacyMessage,
            conversation = conversation,
            requestMessage = requestMessage,
            responseMessage = responseMessage,
            parseErrors = parseErrors,
        )
    }

    private fun mapCreatedEvent(root: JsonObject): StreamEvent.Created {
        // "created" can be either a boolean or an object containing message info
        val createdObj = root["created"]
        val messageObj = when {
            // Format: {"created": {"message": {...}}}
            createdObj is JsonObject -> createdObj["message"]?.jsonObject
            // Format: {"created": true, "message": {...}}
            else -> root["message"]?.jsonObject
        }

        val conversationId = messageObj?.get("conversationId")?.jsonPrimitive?.contentOrNull ?: ""
        val messageId = messageObj?.get("messageId")?.jsonPrimitive?.contentOrNull ?: ""
        val parentMessageId = messageObj?.get("parentMessageId")?.jsonPrimitive?.contentOrNull ?: ""

        return StreamEvent.Created(
            conversationId = conversationId,
            messageId = messageId,
            parentMessageId = parentMessageId,
        )
    }

    private fun mapSyncEvent(root: JsonObject): StreamEvent? {
        val resumeState = root["resumeState"]?.jsonObject ?: return null

        // Replay run steps as individual events first — the caller will handle them
        // The web client replays these as on_run_step events; for our mapper we
        // focus on the aggregatedContent snapshot that replaces current message state.
        val runSteps = resumeState["runSteps"]?.jsonArray

        // Parse aggregatedContent — this is the full content array that replaces
        // whatever the client currently has for this response message.
        val aggregatedContent = resumeState["aggregatedContent"]?.jsonArray
        if (aggregatedContent == null && runSteps == null) return null

        val contentParts = if (aggregatedContent != null) {
            aggregatedContent.mapNotNull { element ->
                try {
                    json.decodeFromJsonElement(
                        com.garfiec.librechat.core.model.MessageContentPart.serializer(),
                        element,
                    )
                } catch (e: Exception) {
                    Logger.w("SSE", e) { "Failed to parse sync aggregatedContent part" }
                    null
                }
            }
        } else {
            emptyList()
        }

        return StreamEvent.Sync(aggregatedContent = contentParts)
    }

    // --- LangGraph events ---

    private fun mapLangGraphEvent(eventType: String, root: JsonObject): StreamEvent? {
        val data = root["data"]?.jsonObject ?: return null
        val metadata = data["metadata"]?.jsonObject

        // Server puts agentId/groupId at top-level of data for on_run_step,
        // and optionally in data.metadata for other events.
        val eventAgentId = data["agentId"]?.jsonPrimitive?.contentOrNull
            ?: metadata?.get("agentId")?.jsonPrimitive?.contentOrNull
        val eventGroupId = data["groupId"]?.jsonPrimitive?.intOrNull
            ?: metadata?.get("groupId")?.jsonPrimitive?.intOrNull

        // Track agent context from run steps for cross-event correlation
        if (eventType == "on_run_step" && (eventAgentId != null || eventGroupId != null)) {
            activeAgentId = eventAgentId
            activeGroupId = eventGroupId
        }

        // Resolve: use event-level values if present, otherwise fall back to tracked state
        val agentId = eventAgentId ?: activeAgentId
        val groupId = eventGroupId ?: activeGroupId

        return when (eventType) {
            "on_message_delta" -> mapMessageDelta(data, agentId, groupId)
            "on_reasoning_delta" -> mapReasoningDelta(data, agentId, groupId)
            "on_run_step" -> mapRunStep(data, agentId, groupId)
            "on_run_step_delta" -> null // Tool call argument streaming - not currently tracked
            "on_run_step_completed" -> mapRunStepCompleted(data, agentId, groupId)
            "on_chat_model_end" -> null
            "on_agent_update" -> null
            "attachment" -> mapAttachment(data)
            else -> null
        }
    }

    private fun mapMessageDelta(
        data: JsonObject,
        agentId: String?,
        groupId: Int?,
    ): StreamEvent? {
        // {"id":"step_xxx","delta":{"content":[{"type":"text","text":"Hello"}]}}
        val delta = data["delta"]?.jsonObject ?: return null
        val contentArray = delta["content"]?.jsonArray ?: return null

        val textBuilder = StringBuilder()
        for (element in contentArray) {
            val part = element.jsonObject
            val type = part["type"]?.jsonPrimitive?.contentOrNull
            if (type == "text") {
                val text = part["text"]?.jsonPrimitive?.contentOrNull
                if (text != null) textBuilder.append(text)
            }
        }

        val text = textBuilder.toString()
        if (text.isEmpty()) return null

        return StreamEvent.ContentDelta(
            chunk = text,
            messageId = data["id"]?.jsonPrimitive?.contentOrNull,
            agentId = agentId,
            groupId = groupId,
        )
    }

    private fun mapReasoningDelta(
        data: JsonObject,
        agentId: String?,
        groupId: Int?,
    ): StreamEvent? {
        // {"id":"step_xxx","delta":{"content":[{"type":"think","think":"..."}]}}
        val delta = data["delta"]?.jsonObject ?: return null
        val contentArray = delta["content"]?.jsonArray ?: return null

        val thinkBuilder = StringBuilder()
        for (element in contentArray) {
            val part = element.jsonObject
            val type = part["type"]?.jsonPrimitive?.contentOrNull
            if (type == "think") {
                val think = part["think"]?.jsonPrimitive?.contentOrNull
                if (think != null) thinkBuilder.append(think)
            }
        }

        val text = thinkBuilder.toString()
        if (text.isEmpty()) return null

        return StreamEvent.ThinkingDelta(
            chunk = text,
            agentId = agentId,
            groupId = groupId,
        )
    }

    private fun mapRunStep(
        data: JsonObject,
        agentId: String?,
        groupId: Int?,
    ): StreamEvent? {
        // {"id":"step_xxx","stepDetails":{"type":"tool_calls","tool_calls":[{"id":"call_xxx","name":"search","args":""}]}}
        val stepDetails = data["stepDetails"]?.jsonObject ?: return null
        val toolCalls = stepDetails["tool_calls"]?.jsonArray ?: return null
        val firstToolCall = toolCalls.firstOrNull()?.jsonObject ?: return null

        val toolCallId = firstToolCall["id"]?.jsonPrimitive?.contentOrNull ?: return null
        val toolName = firstToolCall["name"]?.jsonPrimitive?.contentOrNull ?: ""
        val args = firstToolCall["args"]?.toStringValue() ?: ""

        return StreamEvent.ToolCallStart(
            toolCallId = toolCallId,
            toolName = toolName,
            input = args,
            agentId = agentId,
            groupId = groupId,
        )
    }

    private fun mapRunStepCompleted(
        data: JsonObject,
        agentId: String?,
        groupId: Int?,
    ): StreamEvent? {
        // {"result":{"id":"step_xxx","tool_call":{"id":"call_xxx","name":"search","output":"..."}}}
        val result = data["result"]?.jsonObject ?: return null
        val toolCall = result["tool_call"]?.jsonObject ?: return null

        val toolCallId = toolCall["id"]?.jsonPrimitive?.contentOrNull ?: return null
        val output = toolCall["output"]?.toStringValue() ?: ""

        return StreamEvent.ToolCallComplete(
            toolCallId = toolCallId,
            output = output,
            agentId = agentId,
            groupId = groupId,
        )
    }

    private fun mapAttachment(data: JsonObject): StreamEvent? {
        // Attachment events carry file metadata for tool-generated artifacts
        // (e.g., images from DALL-E / image_edit_oai). The data object has:
        //   file_id, filename, filepath, type, width, height, toolCallId, messageId, etc.
        val fileId = data["file_id"]?.jsonPrimitive?.contentOrNull ?: ""
        val filename = data["filename"]?.jsonPrimitive?.contentOrNull ?: ""
        val type = data["type"]?.jsonPrimitive?.contentOrNull ?: ""
        if (fileId.isBlank() && filename.isBlank()) return null
        return StreamEvent.AttachmentCreated(
            fileId = fileId,
            filename = filename,
            type = type,
            filepath = data["filepath"]?.jsonPrimitive?.contentOrNull,
            toolCallId = data["toolCallId"]?.jsonPrimitive?.contentOrNull,
            width = data["width"]?.jsonPrimitive?.intOrNull,
            height = data["height"]?.jsonPrimitive?.intOrNull,
        )
    }

    // --- Legacy flat format fallback ---

    private fun mapLegacyEvent(root: JsonObject): StreamEvent? {
        val type = root["type"]?.jsonPrimitive?.contentOrNull
        val metadata = root["metadata"]?.jsonObject
        val agentId = metadata?.get("agentId")?.jsonPrimitive?.contentOrNull
        val groupId = metadata?.get("groupId")?.jsonPrimitive?.intOrNull

        return when (type) {
            "content", "text" -> {
                val text = root["text"]?.jsonPrimitive?.contentOrNull ?: return null
                StreamEvent.ContentDelta(
                    chunk = text,
                    messageId = root["messageId"]?.jsonPrimitive?.contentOrNull,
                    agentId = agentId,
                    groupId = groupId,
                )
            }
            "thinking", "think" -> {
                val text = root["text"]?.jsonPrimitive?.contentOrNull ?: return null
                StreamEvent.ThinkingDelta(
                    chunk = text,
                    agentId = agentId,
                    groupId = groupId,
                )
            }
            "tool_call_start" -> {
                StreamEvent.ToolCallStart(
                    toolCallId = root["toolCallId"]?.jsonPrimitive?.contentOrNull ?: "",
                    toolName = root["toolName"]?.jsonPrimitive?.contentOrNull ?: "",
                    input = root["input"]?.toStringValue() ?: "",
                    agentId = agentId,
                    groupId = groupId,
                )
            }
            "tool_call_complete" -> {
                StreamEvent.ToolCallComplete(
                    toolCallId = root["toolCallId"]?.jsonPrimitive?.contentOrNull ?: "",
                    output = root["output"]?.toStringValue() ?: "",
                    agentId = agentId,
                    groupId = groupId,
                )
            }
            else -> {
                // Check for message property (another legacy format)
                val text = root["text"]?.jsonPrimitive?.contentOrNull
                    ?: root["response"]?.jsonPrimitive?.contentOrNull
                if (text != null) {
                    StreamEvent.ContentDelta(
                        chunk = text,
                        agentId = agentId,
                        groupId = groupId,
                    )
                } else {
                    null
                }
            }
        }
    }
}
