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

package com.example.aiapp

import androidx.compose.ui.graphics.Color
import androidx.compose.ui.text.AnnotatedString
import androidx.compose.ui.text.SpanStyle
import androidx.compose.ui.text.buildAnnotatedString

/** What a span of code is, in the terms the palette has a colour for. */
enum class Kind {
    KEYWORD,
    STRING,
    LITERAL,
    COMMENT,
    METADATA,
    PUNCTUATION,
    MARK,
}

/** A run of [Kind] in the code, as a half-open range. */
data class Span(val start: Int, val end: Int, val kind: Kind)

/**
 * The colours the highlighter draws with, ours rather than a library's; [catppuccinSyntax] is the
 * one instance and lives with the rest of the palette.
 */
data class SyntaxPalette(
    val keyword: Color,
    val string: Color,
    val literal: Color,
    val comment: Color,
    val metadata: Color,
    val punctuation: Color,
    val mark: Color,
) {
    fun of(kind: Kind): Color =
        when (kind) {
            Kind.KEYWORD -> keyword
            Kind.STRING -> string
            Kind.LITERAL -> literal
            Kind.COMMENT -> comment
            Kind.METADATA -> metadata
            Kind.PUNCTUATION -> punctuation
            Kind.MARK -> mark
        }
}

/**
 * [code] with its keywords, strings and comments coloured, or plain if there is no language for it.
 *
 * Shared by a tool call's input and a reply's fences, so the same code is the same colours wherever
 * it appears.
 *
 * Not a composable, and it takes no colour from the theme, because that is what lets [warm] run it
 * off the drawing thread.
 *
 * The timing is the number the highlighter is judged by: the library this replaced took **174ms**
 * on the emulator for a two-hundred-line Kotlin fence, which is why [ParsedReplies.highlighted]
 * caches the answer rather than a `remember` inside the fence recomputing it on every scroll back.
 */
fun highlight(code: String, language: Language?): AnnotatedString {
    if (language == null) return AnnotatedString(code)
    val spans = DebugStats.timed("code highlighted") { spansOf(code, language) }
    val palette = catppuccinSyntax()
    return buildAnnotatedString {
        append(code)
        spans.forEach { addStyle(SpanStyle(color = palette.of(it.kind)), it.start, it.end) }
    }
}

/**
 * [code] read once, left to right, into the spans that carry a colour.
 *
 * One pass with a small state -- in a comment, in a string, or in ordinary code -- rather than a
 * locator per token kind over the whole text, which is what the library did and is why it found
 * comments before it knew the language: a `#` inside a shell string, a `//` inside a URL and a
 * block-comment opener inside a shell glob each commented out the rest of a line that was nothing
 * of the sort.
 *
 * Every span is produced by advancing an index forward, so the result is ordered, non-overlapping
 * and inside the code by construction. Nothing here throws: an unterminated string or comment runs
 * to the end of the code, which is also what it looks like while a fence is still being written.
 *
 * In ordinary code the order of recognition is comment, string, attribute, number, word, and
 * finally a single punctuation or mark character, which are coloured only in ordinary code.
 */
fun scan(code: String, rules: Rules): List<Span> = Scanner(code, rules).run()

/** Characters coloured as punctuation, and as marks. Both sets are the ones the library used. */
private const val PUNCTUATION = ",.:;"
private const val MARKS = "()={}<>-+[]|&"

private class Scanner(private val code: String, private val rules: Rules) {
    private val spans = ArrayList<Span>()
    private var at = 0

    fun run(): List<Span> {
        while (at < code.length) {
            // Every branch that answers true has advanced `at`, so this terminates.
            val consumed =
                blockComment() ||
                    lineComment() ||
                    rawString() ||
                    characterOrLifetime() ||
                    string() ||
                    attribute() ||
                    number() ||
                    word() ||
                    singleCharacter()
            if (!consumed) at++
        }
        return spans
    }

    private fun emit(start: Int, kind: Kind) {
        if (at > start) spans.add(Span(start, at, kind))
    }

    private fun starts(token: String) = code.startsWith(token, at)

    /** Whether a line comment token here opens one; see [Rules.lineCommentsAtWordStart]. */
    private fun atWordStart() = at == 0 || code[at - 1].isWhitespace() || code[at - 1] in ";|&("

    /** Whether only whitespace stands between the start of this line and here. */
    private fun atLineStart(): Boolean {
        var back = at - 1
        while (back >= 0 && code[back] != '\n') {
            if (!code[back].isWhitespace()) return false
            back--
        }
        return true
    }

    private fun toEndOfLine() {
        while (at < code.length && code[at] != '\n') at++
    }

    /** From an open bracket through the one that matches it, or to the end if none does. */
    private fun toMatchingBracket() {
        var depth = 0
        while (at < code.length) {
            when (code[at]) {
                '[' -> depth++
                ']' -> depth--
            }
            at++
            if (depth == 0) return
        }
    }

    private fun blockComment(): Boolean {
        val comment = rules.blockComment ?: return false
        if (!starts(comment.open)) return false
        val start = at
        at += comment.open.length
        var depth = 1
        while (at < code.length && depth > 0) {
            // The closer is tried first so that a language whose two delimiters are the same string
            // -- CoffeeScript's `###` -- closes rather than nesting forever.
            if (starts(comment.close)) {
                depth--
                at += comment.close.length
            } else if (comment.nests && starts(comment.open)) {
                depth++
                at += comment.open.length
            } else {
                at++
            }
        }
        emit(start, Kind.COMMENT)
        return true
    }

    private fun lineComment(): Boolean {
        if (rules.lineComments.none { starts(it) }) return false
        if (rules.lineCommentsAtWordStart && !atWordStart()) return false
        val start = at
        toEndOfLine()
        emit(start, Kind.COMMENT)
        return true
    }

    /** Rust and RON: `b`? `r` `#`* `"` … `"` `#`*, with no escapes inside. */
    private fun rawString(): Boolean {
        if (!rules.rawStrings) return false
        var ahead = at
        if (code.getOrNull(ahead) == 'b') ahead++
        if (code.getOrNull(ahead) != 'r') return false
        ahead++
        var hashes = 0
        while (code.getOrNull(ahead) == '#') {
            ahead++
            hashes++
        }
        if (code.getOrNull(ahead) != '"') return false
        val start = at
        val closer = "\"" + "#".repeat(hashes)
        val closed = code.indexOf(closer, ahead + 1)
        at = if (closed < 0) code.length else closed + closer.length
        emit(start, Kind.STRING)
        return true
    }

    /** See [Rules.lifetimes]: an apostrophe that is not a character literal opens nothing. */
    private fun characterOrLifetime(): Boolean {
        if (!rules.lifetimes || code[at] != '\'') return false
        val next = code.getOrNull(at + 1) ?: return false
        if (next == '\\' || code.getOrNull(at + 2) == '\'') {
            quoted(Quote("'", "'", escapes = true))
        } else {
            at++
        }
        return true
    }

    private fun string(): Boolean {
        // Longest opener wins, so Kotlin's `"""` is one delimiter rather than an empty string
        // followed by a quote. A loop rather than filter/maxBy: this runs at every character of
        // ordinary code, and the pair of lists that would allocate is the whole cost of the scan.
        var quote: Quote? = null
        for (candidate in rules.quotes) {
            if (starts(candidate.open) && candidate.open.length > (quote?.open?.length ?: 0)) {
                quote = candidate
            }
        }
        quoted(quote ?: return false)
        return true
    }

    private fun quoted(quote: Quote) {
        val start = at
        at += quote.open.length
        while (at < code.length) {
            if (quote.escapes && code[at] == '\\' && at + 1 < code.length) {
                at += 2
                continue
            }
            if (starts(quote.close)) {
                at += quote.close.length
                break
            }
            at++
        }
        at = at.coerceAtMost(code.length)
        emit(start, Kind.STRING)
    }

    private fun attribute(): Boolean {
        val start = at
        when (rules.attributes) {
            Attributes.NONE -> return false
            Attributes.AT_WORD -> {
                if (code[at] != '@' || !isWordStart(code.getOrNull(at + 1))) return false
                at++
                while (at < code.length && isWordPart(code[at])) at++
            }
            Attributes.HASH_BRACKET -> {
                if (code[at] != '#') return false
                var ahead = at + 1
                if (code.getOrNull(ahead) == '!') ahead++
                if (code.getOrNull(ahead) != '[') return false
                at = ahead
                toMatchingBracket()
            }
            Attributes.HASH_LINE -> {
                if (code[at] != '#' || !atLineStart()) return false
                toEndOfLine()
            }
            Attributes.LINE_BRACKET -> {
                if (code[at] != '[' || !atLineStart()) return false
                toMatchingBracket()
            }
        }
        emit(start, Kind.METADATA)
        return true
    }

    /**
     * A number is a run starting with a digit and carrying on through letters, digits, `_` and `.`
     * -- which covers `0xFF`, `1_000`, `1u32` and `3.14` without a grammar for any of them.
     */
    private fun number(): Boolean {
        if (!code[at].isDigit()) return false
        val start = at
        while (
            at < code.length && (code[at].isLetterOrDigit() || code[at] == '_' || code[at] == '.')
        ) {
            at++
        }
        emit(start, Kind.LITERAL)
        return true
    }

    private fun word(): Boolean {
        if (!isWordStart(code[at])) return false
        val start = at
        while (at < code.length && isWordPart(code[at])) at++
        if (code.substring(start, at) in rules.keywords) emit(start, Kind.KEYWORD)
        return true
    }

    private fun singleCharacter(): Boolean {
        val kind =
            when (code[at]) {
                in PUNCTUATION -> Kind.PUNCTUATION
                in MARKS -> Kind.MARK
                else -> return false
            }
        at++
        emit(at - 1, kind)
        return true
    }
}

private fun isWordStart(c: Char?) = c != null && (c.isLetter() || c == '_')

private fun isWordPart(c: Char) = c.isLetterOrDigit() || c == '_'