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, through [highlighted] and one palette, 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: [fenceLanguage] answers null for those, and
 * plain is what the reader would have seen before.
 *
 * 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, and [MarkdownCodeFence]
 * hands out the code and the language and leaves the drawing to the block it is given.
 */
@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.
 *
 * Which children of the node are the fence markers, the language word and the code between them is
 * the library's knowledge of the parser, copied from its `MarkdownCodeFence` rather than called:
 * that one is a composable, and the whole point of this function is that [warm] can run it 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, which the
 * library skips as invalid.
 */
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.
 * The vertical margin is the renderer's too, kept so a reply's fences sit where they always have.
 */
@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;
            // see [MarkdownRoot]'s `streaming`.
            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]

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,
    )

/**
 * 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
}