package com.aryan.reader.desktop import com.aryan.reader.shared.AiAdapter import com.aryan.reader.shared.AiDefinitionResult import com.aryan.reader.shared.RecapResult import com.aryan.reader.shared.SummarizationResult import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.withContext import kotlinx.serialization.json.Json import kotlinx.serialization.json.JsonPrimitive import kotlinx.serialization.json.buildJsonArray import kotlinx.serialization.json.buildJsonObject import kotlinx.serialization.json.contentOrNull import kotlinx.serialization.json.jsonObject import kotlinx.serialization.json.jsonPrimitive import java.io.InputStream import java.net.HttpURLConnection import java.net.URL internal class DesktopPaidAiAdapter( private val config: DesktopCloudConfig, private val networkAccess: () -> Boolean, private val hideReaderAiFeatures: () -> Boolean, private val currentAuthToken: suspend () -> String?, private val currentSignedIn: () -> Boolean, private val currentIsProUser: () -> Boolean, private val currentCredits: () -> Int, private val onUsageCompleted: suspend () -> Unit = {} ) : AiAdapter { override val isAvailable: Boolean get() = networkAccess() && config.isAiWorkerConfigured && !hideReaderAiFeatures() override suspend fun define(text: String, context: String?): AiDefinitionResult { val trimmed = text.trim() if (trimmed.isBlank()) return AiDefinitionResult(error = "There is no text to define.") val multiWord = wordCount(trimmed) > 1 if (multiWord && !currentSignedIn()) { return AiDefinitionResult(error = "Sign in with Google to use multi-word smart dictionary.") } if (multiWord && !currentIsProUser()) { return AiDefinitionResult(error = "Multi-word smart dictionary requires Pro. Pro can only be purchased from the Android app.") } val result = callWorker( path = "/define", body = buildJsonObject { put("text", JsonPrimitive(trimmed.take(2400))) }.toString(), authRequired = multiWord ) return AiDefinitionResult(definition = result.getOrNull()?.text, error = result.exceptionOrNull()?.message) } override suspend fun defineStreaming( text: String, context: String?, onUpdate: (String) -> Unit ): AiDefinitionResult { val trimmed = text.trim() if (trimmed.isBlank()) return AiDefinitionResult(error = "There is no text to define.") val multiWord = wordCount(trimmed) > 1 if (multiWord && !currentSignedIn()) { return AiDefinitionResult(error = "Sign in with Google to use multi-word smart dictionary.") } if (multiWord && !currentIsProUser()) { return AiDefinitionResult(error = "Multi-word smart dictionary requires Pro. Pro can only be purchased from the Android app.") } val result = callWorker( path = "/define", body = buildJsonObject { put("text", JsonPrimitive(trimmed.take(2400))) }.toString(), authRequired = multiWord, onChunk = onUpdate ) return AiDefinitionResult(definition = result.getOrNull()?.text, error = result.exceptionOrNull()?.message) } override suspend fun summarize(text: String): SummarizationResult { val trimmed = text.trim() if (trimmed.isBlank()) return SummarizationResult(error = "There is no text to summarize.") val gate = paidGenerationGate(freeProSummaryAllowed = true) if (gate != null) return SummarizationResult(error = gate) val result = callWorker( path = "/summarize", body = buildJsonObject { put("content_type", JsonPrimitive("text")) put("data", JsonPrimitive(trimmed)) }.toString(), authRequired = true ) val response = result.getOrNull() return SummarizationResult( summary = response?.text, error = result.exceptionOrNull()?.message, cost = response?.cost, freeRemaining = response?.freeRemaining ) } override suspend fun summarizeStreaming( text: String, onUsageReceived: (cost: Double?, freeRemaining: Int?) -> Unit, onUpdate: (String) -> Unit ): SummarizationResult { val trimmed = text.trim() if (trimmed.isBlank()) return SummarizationResult(error = "There is no text to summarize.") val gate = paidGenerationGate(freeProSummaryAllowed = true) if (gate != null) return SummarizationResult(error = gate) val result = callWorker( path = "/summarize", body = buildJsonObject { put("content_type", JsonPrimitive("text")) put("data", JsonPrimitive(trimmed)) }.toString(), authRequired = true, onChunk = onUpdate, onUsageReceived = onUsageReceived ) val response = result.getOrNull() return SummarizationResult( summary = response?.text, error = result.exceptionOrNull()?.message, cost = response?.cost, freeRemaining = response?.freeRemaining ) } override suspend fun recap(textBeforeCurrentLocation: String): RecapResult { return recapWithContext(emptyList(), textBeforeCurrentLocation) } suspend fun recapWithContext(pastSummaries: List, currentText: String): RecapResult { val trimmed = currentText.trim() if (trimmed.isBlank()) return RecapResult(error = "There is no reading context for a recap.") val gate = paidGenerationGate(freeProSummaryAllowed = false) if (gate != null) return RecapResult(error = gate) val result = callWorker( path = "/recap", body = buildJsonObject { put( "past_summaries", buildJsonArray { pastSummaries.filter { it.isNotBlank() }.forEach { summary -> add(JsonPrimitive(summary)) } } ) put("current_text", JsonPrimitive(trimmed)) }.toString(), authRequired = true ) val response = result.getOrNull() return RecapResult( recap = response?.text, error = result.exceptionOrNull()?.message, cost = response?.cost, freeRemaining = response?.freeRemaining ) } private fun paidGenerationGate(freeProSummaryAllowed: Boolean): String? { if (!config.isAiWorkerConfigured) return "Desktop AI is not configured." if (!networkAccess()) return "AI features are unavailable in this desktop build." if (hideReaderAiFeatures()) return "Reader AI features are hidden." if (!currentSignedIn()) return "Sign in with Google to use this AI feature." if (!(freeProSummaryAllowed && currentIsProUser()) && currentCredits() <= 0) { return "This action needs credits. Pro and credits can only be purchased from the Android app." } return null } private suspend fun callWorker( path: String, body: String, authRequired: Boolean, onChunk: (String) -> Unit = {}, onUsageReceived: (cost: Double?, freeRemaining: Int?) -> Unit = { _, _ -> } ): Result = withContext(Dispatchers.IO) { if (!isAvailable) return@withContext Result.failure(IllegalStateException("AI features are unavailable.")) val token = currentAuthToken() if (authRequired && token.isNullOrBlank()) { return@withContext Result.failure(IllegalStateException("Sign in with Google to use this AI feature.")) } runCatching { val url = URL(config.aiWorkerUrl.removeSuffix("/") + path) val connection = (url.openConnection() as HttpURLConnection).apply { requestMethod = "POST" setRequestProperty("Content-Type", "application/json; charset=UTF-8") setRequestProperty("Accept", "application/json") if (!token.isNullOrBlank()) setRequestProperty("Authorization", "Bearer $token") connectTimeout = 15_000 readTimeout = 120_000 doOutput = true doInput = true } try { connection.outputStream.use { it.write(body.toByteArray(Charsets.UTF_8)) } val responseCode = connection.responseCode val stream = if (responseCode in 200..299) connection.inputStream else connection.errorStream if (responseCode in 200..299) { val parsed = readWorkerStream(stream, onChunk, onUsageReceived) if (parsed.text.isBlank()) throw IllegalStateException("The AI service returned an empty response.") onUsageCompleted() return@runCatching parsed } val responseText = stream?.bufferedReader(Charsets.UTF_8)?.use { it.readText() }.orEmpty() if (connection.responseCode == 402 || responseText.contains("INSUFFICIENT_CREDITS")) { throw IllegalStateException("Out of credits. Pro and credits can only be purchased from the Android app.") } if (connection.responseCode == 401) { throw IllegalStateException("Sign in again to use this AI feature.") } if (connection.responseCode == 403 && responseText.contains("MULTI_WORD_REQUIRES_PRO")) { throw IllegalStateException("Multi-word smart dictionary requires Pro. Pro can only be purchased from the Android app.") } if (connection.responseCode !in 200..299) { throw IllegalStateException(workerErrorMessage(responseText) ?: "AI request failed: HTTP ${connection.responseCode}") } error("Unreachable AI response state.") } finally { connection.disconnect() } } } } private data class DesktopPaidAiResponse( val text: String, val cost: Double? = null, val freeRemaining: Int? = null ) private val DesktopPaidAiJson = Json { ignoreUnknownKeys = true } private fun readWorkerStream( stream: InputStream?, onChunk: (String) -> Unit, onUsageReceived: (cost: Double?, freeRemaining: Int?) -> Unit ): DesktopPaidAiResponse { val output = StringBuilder() var cost: Double? = null var freeRemaining: Int? = null stream?.bufferedReader(Charsets.UTF_8)?.useLines { lines -> lines.forEach { line -> val parsed = parseWorkerStreamLine(line) ?: return@forEach parsed.cost?.let { cost = it } parsed.freeRemaining?.let { freeRemaining = it } if (parsed.cost != null || parsed.freeRemaining != null) { onUsageReceived(parsed.cost, parsed.freeRemaining) } parsed.chunk?.let { chunk -> output.append(chunk) onChunk(chunk) } } } return DesktopPaidAiResponse(text = output.toString().trim(), cost = cost, freeRemaining = freeRemaining) } private fun parseWorkerStreamLine(line: String): DesktopPaidAiStreamLine? { val trimmed = line.trim() if (trimmed.isBlank()) return null val parsed = runCatching { DesktopPaidAiJson.parseToJsonElement(trimmed).jsonObject }.getOrNull() ?: return null parsed.get("error")?.jsonPrimitive?.contentOrNull?.takeIf { it.isNotBlank() }?.let { error -> throw IllegalStateException(workerErrorMessage(error) ?: error) } return DesktopPaidAiStreamLine( chunk = parsed.get("chunk")?.jsonPrimitive?.contentOrNull, cost = parsed.get("cost_deducted")?.jsonPrimitive?.contentOrNull?.toDoubleOrNull(), freeRemaining = parsed.get("free_summaries_remaining")?.jsonPrimitive?.contentOrNull?.toIntOrNull() ) } private data class DesktopPaidAiStreamLine( val chunk: String? = null, val cost: Double? = null, val freeRemaining: Int? = null ) private fun workerErrorMessage(errorBody: String): String? { return when { errorBody.contains("INSUFFICIENT_CREDITS") -> "Out of credits. Pro and credits can only be purchased from the Android app." errorBody.contains("SUMMARY_LIMIT") || (errorBody.contains("free summar", ignoreCase = true) && errorBody.contains("limit", ignoreCase = true)) -> "Free summaries are used up for today. More summaries need credits, and credits can only be purchased from the Android app." errorBody.contains("MULTI_WORD_REQUIRES_PRO") -> "Multi-word smart dictionary requires Pro. Pro can only be purchased from the Android app." errorBody.contains("Authentication required") -> "Sign in with Google to use this AI feature." else -> runCatching { DesktopPaidAiJson.parseToJsonElement(errorBody) .jsonObject["error"] ?.jsonPrimitive ?.contentOrNull }.getOrNull() } } private fun wordCount(text: String): Int { return text.trim().split(Regex("\\s+")).count { it.isNotBlank() } }