ai-app/app/androidApp/src/main/kotlin/com/example/aiapp/Api.kt

package com.example.aiapp

import java.io.IOException
import java.net.HttpURLConnection
import java.net.URL
import org.json.JSONArray
import org.json.JSONObject

// The REST half of the backend's surface (see server/src/routes.rs for the
// table); the SSE half is EventStream.kt. All blocking network calls --
// invoke from a background dispatcher. Each throws ApiException on failure,
// carrying the server's own explanation where it sent one, since those
// messages are written to be read on this screen.

// Shared with EventStream.kt, which connects the same way but then reads
// without a deadline.
const val CONNECT_TIMEOUT_MS = 5000
private const val READ_TIMEOUT_MS = 5000

class ApiException(message: String, cause: Throwable? = null) : Exception(message, cause)

/**
 * Runs one request against the backend, with the pinned TLS setup, the bearer token, and the
 * failure translation every call needs. [readBody] gets the connected, already-status-checked
 * connection to read from.
 *
 * @param readTimeoutMs how long to wait on the response body. The SSE stream doesn't come through
 *   here -- an event stream has no bounded read time (see EventStream.kt).
 */
fun <T> requestFromServer(
    settings: ServerSettings,
    path: String,
    method: String = "GET",
    jsonBody: String? = null,
    /** Raw request body as content-type to bytes -- the upload path. */
    binaryBody: Pair<String, ByteArray>? = null,
    readTimeoutMs: Int = READ_TIMEOUT_MS,
    readBody: (HttpURLConnection) -> T,
): T {
    val connection = URL("${settings.baseUrl}$path").openConnection() as HttpURLConnection
    try {
        connection.applyPinnedTls()
        connection.requestMethod = method
        connection.connectTimeout = CONNECT_TIMEOUT_MS
        connection.readTimeout = readTimeoutMs
        connection.setRequestProperty("Authorization", "Bearer ${settings.token}")
        if (jsonBody != null) {
            connection.doOutput = true
            connection.setRequestProperty("Content-Type", "application/json")
            connection.outputStream.use { it.write(jsonBody.encodeToByteArray()) }
        } else if (binaryBody != null) {
            connection.doOutput = true
            connection.setRequestProperty("Content-Type", binaryBody.first)
            connection.outputStream.use { it.write(binaryBody.second) }
        }
        if (connection.responseCode !in 200..299) {
            val detail = connection.errorStream?.bufferedReader()?.readText()?.trim()
            throw ApiException(
                when {
                    connection.responseCode == 401 ->
                        "The server rejected this device's token. Re-enroll by scanning " +
                            "the server's QR (or rotate with --rotate-token and scan the new one)."
                    detail.isNullOrEmpty() ->
                        "Server returned HTTP ${connection.responseCode} for $path"
                    else -> detail
                }
            )
        }
        return readBody(connection)
    } catch (e: ApiException) {
        throw e
    } catch (e: IOException) {
        // Surfacing the real exception (rather than one canned message for
        // every failure mode) is what lets this be diagnosed on a device
        // with no logcat access.
        throw ApiException(
            "Couldn't reach the server at ${settings.baseUrl} " +
                "(${e::class.simpleName}: ${e.message}) -- is ai-server running, and is " +
                "this device able to reach that address (WireGuard up)?",
            e,
        )
    } catch (e: Exception) {
        throw ApiException(
            "Reached ${settings.baseUrl}$path but couldn't read its response " +
                "(${e::class.simpleName}: ${e.message})",
            e,
        )
    } finally {
        connection.disconnect()
    }
}

/** The response body as one JSON object. */
private fun HttpURLConnection.jsonObject(): JSONObject =
    JSONObject(inputStream.bufferedReader().readText())

/** The response body as a JSON array of objects, each mapped through [parse]. */
private fun <T> HttpURLConnection.jsonObjects(parse: (JSONObject) -> T): List<T> =
    JSONArray(inputStream.bufferedReader().readText()).mapObjects(parse)

private fun <T> JSONArray.mapObjects(parse: (JSONObject) -> T): List<T> =
    (0 until length()).map { parse(getJSONObject(it)) }

private fun JSONArray.strings(): List<String> = (0 until length()).map { getString(it) }

// One row of GET /sessions. `provider` is what runs it, `host` where --
// the two are independent, so a session names both.
data class SessionSummary(
    val id: String,
    val provider: String,
    val title: String,
    val host: String?,
    val model: String?,
    val status: String,
    val lastActivity: Double,
)

private fun parseSession(session: JSONObject) =
    SessionSummary(
        id = session.getString("id"),
        provider = session.getString("provider"),
        title = session.getString("title"),
        host = session.optString("host").ifEmpty { null },
        model = session.optString("model").ifEmpty { null },
        status = session.getString("status"),
        lastActivity = session.getDouble("lastActivity"),
    )

fun fetchSessions(settings: ServerSettings): List<SessionSummary> =
    requestFromServer(settings, "/sessions") { it.jsonObjects(::parseSession) }

// What the server offers, so the spawn screen has no hardcoded lists: a
// provider or host added to the server's config.ron appears here with no
// app rebuild.
data class Provider(val name: String, val kind: String, val models: List<String>)

data class RemoteHost(val name: String, val address: String)

fun fetchProviders(settings: ServerSettings): List<Provider> =
    requestFromServer(settings, "/providers") { connection ->
        connection.jsonObjects { provider ->
            Provider(
                name = provider.getString("name"),
                kind = provider.getString("kind"),
                // Omitted entirely when the provider offers none.
                models = provider.optJSONArray("models")?.strings().orEmpty(),
            )
        }
    }

fun fetchHosts(settings: ServerSettings): List<RemoteHost> =
    requestFromServer(settings, "/hosts") { connection ->
        connection.jsonObjects { host ->
            RemoteHost(name = host.getString("name"), address = host.getString("address"))
        }
    }

/**
 * Spawns a session and returns it as the list would show it. [host] is the name of a configured
 * host, or null to run on the backend machine itself.
 */
fun spawnSession(
    settings: ServerSettings,
    provider: String,
    title: String,
    host: String? = null,
    model: String? = null,
    cwd: String? = null,
    permissionMode: String? = null,
): SessionSummary =
    requestFromServer(
        settings,
        "/sessions",
        method = "POST",
        jsonBody =
            JSONObject()
                .put("provider", provider)
                .put("title", title)
                .apply {
                    if (!host.isNullOrBlank()) put("host", host)
                    if (!model.isNullOrBlank()) put("model", model)
                    if (!cwd.isNullOrBlank()) put("cwd", cwd)
                    if (!permissionMode.isNullOrBlank()) put("permissionMode", permissionMode)
                }
                .toString(),
        readTimeoutMs = 30000,
    ) { connection ->
        parseSession(connection.jsonObject())
    }

fun sendMessage(
    settings: ServerSettings,
    sessionId: String,
    text: String,
    attachmentIds: List<String> = emptyList(),
) {
    requestFromServer(
        settings,
        "/sessions/$sessionId/message",
        method = "POST",
        jsonBody =
            JSONObject()
                .put("text", text)
                .put("attachmentIds", JSONArray(attachmentIds))
                .toString(),
    ) {}
}

/** Uploads one picked image; the returned id goes into [sendMessage]. */
fun uploadAttachment(
    settings: ServerSettings,
    sessionId: String,
    bytes: ByteArray,
    mime: String,
): String {
    val boundary = "----aiapp-${System.currentTimeMillis()}"
    val head =
        ("--$boundary\r\n" +
                "Content-Disposition: form-data; name=\"file\"; filename=\"image\"\r\n" +
                "Content-Type: $mime\r\n\r\n")
            .encodeToByteArray()
    val tail = "\r\n--$boundary--\r\n".encodeToByteArray()
    return requestFromServer(
        settings,
        "/sessions/$sessionId/attachments",
        method = "POST",
        binaryBody = "multipart/form-data; boundary=$boundary" to (head + bytes + tail),
        readTimeoutMs = 60000,
    ) { connection ->
        connection.jsonObject().getString("id")
    }
}

/** Fetches an image the transcript references (produced or uploaded). */
fun fetchSessionFile(settings: ServerSettings, sessionId: String, name: String): ByteArray =
    requestFromServer(settings, "/sessions/$sessionId/files/$name", readTimeoutMs = 30000) {
        it.inputStream.readBytes()
    }

// One rate-limit window, rendered as a labeled bar on the usage screen.
data class UsageWindow(
    val label: String,
    val percent: Double,
    val resetsAt: String?,
    val active: Boolean,
)

data class UsageSnapshot(
    val provider: String,
    val available: Boolean,
    val windows: List<UsageWindow>,
    val error: String?,
)

/** The backend caches; refreshing more often than its poll interval just re-reads the cache. */
fun fetchUsage(settings: ServerSettings): List<UsageSnapshot> =
    requestFromServer(settings, "/usage", readTimeoutMs = 30000) { connection ->
        connection.jsonObjects { snapshot ->
            UsageSnapshot(
                provider = snapshot.getString("provider"),
                available = snapshot.getBoolean("available"),
                error = snapshot.optString("error").ifEmpty { null },
                windows =
                    snapshot.getJSONArray("windows").mapObjects { window ->
                        UsageWindow(
                            label = window.getString("label"),
                            percent = window.getDouble("percent"),
                            resetsAt = window.optString("resetsAt").ifEmpty { null },
                            active = window.getBoolean("active"),
                        )
                    },
            )
        }
    }

fun answerQuestion(
    settings: ServerSettings,
    sessionId: String,
    questionId: String,
    answer: String,
) {
    requestFromServer(
        settings,
        "/sessions/$sessionId/answer",
        method = "POST",
        jsonBody = JSONObject().put("questionId", questionId).put("answer", answer).toString(),
    ) {}
}

fun interruptSession(settings: ServerSettings, sessionId: String) {
    requestFromServer(settings, "/sessions/$sessionId/interrupt", method = "POST") {}
}

fun deleteSession(settings: ServerSettings, sessionId: String) {
    requestFromServer(settings, "/sessions/$sessionId", method = "DELETE") {}
}