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

/**
 * Gemini Developer API provider (REST, no Firebase SDK dependency).
 *
 * Streaming endpoint: POST .../streamGenerateContent?alt=sse&key={apiKey}
 * Each SSE line: data: { candidates[0].content.parts[0].text | functionCall }
 */
class GeminiProvider(
    private val apiKey: String,
    private val model: String = "gemini-2.0-flash",
) : LlmProvider {

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

    private val baseUrl = "https://generativelanguage.googleapis.com/v1beta/models/$model"

    override suspend fun generate(messages: List<AiMessage>, schemas: List<ToolSchema>): String {
        val url = "$baseUrl:generateContent?key=$apiKey"
        val body = buildRequestBody(messages, schemas)
        val request = Request.Builder()
            .url(url)
            .post(body.toRequestBody("application/json".toMediaType()))
            .build()

        return client.newCall(request).execute().use { response ->
            val json = gson.fromJson(response.body?.string() ?: "{}", JsonObject::class.java)
            json.getAsJsonArray("candidates")
                ?.get(0)?.asJsonObject
                ?.getAsJsonObject("content")
                ?.getAsJsonArray("parts")
                ?.get(0)?.asJsonObject
                ?.get("text")?.asString ?: ""
        }
    }

    override fun stream(messages: List<AiMessage>, schemas: List<ToolSchema>): Flow<ProviderEvent> = flow {
        val url = "$baseUrl:streamGenerateContent?alt=sse&key=$apiKey"
        val body = buildRequestBody(messages, schemas)
        val request = Request.Builder()
            .url(url)
            .post(body.toRequestBody("application/json".toMediaType()))
            .build()

        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 json = line.removePrefix("data: ")
                        parseStreamChunk(json)?.let { emit(it) }
                    }
                }
            }
        } catch (e: IOException) {
            Timber.e(e, "Gemini stream IO error")
            throw e
        }
        emit(ProviderEvent.Done)
    }.flowOn(Dispatchers.IO)

    // ------------------------------------------------------------------
    // Request / response helpers
    // ------------------------------------------------------------------

    private fun buildRequestBody(messages: List<AiMessage>, schemas: List<ToolSchema>): String {
        val contents = messages
            .filter { it.role != MessageRole.SYSTEM }
            .map { msg -> msg.toGeminiContent() }

        val payload = mutableMapOf<String, Any>("contents" to contents)

        if (schemas.isNotEmpty()) {
            payload["tools"] = listOf(
                mapOf("function_declarations" to schemas.map { it.toGeminiDeclaration() })
            )
        }

        // Inject system prompt as first user turn if present
        val systemMsg = messages.firstOrNull { it.role == MessageRole.SYSTEM }
        if (systemMsg != null) {
            payload["system_instruction"] = mapOf(
                "parts" to listOf(mapOf("text" to systemMsg.content))
            )
        }

        return gson.toJson(payload)
    }

    private fun parseStreamChunk(json: String): ProviderEvent? {
        return try {
            val obj = gson.fromJson(json, JsonObject::class.java)
            val part = obj.getAsJsonArray("candidates")
                ?.get(0)?.asJsonObject
                ?.getAsJsonObject("content")
                ?.getAsJsonArray("parts")
                ?.get(0)?.asJsonObject ?: return null

            when {
                part.has("text") ->
                    ProviderEvent.Token(part.get("text").asString)
                part.has("functionCall") -> {
                    val fc = part.getAsJsonObject("functionCall")
                    val args = gson.fromJson(fc.getAsJsonObject("args"), Map::class.java)
                        ?.mapKeys { it.key.toString() }
                        ?.mapValues { it.value ?: "" }
                        ?: emptyMap()
                    @Suppress("UNCHECKED_CAST")
                    ProviderEvent.ToolCall(
                        id = fc.get("name").asString,
                        toolName = fc.get("name").asString,
                        arguments = args as Map<String, Any>,
                    )
                }
                else -> null
            }
        } catch (e: Exception) {
            Timber.w(e, "Gemini parse chunk failed: %s", json)
            null
        }
    }

    private fun AiMessage.toGeminiContent(): Map<String, Any> = when (role) {
        MessageRole.USER -> mapOf(
            "role" to "user",
            "parts" to listOf(mapOf("text" to content)),
        )
        MessageRole.ASSISTANT -> mapOf(
            "role" to "model",
            "parts" to listOf(mapOf("text" to content)),
        )
        MessageRole.TOOL_CALL -> mapOf(
            "role" to "model",
            "parts" to listOf(
                mapOf(
                    "functionCall" to mapOf(
                        "name" to toolName.orEmpty(),
                        "args" to toolArguments,
                    )
                )
            ),
        )
        MessageRole.TOOL_RESULT -> mapOf(
            "role" to "user",
            "parts" to listOf(
                mapOf(
                    "functionResponse" to mapOf(
                        "name" to toolName.orEmpty(),
                        "response" to mapOf("result" to content),
                    )
                )
            ),
        )
        MessageRole.SYSTEM -> error("System messages are handled separately")
    }

    private fun ToolSchema.toGeminiDeclaration(): Map<String, Any> = 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(),
        ),
    )
}
