KotlinSourceIndex.kt

    fun startBackgroundSync() {
        thread(isDaemon = true, name = "KotlinSourceIndex") {
            Thread.currentThread().priority = Thread.MIN_PRIORITY
            runSync()
        }
    }

    private fun runSync() {
        val root = Paths.get(System.getProperty("user.dir")).normalize()
        val kitRoots = System.getProperty("afield.kit.sources").orEmpty().split(java.io.File.pathSeparator)
            .filter { it.isNotBlank() }.map { Paths.get(it) }
        val sourceRoots = (listOf(
            root.resolve("./src/jvmMain/kotlin"),
            root.resolve("./src/personal/kotlin"),
            root.resolve("./src/generatedLinalg/kotlin"),
        ) + kitRoots).distinct().filter { Files.exists(it) }

        if (sourceRoots.isEmpty()) {
            lprintln("warning,sourceindex", "no source roots found under ${root.toAbsolutePath()}")
            return
        }

        val previous = loadIndex()
        val files = mutableListOf<KotlinSourceFileRecord>()
        val definitions = mutableListOf<KotlinDefinitionRecord>()

        val ktFiles = collectKotlinFiles(sourceRoots)
        ktFiles.forEach { file ->
            val rel = root.relativize(file).toString().replace("\\", "/")
            val bytes = runCatching { Files.readAllBytes(file) }.getOrElse { return@forEach }
            val rawText = runCatching { bytes.toString(StandardCharsets.UTF_8) }.getOrDefault("")
            val isKotlinFile = rel.endsWith(".kt", ignoreCase = true) || rel.endsWith(".kts", ignoreCase = true)
            val indexedText = if (isKotlinFile) {
                buildString {
                    appendLine("// @language=kotlin")
                    appendLine("// @file $rel")
                    appendLine()
                    append(rawText)
                }
            } else {
                rawText
            }
            val indexedBytes = indexedText.toByteArray(StandardCharsets.UTF_8)
            val contentHash = Sha256.hex(indexedBytes)
            val modified = runCatching { Files.getLastModifiedTime(file).toMillis() }.getOrDefault(0L)
            files += KotlinSourceFileRecord(
                path = rel,
                contentHash = contentHash,
                fileId = null,
                bytes = indexedBytes.size.toLong(),
                modifiedEpochMs = modified
            )

            parseTopLevelDefinitions(rawText).forEach { def ->
                definitions += KotlinDefinitionRecord(
                    name = def.name,
                    kind = def.kind,
                    filePath = rel,
                    line = def.line,
                    fileId = null,
                    contentHash = contentHash,
                    kdoc = def.kdoc,
                    completionCategory = def.completionCategory
                )
            }
        }

        val index = KotlinSourceIndexDb(
            root = ".",
            apiKeyHash = null,
            updatedAtEpochMs = System.currentTimeMillis(),
            files = files,
            definitions = definitions,
            vectorStoreId = null,
            vectorStoreFileHash = null
        )

        if (previous != null && previous.materiallySameAs(index)) {
            latestIndex = previous
            lprintln("info,sourceindex", "index unchanged: ${files.size} files, ${definitions.size} definitions")
        } else {
            runCatching {
                Files.createDirectories(indexPath.parent)
                mapper.writerWithDefaultPrettyPrinter().writeValue(indexPath.toFile(), index)
            }.onFailure {
                lprintln("error,sourceindex", "failed to persist index: ${it.message}")
            }

            latestIndex = index
            lprintln("info,sourceindex", "indexed ${files.size} files, ${definitions.size} definitions")
        }
    }

    private fun KotlinSourceIndexDb.materiallySameAs(other: KotlinSourceIndexDb): Boolean =
        root == other.root &&
            apiKeyHash == other.apiKeyHash &&
            vectorStoreId == other.vectorStoreId &&
            vectorStoreFileHash == other.vectorStoreFileHash &&
            files.map { it.copy(modifiedEpochMs = 0L) } == other.files.map { it.copy(modifiedEpochMs = 0L) } &&
            definitions == other.definitions

    fun getCompletionCategoryForSymbol(
        symbolName: String,
        sourceName: String? = null,
        sourcePath: String? = null,
        declaredIn: String? = null
    ): String? {
        val index = currentIndex() ?: return null
        val simpleName = symbolName.substringAfterLast('.').trim()
        if (simpleName.isBlank()) return null

        fun KotlinDefinitionRecord.categoryOrNull(): String? =
            completionCategory?.trim()?.takeIf { it.isNotEmpty() }

        if (!sourceName.isNullOrBlank()) {
            index.definitions.firstOrNull { it.name == sourceName }?.categoryOrNull()?.let { return it }
        }

        val matches = index.definitions.filter {
            it.name.substringAfterLast('.') == simpleName ||
                (!sourceName.isNullOrBlank() && it.name == sourceName)
        }

        if (!sourcePath.isNullOrBlank()) {
            matches.firstOrNull { it.filePath == sourcePath }?.categoryOrNull()?.let { return it }
        }

        val container = declaredIn?.trim()?.takeIf { it.isNotEmpty() }
        if (container != null) {
            // A member: ask for the exact declaration, `Owner.member`. If it is in the index, its answer is
            // the only correct one — including "no category", which must NOT fall through to the name vote.
            val ownerSimple = container.substringAfterLast('.')
            val exact = matches.filter { it.name == "$ownerSimple.$simpleName" }
            if (exact.isNotEmpty()) return exact.firstNotNullOfOrNull { it.categoryOrNull() }

            // A top-level extension: `declaredIn` is a package, which scopes the vote to the right file
            // instead of choosing between 42 of them. A package can still contain several categorised
            // files, so answer only when those files agree rather than inheriting filesystem walk order.
            val packagePath = container.replace('.', '/')
            val inPackage = matches.filter { it.filePath.contains("/$packagePath/") }
            if (inPackage.isNotEmpty()) {
                return inPackage.mapNotNull { it.categoryOrNull() }.distinct().singleOrNull()
            }

            // `declaredIn` names something outside this tree — `java.lang.Math`, a jar. This tree's
            // namesakes say nothing about it, so decline instead of borrowing their category.
            return null
        }

        // No declaring container to go on: answer only if the tree is unanimous about the name.
        return matches.mapNotNull { it.categoryOrNull() }.distinct().singleOrNull()
    }

    internal fun parseTopLevelDefinitions(text: String): List<ParsedDefinition> {
        val out = mutableListOf<ParsedDefinition>()

        // Extract KDocs using the KDocExtractor
        val kdocs = KDocExtractor.extractKDocs(text)
        val kdocByLine = kdocs.associateBy { it.declarationLine }
        val fileCategory = fileCompletionCategory(text)

        val state = KotlinParseState()
        val classStack = ArrayDeque<ClassContext>()
        var pendingClass: PendingClass? = null
        var pendingCategory: String? = null

        val lines = text.split("\n")
        lines.forEachIndexed { index, rawLine ->
            val lineNumber = index + 1
            val depthAtStart = state.depth
            val currentClass = classStack.lastOrNull()
            val canParse = !state.inBlockComment &&
                !state.inRawString &&
                !state.inString &&
                (depthAtStart == 0 || (currentClass != null && depthAtStart == currentClass.depth))

            if (canParse) {
                val def = parseDefinitionFromLine(rawLine)
                if (def == null && rawLine.trimStart().startsWith("@")) {
                    completionCategoryFromAnnotation(rawLine)?.let { pendingCategory = it }
                }
                if (pendingClass != null &&
                    pendingClass.lineNumber != lineNumber &&
                    pendingClass.headerParenDepth == 0 &&
                    def != null
                ) {
                    // A new declaration appeared before the pending class body opened.
                    pendingClass = null
                }
                if (pendingClass != null &&
                    pendingClass.lineNumber != lineNumber &&
                    pendingClass.headerParenDepth == 0 &&
                    INIT_BLOCK.containsMatchIn(rawLine)
                ) {
                    // `data class Foo(…)` with no body, then the enclosing class's `init {`: that
                    // brace is not Foo's body, and taking it for one attributed the init block's
                    // locals to Foo (CommandBuffer.SecondaryInheritance, JoltPhysicsWorld.SurfaceQuery,
                    // StudyBuilds.Manifest). No class header continues with `init`.
                    pendingClass = null
                }
                def?.let { parsed ->
                    val kind = parsed.kind
                    val name = parsed.name
                    val kdoc = kdocByLine[lineNumber]?.text
                    val category = parsed.completionCategory ?: pendingCategory ?: fileCategory
                    pendingCategory = null
                    when (kind) {
                        "class", "interface", "object" -> {
                            val qualifiedName = if (currentClass != null) {
                                "${currentClass.name}.$name"
                            } else {
                                name
                            }
                            out += ParsedDefinition(
                                kind, qualifiedName, lineNumber, kdoc, category,
                                parsed.visibility, currentClass?.name
                            )
                            pendingClass = PendingClass(qualifiedName, lineNumber)
                        }
                        "fun", "val", "var", "typealias" -> {
                            // A `val`/`var` on a continuation line of a class's primary constructor
                            // is that class's property, not a top-level declaration. Without this the
                            // index gained bogus bare entries (`re`, `im`, `sampleRate`, …) that any
                            // simple-name KDoc lookup could match by accident.
                            val constructorOwner = pendingClass
                                ?.takeIf { (kind == "val" || kind == "var") && it.headerParenDepth > 0 }
                                ?.name
                            val owner = constructorOwner
                                ?: currentClass?.takeIf { depthAtStart == it.depth }?.name
                            val qualifiedName = if (owner != null) "$owner.$name" else name
                            out += ParsedDefinition(
                                kind, qualifiedName, lineNumber, kdoc, category,
                                parsed.visibility, owner
                            )
                        }
                    }
                }
            }

            var openedClassThisLine = false
            state.consumeLine(rawLine) { ch, _, _, after ->
                val pending = pendingClass
                if (pending != null && !openedClassThisLine) {
                    when (ch) {
                        '(' -> pending.headerParenDepth++
                        ')' -> if (pending.headerParenDepth > 0) pending.headerParenDepth--
                        '{' -> if (pending.headerParenDepth == 0) {
                            classStack.addLast(ClassContext(pending.name, after))
                            pendingClass = null
                            openedClassThisLine = true
                        }
                    }
                }
            }

            while (classStack.isNotEmpty() && state.depth < classStack.last().depth) {
                classStack.removeLast()
            }
        }
        return out
    }

    private fun parseDefinitionFromLine(line: String): ParsedLineDefinition? {
        var s = line.trimStart()
        if (s.isEmpty()) return null
        if (s.startsWith("package ") || s.startsWith("import ")) return null
        if (s.startsWith("//")) return null

        val annotationRegex = Regex("^@\\w[\\w\\.]*\\s*(\\([^)]*\\))?\\s*")
        var category: String? = null
        while (true) {
            val match = annotationRegex.find(s)
            if (match != null && match.range.first == 0) {
                completionCategoryFromAnnotation(match.value)?.let { category = it }
                s = s.substring(match.range.last + 1).trimStart()
            } else {
                break
            }
        }

        val tokens = s.split(Regex("\\s+"))
        var idx = 0
        var sawCompanion = false
        var visibility = "public"
        while (idx < tokens.size && modifierTokens.contains(tokens[idx])) {
            if (tokens[idx] == "companion") sawCompanion = true
            if (tokens[idx] in visibilityTokens) visibility = tokens[idx]
            idx++
        }
        if (idx >= tokens.size) return null

        var keyword = tokens[idx]
        var remainder = s.substringAfter(keyword, "").trimStart()

        // `fun interface Foo` declares a SAM interface, not a function. Skipping it hid a whole
        // class of public API (FnMetric, TuioListener, QuantileFunction…) from the index.
        if (keyword == "fun" && FUN_INTERFACE.containsMatchIn(remainder)) {
            keyword = "interface"
            remainder = remainder.removePrefix("interface").trimStart()
        }

        val rawName = when (keyword) {
            "fun" -> {
                val head = remainder.substringBefore("(").trim()
                parseQualifiedName(head)
            }
            "class", "interface", "object" -> {
                Regex("^(`[^`]+`|[A-Za-z_][A-Za-z0-9_]*)")
                    .find(remainder)?.groupValues?.get(1)?.let { unwrapIdentifier(it) }
            }
            "val", "var" -> {
                val head = propertyNameHead(remainder)
                parseQualifiedName(head)
            }
            "typealias" -> {
                val head = remainder.substringBefore("=").trim()
                parseQualifiedName(head)
            }
            else -> null
        }
        val effectiveName = when {
            keyword == "object" && (rawName == null || rawName.isBlank()) && sawCompanion -> "Companion"
            rawName == null -> return null
            else -> rawName
        }
        val name = effectiveName
        if (name.isBlank()) return null
        if (keyword == "class" || keyword == "interface" || keyword == "object" || keyword == "fun" || keyword == "val" || keyword == "var" || keyword == "typealias") {
            return ParsedLineDefinition(keyword, name, category, visibility)
        }
        return null
    }

    private fun fileCompletionCategory(text: String): String? =
        Regex("""(?m)^\s*@file:(?:[\w.]+\.)?CompletionCategory\s*\(\s*"([^"]+)"""")
            .find(text)
            ?.groupValues
            ?.getOrNull(1)
            ?.trim()
            ?.takeIf { it.isNotEmpty() }

    private fun completionCategoryFromAnnotation(text: String): String? =
        Regex("""@(?:[\w.]+\.)?CompletionCategory\s*\(\s*"([^"]+)"""")
            .find(text)
            ?.groupValues
            ?.getOrNull(1)
            ?.trim()
            ?.takeIf { it.isNotEmpty() }

    private fun parseQualifiedName(raw: String): String? {
        if (raw.isBlank()) return null
        val noGenerics = stripGenericArgs(raw)
        val compact = removeWhitespaceOutsideBackticks(noGenerics).replace("?", "")
        if (compact.isBlank()) return null
        val regex = Regex("^(`[^`]+`|[A-Za-z_][A-Za-z0-9_]*)(\\.`[^`]+`|\\.[A-Za-z_][A-Za-z0-9_]*)*$")
        val match = regex.find(compact) ?: return null
        return match.value.split(".").joinToString(".") { unwrapIdentifier(it) }
    }

    private fun unwrapIdentifier(raw: String): String =
        if (raw.startsWith("`") && raw.endsWith("`") && raw.length > 1) {
            raw.removePrefix("`").removeSuffix("`")
        } else {
            raw
        }

    private fun stripGenericArgs(raw: String): String {
        val sb = StringBuilder(raw.length)
        var depth = 0
        for (ch in raw) {
            when (ch) {
                '<' -> depth++
                '>' -> if (depth > 0) depth--
                else -> if (depth == 0) sb.append(ch)
            }
        }
        return sb.toString()
    }

    private fun removeWhitespaceOutsideBackticks(raw: String): String {
        val sb = StringBuilder(raw.length)
        var inBackticks = false
        for (ch in raw) {
            if (ch == '`') {
                inBackticks = !inBackticks
                sb.append(ch)
                continue
            }
            if (!inBackticks && ch.isWhitespace()) continue
            sb.append(ch)
        }
        return sb.toString()
    }

    private class KotlinParseState {
        /** Brace nesting depth in code: braces inside strings, comments and templates do not count. */
        var depth: Int = 0
            private set

        /** Nesting count for block comments, which nest in Kotlin. Zero means "not in a comment". */
        private var blockCommentDepth: Int = 0

        private enum class Frame { CODE, STRING, RAW_STRING }

        /**
         * The lexical context stack, innermost last, with a `CODE` frame at the bottom that is never
         * popped. A string pushes `STRING`/`RAW_STRING`; a `${` inside one pushes a `CODE` frame back
         * on top, which is what makes `"${f("x")}"` come out right.
         */
        private val frames = ArrayDeque<Frame>().apply { addLast(Frame.CODE) }

        /** The opening quote of each `STRING` frame, so a `'` literal is not closed by a `"`. */
        private val quotes = ArrayDeque<Char>()

        /** Open brace count *within* the innermost `CODE` frame, so a template knows its own `}`. */
        private val braces = ArrayDeque<Int>().apply { addLast(0) }

        val inBlockComment: Boolean get() = blockCommentDepth > 0

        /** True anywhere lexically inside a raw string, including inside one of its templates. */
        val inRawString: Boolean get() = frames.contains(Frame.RAW_STRING)

        /** True anywhere lexically inside a normal string or character literal. */
        val inString: Boolean get() = frames.contains(Frame.STRING)

        private fun pushCode() {
            frames.addLast(Frame.CODE)
            braces.addLast(0)
        }

        private fun popFrame() {
            when (frames.removeLast()) {
                Frame.CODE -> braces.removeLast()
                Frame.STRING -> quotes.removeLast()
                Frame.RAW_STRING -> Unit
            }
        }

        fun consumeLine(line: String, onToken: ((ch: Char, index: Int, before: Int, after: Int) -> Unit)? = null) {
            var i = 0
            while (i < line.length) {
                if (blockCommentDepth > 0) {
                    when {
                        line.startsWith("*/", i) -> { blockCommentDepth--; i += 2 }
                        line.startsWith("/*", i) -> { blockCommentDepth++; i += 2 }
                        else -> i++
                    }
                    continue
                }
                when (frames.last()) {
                    Frame.RAW_STRING -> {
                        when {
                            line.startsWith("\${", i) -> { pushCode(); i += 2 }
                            line.startsWith("\"\"\"", i) -> {
                                // A run of N ≥ 3 quotes closes with its last three; the rest are content.
                                var run = 0
                                while (i + run < line.length && line[i + run] == '"') run++
                                popFrame()
                                i += run
                            }
                            else -> i++
                        }
                        continue
                    }
                    Frame.STRING -> {
                        when {
                            line[i] == '\\' -> i += 2
                            line.startsWith("\${", i) -> { pushCode(); i += 2 }
                            line[i] == quotes.last() -> { popFrame(); i++ }
                            else -> i++
                        }
                        continue
                    }
                    Frame.CODE -> Unit
                }

                if (line.startsWith("//", i)) break
                if (line.startsWith("/*", i)) {
                    blockCommentDepth = 1
                    i += 2
                    continue
                }
                if (line.startsWith("\"\"\"", i)) {
                    frames.addLast(Frame.RAW_STRING)
                    i += 3
                    continue
                }
                val ch = line[i]
                if (ch == '`') {
                    // A backtick-quoted identifier. It cannot span a line, so an unclosed one is a
                    // lexical error; treat that backtick as an ordinary character rather than eating
                    // the rest of the line.
                    val close = line.indexOf('`', i + 1)
                    i = if (close < 0) i + 1 else close + 1
                    continue
                }
                if (ch == '"' || ch == '\'') {
                    frames.addLast(Frame.STRING)
                    quotes.addLast(ch)
                    i++
                    continue
                }
                // Tokens are only reported from the outermost code frame: a brace inside a string
                // template is a lambda in an interpolation, never a class body.
                val outermost = frames.size == 1
                if (ch == '{') {
                    val before = depth
                    depth++
                    braces.addLast(braces.removeLast() + 1)
                    if (outermost) onToken?.invoke(ch, i, before, depth)
                } else if (ch == '}') {
                    if (braces.last() == 0 && frames.size > 1) {
                        // The `}` that ends a `${…}`; the string it interrupted resumes.
                        popFrame()
                        i++
                        continue
                    }
                    val before = depth
                    depth = (depth - 1).coerceAtLeast(0)
                    braces.addLast((braces.removeLast() - 1).coerceAtLeast(0))
                    if (outermost) onToken?.invoke(ch, i, before, depth)
                } else if (ch == '(' || ch == ')') {
                    if (outermost) onToken?.invoke(ch, i, depth, depth)
                }
                i++
            }
            // A normal string or character literal cannot contain a newline. One still open here is
            // a mis-lex; end it, so the damage stops at this line. A `${…}` left open is not dropped:
            // a template expression legitimately spans lines.
            while (frames.size > 1 && frames.last() == Frame.STRING) popFrame()
        }
    }
}