Jetpack Compose - trying to mimic Intellij's code design into a Text block - string

Question

I have a UI that I am creating that replicated Android Studio's (IntelliJ) design system for text as code. What I am trying to achieve is a text that changes it's color multiple times when certain words or patterns happen, such as Kotlin keywords (orange text), extension methods(yellowish text) and variables (purple text).
What do I mean? the following functions text design is what I want to achieve, and they are examples of using all 3 colors -
So I tried to implement it using a buildAnnotatedString { } block and things got A) very messy and B) didn't even reach my goal. Here is my current code of trying to implement keywords and extensions -
// A dummy function that represents only the text part of my entire screen, everything else is irrelevant for the question
#Composable
fun TextAsIntellij(solution : String) {
Text(text = buildAnnotatedString {
val kotlinKeywords = mutableListOf("fun", "val", "var", "if", "return", "null", ",", "class", "while", "break", "continue")
val kotlinExtensions = mutableListOf("forEachIndexed") // will add more as needed
solution.split(" ", ".").forEach { string ->
if (kotlinKeywords.contains(string)) {
withStyle(
style = SpanStyle(
color = Orange
)
) {
append(string)
append(' ')
}
return#forEach
} else if (kotlinExtensions.contains(string)) {
withStyle(
style = SpanStyle(
color = Yellow
)
) {
append(string)
append(" ")
}
return#forEach
}
append(string)
append(' ')
}
})
}
The results in the UI are the following -
As you can see, in the first question I have a problem that before an extension I put a redundant space because I used split() method to identify all words by spaces and dots. For the 2nd picture, I don't yet have a proper solution for identifying varibles from an "object", because I have already split my String using spaces and dots, so diving deeper into another split would make my code really inefficient and I though better can be done. 😁
Hopefully someone has a good idea...any Jetbrains spy in the crowd? ;)