ai_macro.kt

package afield.runtime.macros

import afield.ai.agent.*
import afield.ai.conversation.*
import afield.ai.tools.*
import afield.animation.gesture.*
import afield.animation.interpolation.*
import afield.animation.keyframes.*
import afield.app.diagnostics.*
import afield.app.notifications.*
import afield.app.protocol.*
import afield.app.server.*
import afield.app.settings.*
import afield.audio.synth.*
import afield.data.table.*
import afield.document.context.*
import afield.document.graph.*
import afield.document.history.*
import afield.document.persistence.*
import afield.document.properties.*
import afield.document.workspace.*
import afield.graphics.color.*
import afield.graphics.paint.*
import afield.graphics.vulkan.*
import afield.math.functions.*
import afield.media.capture.*
import afield.media.stream.*
import afield.media.video.*
import afield.runtime.completion.*
import afield.runtime.execution.*
import afield.runtime.kotlin.*
import afield.runtime.lifecycle.*
import afield.tools.automation.*
import afield.tools.docs.*
import afield.ui.controls.*
import afield.ui.layout.*
import afield.ui.notation.*
import afield.ui.output.*
import afield.ui.output.templates.*
import afield.ui.slides.*
import afield.ui.visualization.*
import afield.ai.tools.DefinitionInfo
import afield.ai.tools.ListDefinitionsResult
import afield.ai.models.Ai4ModelDefaults
import afield.ai.models.LlmBridge
import afield.ai.models.LlmReasoningEffort
import afield.ai.models.LlmThreadState
import afield.ai.diagnostics.StdoutLog
import afield.ai.tools.ToolExecutionContext
import afield.ai.models.aiProviderConfigRevision
import afield.ai.models.createLlmBridge
import afield.app.diagnostics.logging.lprintln
import afield.runtime.macros.MacroRegistry.MacroExpansionProblem
import afield.runtime.macros.MacroRegistry.callRewriteWithContext
import kotlin.script.experimental.api.ScriptDiagnostic
import kotlinx.coroutines.runBlocking
import java.util.concurrent.ConcurrentHashMap

fun Box.macro_ai(args: List<String>, containingFunction: ContainingFunction?): Insert? {
    if (containingFunction == null) {
        throw MacroExpansionProblem(
            "ai() must be called inside a function with explicit parameter types and return type, e.g.:\n" +
            "  fun foo(a: Double): Double { ai(\"description\") }"
        )
    }

    // Parse the prompt from args — args looks like: ["(\"some prompt text\")"]
    val rawArgs = args.joinToString("")
    val promptMatch = Regex("""^\(\s*"(.*)"\s*\)$""", RegexOption.DOT_MATCHES_ALL).find(rawArgs)
        ?: throw MacroExpansionProblem("ai() expects a single string argument, e.g. ai(\"description of what this function does\")")
    val prompt = promptMatch.groupValues[1]

    val funcName = containingFunction.name
    val implName = "__$funcName"
    val paramList = containingFunction.parameterList
    val returnType = containingFunction.returnType
    val signature = "$paramList: $returnType"
    val paramNames = extractParamNames(paramList)
    val forwardCall = if (paramNames.isBlank()) "return $implName()" else "return $implName($paramNames)"

    // Check attic for existing matching block
    val currentAttic = this.attic
    val existing = parseAtticBlock(currentAttic, implName)

    if (existing != null && existing.signature == signature && existing.prompt == prompt) {
        // Cache hit — signature and prompt match
        lprintln("info,ai.macro", "ai() cache hit for $implName in box ${this.id}")
        safeNotify("ai($funcName): using cached implementation")
        return Insert(here = forwardCall)
    }

    // Remove stale block if it exists
    val cleanedAttic = if (existing != null) {
        lprintln("info,ai.macro", "ai() removing stale block for $implName (sig or prompt changed)")
        removeAtticBlock(currentAttic, implName)
    } else {
        currentAttic
    }

    // Build AI context
    val atticContext = stripAiGenMetadata(cleanedAttic).trim()

    // Get preview-expanded code with this ai() macro skipped and the containing function removed
    val rawCode = this.code.orEmpty()
    val codeWithoutFunction = removeFunctionByName(rawCode, funcName)
    val expandedCode = try {
        this.ki.applyMacros(this, codeWithoutFunction, skipMacros = setOf("ai"), previewOnly = true)
    } catch (e: Exception) {
        lprintln("warning,ai.macro", "Failed to preview-expand code for AI context: ${e.message}")
        codeWithoutFunction
    }

    // Build the prompt for AI
    val userPrompt = buildString {
        appendLine("Write a Kotlin function with the following specification:")
        appendLine()
        appendLine("Function name: $implName")
        appendLine("Parameters: $paramList")
        appendLine("Return type: $returnType")
        appendLine("Description: $prompt")
        appendLine()
        if (atticContext.isNotBlank()) {
            appendLine("Here are other definitions already available in this scope (attic):")
            appendLine("```")
            appendLine(atticContext)
            appendLine("```")
            appendLine()
        }
        if (expandedCode.isNotBlank()) {
            appendLine("Here is the rest of the code in this box for context:")
            appendLine("```")
            appendLine(expandedCode)
            appendLine("```")
            appendLine()
        }
        appendLine("Write ONLY the function definition for $implName. No markdown, no explanation.")
    }

    // Notify the user that generation is in progress
    lprintln("info,ai.macro", "ai() generating $implName for box ${this.id}: \"$prompt\"")
    safeNotify("ai($funcName): generating via AI — \"$prompt\"")

    val contextForThread = buildString {
        if (atticContext.isNotBlank()) {
            appendLine(atticContext)
            appendLine()
        }
        appendLine(expandedCode)
    }

    val maxAttempts = 3
    var currentPrompt = userPrompt
    var generatedCode: String? = null

    for (attempt in 1..maxAttempts) {
        // Generate code from AI
        safeNotify("ai($funcName): calling AI${if (attempt > 1) " (retry $attempt/$maxAttempts)" else ""}...")
        val candidate = try {
            val thread = AiMacroBridge.getOrCreateThread(this.id, contextForThread)

            // Build definitions from this box's KotlinInterface context
            val box = this
            val ctx = ToolExecutionContext(
                listDefinitions = {
                    val defs = mutableListOf<DefinitionInfo>()
                    try {
                        box.ki.context.vars.values.mapTo(defs) { v ->
                            DefinitionInfo(name = v.name, type = "variable", fileRef = "box:${box.id}", language = "kotlin")
                        }
                        box.ki.context.functions.values.mapTo(defs) { f ->
                            DefinitionInfo(name = f.toString(true), type = "function", fileRef = "box:${box.id}", language = "kotlin")
                        }
                    } catch (e: Exception) {
                        lprintln("warning,ai.macro", "Failed to gather definitions for ListDefinitions: ${e.message}")
                    }
                    ListDefinitionsResult(definitions = defs)
                }
            )

            val result = AiMacroBridge.bridge.send(
                state = thread,
                userPrompt = currentPrompt,
                ctx = ctx,
                maxOutputTokens = 4096
            )
            result.text.trim()
                .removePrefix("```kotlin").removePrefix("```")
                .removeSuffix("```")
                .trim()
        } catch (e: Exception) {
            lprintln("error,ai.macro", "AI generation failed for $implName (attempt $attempt): ${e.message}")
            if (attempt == maxAttempts) {
                safeNotify("ai($funcName): AI generation failed after $maxAttempts attempts")
                throw MacroExpansionProblem("AI generation failed for $implName: ${e.message}")
            }
            continue
        }

        // Validate that the generated code looks like a function
        if (!candidate.contains("fun $implName")) {
            lprintln("error,ai.macro", "AI did not generate expected function $implName (attempt $attempt). Got:\n$candidate")
            if (attempt == maxAttempts) {
                safeNotify("ai($funcName): AI failed to produce valid function after $maxAttempts attempts")
                throw MacroExpansionProblem("AI did not generate expected function $implName. Response:\n$candidate")
            }
            currentPrompt = buildString {
                appendLine("Your previous response did not contain the expected function definition 'fun $implName'.")
                appendLine("Here is what you returned:")
                appendLine("```")
                appendLine(candidate)
                appendLine("```")
                appendLine()
                appendLine("Please try again. Write ONLY the function definition for $implName with the exact signature: fun $implName$paramList: $returnType")
            }
            continue
        }

        // Compile-check the generated code
        safeNotify("ai($funcName): compile-checking generated code...")
        val compileErrors = compileCheckGenerated(candidate, cleanedAttic, expandedCode)

        if (compileErrors == null) {
            // Compilation succeeded
            lprintln("info,ai.macro", "ai() compile check passed for $implName on attempt $attempt")
            safeNotify("ai($funcName): compile check passed" + if (attempt > 1) " (after $attempt attempts)" else "")
            generatedCode = candidate
            break
        } else {
            // Compilation failed — show the first error in the notification
            val firstError = compileErrors.lines().firstOrNull() ?: compileErrors
            lprintln("warning,ai.macro", "ai() compile check failed for $implName (attempt $attempt):\n$compileErrors")
            safeNotify("ai($funcName): compile error (attempt $attempt): $firstError")

            if (attempt == maxAttempts) {
                // Last attempt — use the code anyway but warn the user
                lprintln("warning,ai.macro", "ai() giving up on compile-clean code for $implName after $maxAttempts attempts, using last result")
                safeNotify("ai($funcName): could not get clean compile after $maxAttempts attempts — using last result")
                generatedCode = candidate
                break
            }

            // Build a fix-it prompt for the next iteration
            currentPrompt = buildString {
                appendLine("The function you generated has Kotlin compiler errors:")
                appendLine()
                appendLine("```")
                appendLine(compileErrors)
                appendLine("```")
                appendLine()
                appendLine("Here is the code you wrote:")
                appendLine("```kotlin")
                appendLine(candidate)
                appendLine("```")
                appendLine()
                appendLine("Please fix these errors and return the corrected function definition for $implName.")
                appendLine("Return ONLY the function definition, no markdown fences, no explanation.")
            }
        }
    }

    if (generatedCode == null) {
        safeNotify("ai($funcName): generation failed")
        throw MacroExpansionProblem("AI generation failed for $implName after $maxAttempts attempts")
    }

    // Build the attic block and patch attic
    val block = AtticBlock(
        functionName = implName,
        signature = signature,
        prompt = prompt,
        body = generatedCode
    )
    val newAttic = appendAtticBlock(cleanedAttic, block)

    // Update attic property (this triggers resync to frontend)
    this.setPropertyRaw("attic", newAttic, doResync = true)
    // _attic is intentionally NOT updated here so checkAtticAndExecute detects the change and executes the new code

    lprintln("info,ai.macro", "ai() generated $implName successfully, attic updated")
    safeNotify("ai($funcName): done ✓")

    return Insert(here = forwardCall)
}

private fun safeNotify(text: String) {
    try {
        _send(Message("notify").payload(Notify(text, "AI")))
    } catch (e: Exception) {
        lprintln("warning,ai.macro", "safeNotify failed: ${e.message}")
    }
}