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

import com.correx.core.artifacts.kind.JsonSchema
import com.correx.core.artifacts.kind.JsonSchemaProperty
import org.junit.jupiter.api.Assertions.assertFalse
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test

class GbnfGrammarConverterTest {

    /** FileWrittenPayload-shaped schema: all required fields. */
    private val allRequiredSchema = JsonSchema(
        type = "object",
        properties = mapOf(
            "path" to JsonSchemaProperty(type = "string", description = "relative file path"),
            "content" to JsonSchemaProperty(type = "string", description = "file content"),
            "mode" to JsonSchemaProperty(type = "string", description = "octal permission mode"),
        ),
        required = listOf("path", "content", "mode"),
    )

    /** Schema with some required and some optional fields. */
    private val mixedSchema = JsonSchema(
        type = "object",
        properties = mapOf(
            "name" to JsonSchemaProperty(type = "string"),
            "count" to JsonSchemaProperty(type = "integer"),
            "description" to JsonSchemaProperty(type = "string"),
        ),
        required = listOf("name"),
    )

    /** Schema with all optional fields (no required entries). */
    private val allOptionalSchema = JsonSchema(
        type = "object",
        properties = mapOf(
            "label" to JsonSchemaProperty(type = "string"),
            "value" to JsonSchemaProperty(type = "integer"),
        ),
        required = emptyList(),
    )

    @Test
    fun `array property maps to the array rule, not string`() {
        // Regression: an array-typed property used to fall through to `string`, so the grammar forced
        // a string where the schema demands an array (analysis.requirements). It must map to `array`.
        val schema = JsonSchema(
            type = "object",
            properties = mapOf(
                "summary" to JsonSchemaProperty(type = "string"),
                "requirements" to JsonSchemaProperty(type = "array", items = JsonSchemaProperty(type = "string")),
            ),
            required = listOf("summary", "requirements"),
        )
        val grammar = GbnfGrammarConverter.convert(schema)
        assertTrue(grammar.contains("array  ::="), "grammar must declare an array rule")
        val reqLine = grammar.lineSequence().first { it.contains("\\\"requirements\\\"") }
        assertTrue(reqLine.trimEnd().endsWith("array"), "requirements must use the array rule, was: $reqLine")
    }

    @Test
    fun `string rule does not contain quoted character class`() {
        val grammar = GbnfGrammarConverter.convert(allRequiredSchema)
        // Guard the original bug: the `string` rule's char class must stay unquoted. The `array` rule
        // legitimately contains a quoted "[" literal, so scope the assertion to the string rule line.
        val stringRuleLine = grammar.lineSequence().first { it.trimStart().startsWith("string ::=") }
        assertFalse(stringRuleLine.contains("\"["), "string rule must not contain quoted character class (\"[)")
    }

    @Test
    fun `all-required schema produces non-empty valid GBNF`() {
        val grammar = GbnfGrammarConverter.convert(allRequiredSchema)

        assertFalse(grammar.isBlank(), "grammar must not be blank")
        assertTrue(grammar.contains("root"), "grammar must declare root rule")
        assertTrue(grammar.contains("members"), "grammar must declare members rule")
        assertTrue(grammar.contains("path"), "grammar must reference 'path' key")
        assertTrue(grammar.contains("content"), "grammar must reference 'content' key")
        assertTrue(grammar.contains("mode"), "grammar must reference 'mode' key")
        // Optional wrapper must not appear for required keys
        assertFalse(grammar.substringAfter("members").trimStart().startsWith("("),
            "first member of all-required schema must not be wrapped in optional ()?")
    }

    @Test
    fun `mixed required and optional schema produces non-empty valid GBNF`() {
        val grammar = GbnfGrammarConverter.convert(mixedSchema)

        assertFalse(grammar.isBlank(), "grammar must not be blank")
        assertTrue(grammar.contains("root"), "grammar must declare root rule")
        assertTrue(grammar.contains("members"), "grammar must declare members rule")
        assertTrue(grammar.contains("name"), "grammar must reference required key 'name'")
        assertTrue(grammar.contains("count"), "grammar must reference optional key 'count'")
        assertTrue(grammar.contains("description"), "grammar must reference optional key 'description'")
        // Optional keys should be wrapped in (...)?
        val membersLine = grammar.lines().first { it.startsWith("members") }
        assertTrue(membersLine.contains(")?"), "optional fields must use (...)? in members rule")
    }

    @Test
    fun `all-optional schema produces non-empty valid GBNF without leading comma on first member`() {
        val grammar = GbnfGrammarConverter.convert(allOptionalSchema)

        assertFalse(grammar.isBlank(), "grammar must not be blank")
        assertTrue(grammar.contains("root"), "grammar must declare root rule")
        assertTrue(grammar.contains("members"), "grammar must declare members rule")
        // The members rule line must not start the first optional key with a comma
        val membersLine = grammar.lines().first { it.startsWith("members") }
        // First token after "members ::= " must be "(" (the optional wrapper), not a comma
        val afterAssign = membersLine.substringAfter("::=").trimStart()
        assertFalse(afterAssign.startsWith(","), "first member must not begin with a comma")
        assertTrue(afterAssign.startsWith("("), "first optional member must be wrapped in (...)?")
    }
}
