package com.optix.android.ai.providers

import com.google.gson.Gson
import com.google.gson.JsonObject
import com.optix.android.domain.model.AiMessage
import com.optix.android.domain.model.MessageRole
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.flow.flowOn
import okhttp3.MediaType.Companion.toMediaType
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.RequestBody.Companion.toRequestBody
import timber.log.Timber
import java.io.IOException
import java.util.concurrent.TimeUnit

/**
 * OpenAI-compatible provider — works with OpenAI, Groq, Mistral, OpenRouter, and Ollama
 * by swapping [baseUrl] and [apiKey].
 *
 * Streaming: POST {baseUrl}/chat/completions with stream=true
 * SSE format: data: {"choices":[{"delta":{"content":"..."}}]}
 */
class OpenAiProvider(
    private val apiKey: String,
    private val model: String,
    private val baseUrl: String = "https://api.openai.com/v1",
) : LlmProvider {

    private val gson = Gson()
    private val client = OkHttpClient.Builder()
        .connectTimeout(30, TimeUnit.SECONDS)
        .readTimeout(120, TimeUnit.SECONDS)
        .build()

    override suspend fun generate(messages: List<AiMessage>, schemas: List<ToolSchema>): String {
        val request = buildRequest(messages, schemas, stream = false)
        return client.newCall(request).execute().use { response ->
            val json = gson.fromJson(response.body?.string() ?: "{}", JsonObject::class.java)
            json.getAsJsonArray("choices")
                ?.get(0)?.asJsonObject
                ?.getAsJsonObject("message")
                ?.get("content")?.asString ?: ""
        }
    }

    override fun stream(messages: List<AiMessage>, schemas: List<ToolSchema>): Flow<ProviderEvent> = flow {
        val request = buildRequest(messages, schemas, stream = true)
        // Accumulate partial tool call arguments across chunks
        val toolCallBuffer = mutableMapOf<Int, ToolCallAccumulator>()

        try {
            client.newCall(request).execute().use { response ->
                response.body?.source()?.use { source ->
                    while (!source.exhausted()) {
                        val line = source.readUtf8Line() ?: continue
                    if (!line.startsWith("data: ")) continue
                    val data = line.removePrefix("data: ")
                    if (data == "[DONE]") break

                    try {
                        val json = gson.fromJson(data, JsonObject::class.java)
                        val delta = json.getAsJsonArray("choices")
                            ?.get(0)?.asJsonObject
                            ?.getAsJsonObject("delta") ?: continue

                        // Text token
                        delta.get("content")?.asString?.let { emit(ProviderEvent.Token(it)) }

                        // Tool call chunks — accumulate until complete
                        delta.getAsJsonArray("tool_calls")?.forEach { tcElem ->
                            val tc = tcElem.asJsonObject
                            val idx = tc.get("index")?.asInt ?: 0
                            val acc = toolCallBuffer.getOrPut(idx) { ToolCallAccumulator() }
                            tc.get("id")?.asString?.let { acc.id = it }
                            tc.getAsJsonObject("function")?.let { fn ->
                                fn.get("name")?.asString?.let { acc.name = it }
                                fn.get("arguments")?.asString?.let { acc.argsBuffer += it }
                            }
                        }
                    } catch (e: Exception) {
                        Timber.w(e, "Skipped malformed SSE chunk: %s", data)
                    }
                }
            }
        }
        } catch (e: IOException) {
            Timber.e(e, "OpenAI stream IO error")
            throw e
        }

        // Emit completed tool calls
        toolCallBuffer.values.forEach { acc ->
            if (acc.name.isNotEmpty()) {
                val args = runCatching {
                    @Suppress("UNCHECKED_CAST")
                    gson.fromJson(acc.argsBuffer, Map::class.java) as Map<String, Any>
                }.getOrDefault(emptyMap())
                emit(ProviderEvent.ToolCall(acc.id, acc.name, args))
            }
        }

        emit(ProviderEvent.Done)
    }.flowOn(Dispatchers.IO)

    // ------------------------------------------------------------------

    private fun buildRequest(
        messages: List<AiMessage>,
        schemas: List<ToolSchema>,
        stream: Boolean,
    ): Request {
        val payload = mutableMapOf<String, Any>(
            "model" to model,
            "stream" to stream,
            "messages" to messages.map { it.toOpenAiMessage() },
        )
        if (schemas.isNotEmpty()) {
            payload["tools"] = schemas.map { it.toOpenAiTool() }
            payload["tool_choice"] = "auto"
        }

        return Request.Builder()
            .url("$baseUrl/chat/completions")
            .header("Authorization", "Bearer $apiKey")
            .header("Content-Type", "application/json")
            .post(gson.toJson(payload).toRequestBody("application/json".toMediaType()))
            .build()
    }

    private fun AiMessage.toOpenAiMessage(): Map<String, Any> = when (role) {
        MessageRole.USER -> mapOf(
            "role" to "user",
            "content" to content,
        )
        MessageRole.ASSISTANT -> mapOf(
            "role" to "assistant",
            "content" to content,
        )
        MessageRole.SYSTEM -> mapOf(
            "role" to "system",
            "content" to content,
        )
        MessageRole.TOOL_CALL -> mapOf(
            "role" to "assistant",
            "content" to "",
            "tool_calls" to listOf(
                mapOf(
                    "id" to (toolCallId ?: toolName.orEmpty()),
                    "type" to "function",
                    "function" to mapOf(
                        "name" to toolName.orEmpty(),
                        "arguments" to gson.toJson(toolArguments),
                    ),
                )
            ),
        )
        MessageRole.TOOL_RESULT -> mapOf(
            "role" to "tool",
            "content" to content,
            "tool_call_id" to (toolCallId ?: toolName.orEmpty()),
            "name" to toolName.orEmpty(),
        )
    }

    private fun ToolSchema.toOpenAiTool(): Map<String, Any> = mapOf(
        "type" to "function",
        "function" to mapOf(
            "name" to name,
            "description" to description,
            "parameters" to mapOf(
                "type" to "object",
                "properties" to parameters.mapValues { (_, p) ->
                    buildMap {
                        put("type", p.type)
                        put("description", p.description)
                        p.enum?.let { put("enum", it) }
                    }
                },
                "required" to parameters.filter { it.value.required }.keys.toList(),
            ),
        ),
    )

    private data class ToolCallAccumulator(
        var id: String = "",
        var name: String = "",
        var argsBuffer: String = "",
    )
}
