package com.correx.infrastructure.inference.llama.cpp

import com.correx.core.artifacts.kind.JsonSchema

internal object GbnfGrammarConverter {
    fun convert(schema: JsonSchema): String {
        require(schema.type == "object") { "Only object schemas supported" }

        val required = schema.required.toSet()
        val allKeys = schema.properties.keys.toList()

        val sb = StringBuilder()
        sb.appendLine("""root   ::= "{" ws members ws "}"""")
        sb.appendLine("""ws     ::= [ \t\n\r]*""")
        val hexDigits = "[0-9a-fA-F]"
        val unicodeEscape = "\"u\" $hexDigits $hexDigits $hexDigits $hexDigits"
        val closingQuote = "\"\\\"\"" // GBNF string literal matching a single "
        val charClass = """[^"\\\x7F\x00-\x1F]"""
        val escapeSeq = """"\\" (["\\/bfnrt] | $unicodeEscape)"""
        val stringRule = """string ::= "\"" ($charClass | $escapeSeq)* $closingQuote"""
        sb.appendLine(stringRule)
        sb.appendLine("""number ::= "-"? [0-9]+ ("." [0-9]+)?""")
        sb.appendLine("""boolean ::= "true" | "false"""")
        sb.appendLine("""null   ::= "null"""")
        // Generic JSON value/object/array rules so non-scalar properties are constrained to the right
        // JSON shape. Without these, an `array` property fell through to `string`, so the grammar forced
        // the model to emit a string where the schema (and post-hoc validation) demand an array — an
        // unwinnable disagreement. Element/field *types* are left to schema validation; the grammar only
        // enforces the JSON shape (brackets/braces).
        sb.appendLine("""value  ::= string | number | boolean | null | object | array""")
        sb.appendLine("""object ::= "{" ws (string ws ":" ws value (ws "," ws string ws ":" ws value)*)? ws "}"""")
        sb.appendLine("""array  ::= "[" ws (value (ws "," ws value)*)? ws "]"""")

        // Build members rule:
        // - Required keys are always present, joined by commas.
        // - Optional keys are wrapped in (...)?; they always carry a leading comma so they
        //   can be appended after whatever came before without complicating the grammar.
        // - When there are no required keys, the first optional key cannot have a leading comma.
        //   We handle this by making the first key (required or optional) comma-free, and all
        //   subsequent keys comma-prefixed (optional ones wrapped in (...)? around the comma too).
        val memberParts = mutableListOf<String>()

        for ((index, key) in allKeys.withIndex()) {
            val valueRule = valueRuleFor(schema, key)
            val isRequired = key in required
            val pair = "\"\\\"$key\\\"\" ws \":\" ws $valueRule"
            if (index == 0) {
                // First key: never a leading comma
                if (isRequired) memberParts.add(pair) else memberParts.add("($pair)?")
            } else {
                // Subsequent keys always have a leading comma, gated by optional wrapper if needed
                if (isRequired) {
                    memberParts.add("ws \",\" ws $pair")
                } else {
                    memberParts.add("(ws \",\" ws $pair)?")
                }
            }
        }

        sb.append("members ::= ")
        if (memberParts.isEmpty()) {
            sb.appendLine("ws")
        } else {
            sb.appendLine(memberParts.joinToString(" "))
        }

        return sb.toString()
    }

    private fun valueRuleFor(schema: JsonSchema, key: String): String {
        val prop = schema.properties[key] ?: return "value"
        return when (prop.type) {
            "integer", "number" -> "number"
            "boolean" -> "boolean"
            "array" -> "array"
            "object" -> "object"
            "string" -> "string"
            else -> "value"
        }
    }
}
