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

class AnthropicProvider(
    private val apiKey: String,
    private val model: String,
) : 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("content")
                ?.get(0)?.asJsonObject
                ?.get("text")?.asString ?: ""
        }
    }

    override fun stream(messages: List<AiMessage>, schemas: List<ToolSchema>): Flow<ProviderEvent> = flow {
        val request = buildRequest(messages, schemas, stream = true)

        var currentToolCallId = ""
        var currentToolCallName = ""
        var currentToolCallArgs = ""

        try {
            client.newCall(request).execute().use { response ->
            response.body?.source()?.use { source ->
                var currentEvent = ""
                while (!source.exhausted()) {
                    val line = source.readUtf8Line() ?: continue
                    if (line.isBlank()) continue

                    if (line.startsWith("event: ")) {
                        currentEvent = line.removePrefix("event: ").trim()
                        continue
                    }

                    if (line.startsWith("data: ")) {
                        val data = line.removePrefix("data: ").trim()
                        if (data.isBlank() || currentEvent.isBlank()) continue

                        try {
                            val json = gson.fromJson(data, JsonObject::class.java)

                            when (currentEvent) {
                                "content_block_start" -> {
                                    val contentBlock = json.getAsJsonObject("content_block")
                                    if (contentBlock?.get("type")?.asString == "tool_use") {
                                        currentToolCallId = contentBlock.get("id")?.asString ?: ""
                                        currentToolCallName = contentBlock.get("name")?.asString ?: ""
                                        currentToolCallArgs = ""
                                    }
                                }
                                "content_block_delta" -> {
                                    val delta = json.getAsJsonObject("delta")
                                    val type = delta?.get("type")?.asString
                                    if (type == "text_delta") {
                                        delta.get("text")?.asString?.let { emit(ProviderEvent.Token(it)) }
                                    } else if (type == "input_json_delta") {
                                        delta.get("partial_json")?.asString?.let { currentToolCallArgs += it }
                                    }
                                }
                                "content_block_stop" -> {
                                    if (currentToolCallName.isNotEmpty()) {
                                        val args = runCatching {
                                            @Suppress("UNCHECKED_CAST")
                                            gson.fromJson(currentToolCallArgs, Map::class.java) as Map<String, Any>
                                        }.getOrDefault(emptyMap())
                                        emit(ProviderEvent.ToolCall(currentToolCallId, currentToolCallName, args))

                                        currentToolCallId = ""
                                        currentToolCallName = ""
                                        currentToolCallArgs = ""
                                    }
                                }
                                "message_stop" -> {
                                    break
                                }
                            }
                        } catch (e: Exception) {
                            Timber.w(e, "Anthropic parse chunk failed: %s", data)
                        }
                    }
                }
            }
        }
        } catch (e: IOException) {
            Timber.e(e, "Anthropic stream IO error")
            throw e
        }
        emit(ProviderEvent.Done)
    }.flowOn(Dispatchers.IO)

    private fun buildRequest(
        messages: List<AiMessage>,
        schemas: List<ToolSchema>,
        stream: Boolean,
    ): Request {
        val systemMessage = messages.firstOrNull { it.role == MessageRole.SYSTEM }?.content ?: ""
        val nonSystemMessages = messages.filter { it.role != MessageRole.SYSTEM }

        val payload = mutableMapOf<String, Any>(
            "model" to model,
            "max_tokens" to 4096,
            "stream" to stream,
            "system" to systemMessage,
            "messages" to nonSystemMessages.map { it.toAnthropicMessage() },
        )

        if (schemas.isNotEmpty()) {
            payload["tools"] = schemas.map { it.toAnthropicTool() }
            // tool_choice could be specified here if needed
            // payload["tool_choice"] = mapOf("type" to "auto") 
        }

        return Request.Builder()
            .url("https://api.anthropic.com/v1/messages")
            .header("x-api-key", apiKey)
            .header("anthropic-version", "2023-06-01")
            .header("content-type", "application/json")
            .post(gson.toJson(payload).toRequestBody("application/json".toMediaType()))
            .build()
    }

    private fun AiMessage.toAnthropicMessage(): 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 -> emptyMap() // Filtered out above
        MessageRole.TOOL_CALL -> mapOf(
            "role" to "assistant",
            "content" to listOf(
                mapOf(
                    "type" to "tool_use",
                    "id" to (toolCallId ?: toolName.orEmpty()),
                    "name" to toolName.orEmpty(),
                    "input" to (toolArguments ?: emptyMap<String, Any>())
                )
            )
        )
        MessageRole.TOOL_RESULT -> mapOf(
            "role" to "user",
            "content" to listOf(
                mapOf(
                    "type" to "tool_result",
                    "tool_use_id" to (toolCallId ?: toolName.orEmpty()),
                    "content" to content,
                )
            )
        )
    }

    private fun ToolSchema.toAnthropicTool(): Map<String, Any> = mapOf(
        "name" to name,
        "description" to description,
        "input_schema" 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(),
        ),
    )
}
