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

package com.example.aiapp

import androidx.compose.foundation.horizontalScroll
import androidx.compose.foundation.layout.Column
import androidx.compose.foundation.layout.fillMaxWidth
import androidx.compose.foundation.layout.padding
import androidx.compose.foundation.rememberScrollState
import androidx.compose.material3.MaterialTheme
import androidx.compose.material3.Text
import androidx.compose.runtime.Composable
import androidx.compose.runtime.remember
import androidx.compose.ui.Modifier
import androidx.compose.ui.text.AnnotatedString
import androidx.compose.ui.text.SpanStyle
import androidx.compose.ui.text.buildAnnotatedString
import androidx.compose.ui.text.font.FontFamily
import androidx.compose.ui.text.font.FontWeight
import androidx.compose.ui.unit.dp
import dev.snipme.highlights.Highlights
import dev.snipme.highlights.model.BoldHighlight
import dev.snipme.highlights.model.ColorHighlight
import dev.snipme.highlights.model.SyntaxLanguage
import org.json.JSONObject

/**
 * A tool call's input, read rather than dumped.
 *
 * Every tool's input arrives as JSON, and showing it raw makes the reader parse `{"command":"…",
 * "timeout":120000}` themselves to find the one line they care about. So the fields that carry the
 * meaning are pulled out -- the command a shell will run, what it is for, how long it may take --
 * and anything left over is still shown, because dropping a field would be claiming the tool has no
 * other input when it might.
 */
data class ToolInput(
    /** The thing that will actually be run or read, if this tool has one. */
    val subject: String?,
    /** The language [subject] is written in, for highlighting. */
    val language: SyntaxLanguage?,
    /** The tool's own one-line summary, when it wrote one. */
    val description: String?,
    /**
     * How long the call may take, as the tool expressed it. Shown apart because it is a limit on
     * the call rather than part of what the call does.
     */
    val timeout: String?,
    /** Everything else, as `name: value` lines. Never dropped. */
    val rest: List<String>,
) {
    /** The one line to show when there is only room for one: what this call is for. */
    val title: String?
        get() = description ?: subject
}

/**
 * Which field of which tool is the subject.
 *
 * A table rather than a chain of `if`s: adding a tool is a row, and the shape stops any of them
 * from being the special case that gets its own code path. Unknown tools fall through to "no
 * subject, everything is rest", which is what the card always did.
 */
private val SUBJECTS: Map<String, Pair<String, SyntaxLanguage?>> =
    mapOf(
        "Bash" to ("command" to SyntaxLanguage.SHELL),
        "Read" to ("file_path" to null),
        "Write" to ("file_path" to null),
        "Edit" to ("file_path" to null),
        "Glob" to ("pattern" to null),
        "Grep" to ("pattern" to null),
        "WebFetch" to ("url" to null),
    )

/** Fields that are the tool's own prose about itself rather than input to it. */
private val DESCRIPTIONS = listOf("description", "prompt")

fun parseToolInput(tool: String, input: String): ToolInput {
    val json =
        try {
            JSONObject(input)
        } catch (_: org.json.JSONException) {
            // Not an object: older transcripts and some tools send a bare
            // string. It is still the input, so it is still shown.
            return ToolInput(
                null,
                null,
                null,
                null,
                input.takeIf { it.isNotBlank() }?.let { listOf(it) }.orEmpty(),
            )
        }
    val (subjectKey, language) = SUBJECTS[tool] ?: (null to null)
    val subject = subjectKey?.let { json.optString(it) }?.takeIf { it.isNotBlank() }
    val description = DESCRIPTIONS.firstNotNullOfOrNull {
        json.optString(it).takeIf { v -> v.isNotBlank() }
    }
    val timeout = json.optString("timeout").takeIf { it.isNotBlank() }
    val rest =
        json
            .keys()
            .asSequence()
            .filter { it != subjectKey || subject == null }
            .filter { it !in DESCRIPTIONS || description == null }
            .filter { it != "timeout" || timeout == null }
            .sorted()
            .map { key -> "$key: ${json.get(key)}" }
            .toList()
    return ToolInput(subject, language, description, timeout, rest)
}

/** A tool call's input: its subject highlighted, its description, then whatever else it carried. */
@Composable
fun ToolInputView(tool: String, input: String, modifier: Modifier = Modifier) {
    val parsed = remember(tool, input) { parseToolInput(tool, input) }
    Column(modifier.fillMaxWidth()) {
        parsed.subject?.let { subject ->
            // Not wrapped: a wrapped command hides where its arguments end,
            // and the long one is the one being read closely.
            Text(
                highlighted(subject, parsed.language),
                style = MaterialTheme.typography.bodySmall,
                fontFamily = FontFamily.Monospace,
                softWrap = false,
                modifier =
                    Modifier.padding(top = 4.dp)
                        .fillMaxWidth()
                        .horizontalScroll(rememberScrollState()),
            )
        }
        parsed.rest.forEach {
            Text(
                it,
                style = MaterialTheme.typography.bodySmall,
                fontFamily = FontFamily.Monospace,
                color = MaterialTheme.colorScheme.onSurfaceVariant,
                modifier = Modifier.padding(top = 2.dp),
            )
        }
    }
}

/**
 * [code] with its keywords and strings coloured, or plain if there is no language for it.
 *
 * The lexing is dev.snipme:highlights. The colours are this app's, mapped in [catppuccinSyntax] --
 * a library's default theme would be the one place in the app whose palette came from somewhere
 * else.
 */
@Composable
private fun highlighted(code: String, language: SyntaxLanguage?): AnnotatedString {
    val theme = catppuccinSyntax()
    val plain = MaterialTheme.colorScheme.onSurface
    return remember(code, language, theme, plain) {
        if (language == null) return@remember AnnotatedString(code)
        val marks =
            Highlights.Builder(code = code, language = language, theme = theme)
                .build()
                .getHighlights()
        buildAnnotatedString {
            append(code)
            marks.forEach { mark ->
                when (mark) {
                    is ColorHighlight ->
                        addStyle(
                            SpanStyle(
                                color =
                                    androidx.compose.ui.graphics.Color(
                                        mark.rgb or 0xFF000000.toInt()
                                    )
                            ),
                            mark.location.start,
                            mark.location.end,
                        )
                    is BoldHighlight ->
                        addStyle(
                            SpanStyle(fontWeight = FontWeight.Bold),
                            mark.location.start,
                            mark.location.end,
                        )
                }
            }
        }
    }
}