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

package com.example.aiapp

import androidx.compose.foundation.background
import androidx.compose.foundation.horizontalScroll
import androidx.compose.foundation.layout.Box
import androidx.compose.foundation.layout.fillMaxWidth
import androidx.compose.foundation.layout.padding
import androidx.compose.foundation.rememberScrollState
import androidx.compose.foundation.shape.RoundedCornerShape
import androidx.compose.foundation.text.BasicText
import androidx.compose.runtime.Composable
import androidx.compose.runtime.remember
import androidx.compose.ui.Modifier
import androidx.compose.ui.semantics.isTraversalGroup
import androidx.compose.ui.semantics.semantics
import androidx.compose.ui.text.TextStyle
import androidx.compose.ui.unit.dp
import com.mikepenz.markdown.compose.LocalMarkdownColors
import com.mikepenz.markdown.compose.LocalMarkdownDimens
import com.mikepenz.markdown.compose.LocalMarkdownPadding
import com.mikepenz.markdown.model.State
import org.intellij.markdown.MarkdownElementTypes
import org.intellij.markdown.MarkdownTokenTypes
import org.intellij.markdown.ast.ASTNode
import org.intellij.markdown.ast.findChildOfType
import org.intellij.markdown.ast.getTextInNode

/**
 * A fenced code block in a reply: the code highlighted, on the dark surface every verbatim thing
 * sits on, scrolling sideways rather than wrapping.
 *
 * The renderer's own fence drew the same block in plain text. The scanner that colours a tool
 * call's command colours a reply's code the same way, so a `kotlin` fence and the Kotlin a tool
 * wrote are the same colours. A fence in a language [scan] has no rules for is plain rather than
 * wrongly coloured.
 *
 * Finding the code is still the library's: which children of the node are the fence markers, the
 * language word and the code between them is its knowledge of the parser.
 */
@Composable
fun CodeFence(
    content: String,
    node: ASTNode,
    style: TextStyle,
    replies: ParsedReplies,
    streaming: Boolean = false,
) {
    val (code, language) = remember(content, node) { fenceContent(content, node) } ?: return
    CodeBlockText(code, language, style, replies, streaming)
}

/** An indented code block, which is a fence with no language word. */
@Composable
fun CodeBlock(
    content: String,
    node: ASTNode,
    style: TextStyle,
    replies: ParsedReplies,
    streaming: Boolean = false,
) {
    val (code, language) = remember(content, node) { fenceContent(content, node) } ?: return
    CodeBlockText(code, language, style, replies, streaming)
}

/**
 * The code inside a fence or indented block, and the highlighter's language for its info word.
 *
 * Copied from the library's `MarkdownCodeFence` rather than called: that one is a composable, and
 * the whole point here is that [warm] can run this on a background thread and highlight the same
 * string the drawing will ask for. Two extractions would be two keys, and the warmed answer would
 * be silently missed at every fence.
 *
 * Null for a fence too short to hold anything -- an unterminated one still arriving.
 */
fun fenceContent(content: String, node: ASTNode): Pair<String, Language?>? {
    val word =
        node.findChildOfType(MarkdownTokenTypes.FENCE_LANG)?.getTextInNode(content)?.toString()
    val language = fenceLanguage(word)
    if (node.type == MarkdownElementTypes.CODE_BLOCK) {
        val start = node.children.firstOrNull()?.startOffset ?: return null
        val end = node.children.lastOrNull()?.endOffset ?: return null
        return content.substring(start, end).replaceIndent() to language
    }
    if (node.children.size < 3) return null
    val start = node.children[2].startOffset
    val fenceCount = if (word != null && node.children.size > 3) 3 else 2
    val end = node.children[(node.children.size - 2).coerceAtLeast(fenceCount)].endOffset
    return content.substring(start, end).replaceIndent() to language
}

/**
 * Plain while [streaming], coloured once the block is finished; see [MarkdownRoot].
 *
 * The renderer's own block, less what nothing here needs: the same background, corner, padding and
 * sideways scroll, without the shadow, the border and the empty pointer handler it also carried.
 */
@Composable
private fun CodeBlockText(
    code: String,
    language: Language?,
    style: TextStyle,
    replies: ParsedReplies,
    streaming: Boolean,
) {
    val colors = LocalMarkdownColors.current
    val dimens = LocalMarkdownDimens.current
    val padding = LocalMarkdownPadding.current
    Box(
        Modifier.fillMaxWidth()
            .padding(vertical = 8.dp)
            .background(colors.codeBackground, RoundedCornerShape(dimens.codeBackgroundCornerSize))
            .semantics { isTraversalGroup = true }
    ) {
        BasicText(
            // No language while the block is still being written, which is what draws it plain.
            replies.highlighted(code, language.takeUnless { streaming }),
            style = style,
            modifier = Modifier.horizontalScroll(rememberScrollState()).padding(padding.codeBlock),
        )
    }
}

/**
 * The highlighter's language for a fence's info word, or null for one it has no rules for.
 *
 * The aliases are what people actually write after the backticks: the file extension as often as
 * the name. A word not here gets no colour rather than the nearest language's, because a fence
 * coloured by the wrong language's rules looks highlighted and is wrong in a way the reader cannot
 * see.
 */
fun fenceLanguage(name: String?): Language? =
    FENCE_LANGUAGES[name?.trim()?.lowercase() ?: return null]

/**
 * The highlighter's language for a *file*, from its name.
 *
 * The same table [fenceLanguage] reads, deliberately: it already keys on the extensions people
 * write after the backticks. One table rather than two, so a language added for fences is a
 * language added for files and neither can be the one somebody forgot.
 *
 * The extension is the part after the *last* dot, which is what makes `build.gradle.kts` Kotlin. A
 * leading dot is not one: `.bashrc` has no extension, it has a name that starts with a dot. A name
 * with no dot at all -- `Makefile` -- is likewise null, and null is drawn plain.
 */
fun fileLanguage(name: String): Language? {
    val dot = name.lastIndexOf('.')
    if (dot < 1) return null
    return fenceLanguage(name.substring(dot + 1))
}

private val FENCE_LANGUAGES: Map<String, Language> =
    mapOf(
        "kotlin" to Language.KOTLIN,
        "kt" to Language.KOTLIN,
        "kts" to Language.KOTLIN,
        "rust" to Language.RUST,
        "rs" to Language.RUST,
        "sh" to Language.SHELL,
        "bash" to Language.SHELL,
        "shell" to Language.SHELL,
        "zsh" to Language.SHELL,
        "console" to Language.SHELL,
        "python" to Language.PYTHON,
        "py" to Language.PYTHON,
        "javascript" to Language.JAVASCRIPT,
        "js" to Language.JAVASCRIPT,
        "jsx" to Language.JAVASCRIPT,
        "typescript" to Language.TYPESCRIPT,
        "ts" to Language.TYPESCRIPT,
        "tsx" to Language.TYPESCRIPT,
        "java" to Language.JAVA,
        "c" to Language.C,
        "h" to Language.C,
        "cpp" to Language.CPP,
        "c++" to Language.CPP,
        "cc" to Language.CPP,
        "hpp" to Language.CPP,
        "csharp" to Language.CSHARP,
        "cs" to Language.CSHARP,
        "c#" to Language.CSHARP,
        "go" to Language.GO,
        "golang" to Language.GO,
        "swift" to Language.SWIFT,
        "dart" to Language.DART,
        "ruby" to Language.RUBY,
        "rb" to Language.RUBY,
        "php" to Language.PHP,
        "perl" to Language.PERL,
        "pl" to Language.PERL,
        "coffeescript" to Language.COFFEESCRIPT,
        "coffee" to Language.COFFEESCRIPT,
        "ron" to Language.RON,
        "toml" to Language.TOML,
        "fish" to Language.FISH,
        "json" to Language.JSON,
        "markdown" to Language.MARKDOWN,
        "md" to Language.MARKDOWN,
    )

/**
 * Every fence in [parse], as the code and language [highlight] will be asked for. Walks the whole
 * tree rather than the top level: a fence inside a list item or a quote is drawn the same way and
 * costs the same to lex.
 */
fun fences(parse: State): List<Pair<String, Language?>> {
    val success = parse as? State.Success ?: return emptyList()
    val out = ArrayList<Pair<String, Language?>>()
    fun walk(node: ASTNode) {
        if (
            node.type == MarkdownElementTypes.CODE_FENCE ||
                node.type == MarkdownElementTypes.CODE_BLOCK
        ) {
            fenceContent(success.content, node)?.let { if (it.second != null) out += it }
            return
        }
        node.children.forEach(::walk)
    }
    walk(success.node)
    return out
}