package app.chompass.services.ai

import app.chompass.R

import kotlinx.serialization.json.Json
import kotlinx.serialization.json.contentOrNull
import kotlinx.serialization.json.jsonArray
import kotlinx.serialization.json.jsonObject
import kotlinx.serialization.json.jsonPrimitive
import okhttp3.MediaType.Companion.toMediaType
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.RequestBody.Companion.toRequestBody
import org.json.JSONArray
import org.json.JSONObject
import java.util.Base64

/**
 * Anthropic Messages format:
 *   POST <base>/messages
 *   Headers: x-api-key: <apiKey>, anthropic-version: 2023-06-01
 *   Body:    {model, max_tokens, system?, messages: [{role, content: [...]}]}
 */
object AnthropicClient {
    private const val API_VERSION = "2023-06-01"
    private val jsonMedia = "application/json; charset=utf-8".toMediaType()

    suspend fun analyze(
        client: OkHttpClient,
        baseUrl: String,
        model: String,
        apiKey: String,
        prompt: String,
        imageBytesList: List<ByteArray>,
        maxTokens: Int
    ): String {
        val url = "$baseUrl/messages"

        suspend fun request(requestPrompt: String): AnthropicTextResponse {
            val content = JSONArray().apply {
                imageBytesList.forEach {
                    put(
                        JSONObject()
                            .put("type", "image")
                            .put(
                                "source",
                                JSONObject()
                                    .put("type", "base64")
                                    .put("media_type", "image/jpeg")
                                    .put("data", Base64.getEncoder().encodeToString(it))
                            )
                    )
                }
                put(JSONObject().put("type", "text").put("text", requestPrompt))
            }

            val body = JSONObject()
                .put("model", model)
                .put("max_tokens", maxTokens)
                .put("messages", JSONArray().put(JSONObject().put("role", "user").put("content", content)))

            val bodyStr = RetryPolicy.execute {
                client.newCall(
                    Request.Builder()
                        .url(url)
                        .addHeader("Content-Type", "application/json")
                        .addHeader("x-api-key", apiKey)
                        .addHeader("anthropic-version", API_VERSION)
                        .post(body.toString().toRequestBody(jsonMedia))
                        .build()
                )
            }
            return AnthropicResponseParser.parse(bodyStr)
        }

        var response = request(prompt)
        if (response.wasTruncated) {
            response = request(compactRetryPrompt(prompt, maxTokens))
            if (response.wasTruncated) {
                throw AiError.Api("The AI response was truncated twice. Try a shorter description or another model.", messageRes = R.string.ai_error_truncated_twice_description)
            }
        }
        return response.text ?: throw AiError.InvalidResponse
    }

    /**
     * Streaming Messages API. Invokes [onDelta] with each text fragment and
     * returns the full concatenated text. Falls back to [analyze] if streaming fails.
     */
    suspend fun analyzeStreaming(
        client: OkHttpClient,
        baseUrl: String,
        model: String,
        apiKey: String,
        prompt: String,
        imageBytesList: List<ByteArray>,
        maxTokens: Int,
        onDelta: (String) -> Unit,
    ): String {
        val url = "$baseUrl/messages"

        suspend fun streamOnce(requestPrompt: String): Pair<String, Boolean> {
            val content = JSONArray().apply {
                imageBytesList.forEach {
                    put(
                        JSONObject()
                            .put("type", "image")
                            .put(
                                "source",
                                JSONObject()
                                    .put("type", "base64")
                                    .put("media_type", "image/jpeg")
                                    .put("data", Base64.getEncoder().encodeToString(it))
                            )
                    )
                }
                put(JSONObject().put("type", "text").put("text", requestPrompt))
            }

            val body = JSONObject()
                .put("model", model)
                .put("max_tokens", maxTokens)
                .put("stream", true)
                .put("messages", JSONArray().put(JSONObject().put("role", "user").put("content", content)))

            val assembled = StringBuilder()
            var stopReason: String? = null
            val response = RetryPolicy.open {
                client.newCall(
                    Request.Builder()
                        .url(url)
                        .addHeader("Content-Type", "application/json")
                        .addHeader("Accept", "text/event-stream")
                        .addHeader("x-api-key", apiKey)
                        .addHeader("anthropic-version", API_VERSION)
                        .post(body.toString().toRequestBody(jsonMedia))
                        .build()
                )
            }
            AiSse.read(response) { payload ->
                val json = runCatching { Json.parseToJsonElement(payload).jsonObject }.getOrNull()
                    ?: return@read
                when (json["type"]?.jsonPrimitive?.contentOrNull) {
                    "content_block_delta" -> {
                        val delta = json["delta"]?.jsonObject ?: return@read
                        if (delta["type"]?.jsonPrimitive?.contentOrNull != "text_delta") return@read
                        val piece = delta["text"]?.jsonPrimitive?.contentOrNull
                        if (!piece.isNullOrEmpty()) {
                            assembled.append(piece)
                            onDelta(piece)
                        }
                    }
                    "message_delta" -> {
                        stopReason = json["delta"]?.jsonObject
                            ?.get("stop_reason")?.jsonPrimitive?.contentOrNull
                            ?: stopReason
                    }
                }
            }
            return assembled.toString() to (stopReason == "max_tokens")
        }

        return try {
            val (text, truncated) = streamOnce(prompt)
            if (text.isBlank() || truncated) {
                return analyze(client, baseUrl, model, apiKey, prompt, imageBytesList, maxTokens)
            }
            text
        } catch (e: AiError) {
            throw e
        } catch (_: Throwable) {
            analyze(client, baseUrl, model, apiKey, prompt, imageBytesList, maxTokens)
        }
    }

    suspend fun chat(
        client: OkHttpClient,
        baseUrl: String,
        model: String,
        apiKey: String,
        systemPrompt: String,
        history: List<Pair<String, String>>, // (role: "user"|"assistant", content)
        userMessage: String,
        maxTokens: Int
    ): String {
        val url = "$baseUrl/messages"

        val messages = JSONArray()
        for ((role, content) in history) {
            messages.put(JSONObject().put("role", role).put("content", content))
        }
        messages.put(JSONObject().put("role", "user").put("content", userMessage))

        val body = JSONObject()
            .put("model", model)
            .put("max_tokens", maxTokens)
            .put("system", systemPrompt)
            .put("messages", messages)

        val bodyStr = RetryPolicy.execute {
            client.newCall(
                Request.Builder()
                    .url(url)
                    .addHeader("Content-Type", "application/json")
                    .addHeader("x-api-key", apiKey)
                    .addHeader("anthropic-version", API_VERSION)
                    .post(body.toString().toRequestBody(jsonMedia))
                    .build()
            )
        }
        return AnthropicResponseParser.parse(bodyStr).text ?: throw AiError.InvalidResponse
    }

    private fun compactRetryPrompt(prompt: String, maxTokens: Int): String =
        "$prompt\n\nIMPORTANT: The previous response was truncated. Return only the requested compact JSON object, with no reasoning, explanation, or markdown. Keep the complete response under $maxTokens tokens."
}

internal data class AnthropicTextResponse(
    val text: String?,
    val stopReason: String?,
) {
    val wasTruncated: Boolean get() = stopReason == "max_tokens"
}

internal object AnthropicResponseParser {
    fun parse(body: String): AnthropicTextResponse {
        val json = runCatching { Json.parseToJsonElement(body).jsonObject }.getOrNull()
            ?: throw AiError.InvalidResponse
        val content = runCatching { json["content"]?.jsonArray }.getOrNull()
            ?: throw AiError.InvalidResponse
        val text = content.mapNotNull { element ->
            val block = runCatching { element.jsonObject }.getOrNull() ?: return@mapNotNull null
            if (block["type"]?.jsonPrimitive?.contentOrNull != "text") return@mapNotNull null
            block["text"]?.jsonPrimitive?.contentOrNull?.trim()?.takeIf { it.isNotEmpty() }
        }.joinToString("\n").takeIf { it.isNotEmpty() }
        return AnthropicTextResponse(
            text = text,
            stopReason = json["stop_reason"]?.jsonPrimitive?.contentOrNull,
        )
    }
}
