dependencyResolutionManagement {
    repositories {
        google()
        mavenCentral()
    }
}

commonMain.dependencies {
    implementation("com.llamatik:library:1.5.0")
}

// Resolve model path (place GGUF in assets / bundle)
val modelPath = LlamaBridge.getModelPath("phi-2.Q4_0.gguf")

// (Optional) tune parameters before loading — contextLength/useMmap/flashAttention
// take effect at model init time; the others can be changed at any time
LlamaBridge.updateGenerateParams(
    temperature    = 0.7f,
    maxTokens      = 512,
    topP           = 0.95f,
    topK           = 40,
    repeatPenalty  = 1.1f,
    contextLength  = 4096,
    numThreads     = 4,
    useMmap        = true,
    flashAttention = false,
)

// Load model
LlamaBridge.initGenerateModel(modelPath)

// Generate text
val output = LlamaBridge.generate(
    "Explain Kotlin Multiplatform in one sentence."
)

@Suppress("EXPECT_ACTUAL_CLASSIFIERS_ARE_IN_BETA_WARNING")
expect object LlamaBridge {
    // Utilities
    fun getModelPath(modelFileName: String): String   // copy asset/bundle model to app files dir and return absolute path
    fun shutdown()                                    // free native resources

    // Embeddings
    fun initEmbedModel(modelPath: String): Boolean    // load embeddings model
    fun embed(input: String): FloatArray              // return embedding vector

    // Text generation (non-streaming)
    fun initGenerateModel(modelPath: String): Boolean // load generation model
    fun generate(prompt: String): String
    fun generateWithContext(
        systemPrompt: String,
        contextBlock: String,
        userPrompt: String
    ): String

    // Text generation (streaming)
    fun generateStream(prompt: String, callback: GenStream)
    fun generateStreamWithContext(
        systemPrompt: String,
        contextBlock: String,
        userPrompt: String,
        callback: GenStream
    )

    // Convenience streaming overload (lambda callbacks)
    fun generateWithContextStream(
        system: String,
        context: String,
        user: String,
        onDelta: (String) -> Unit,
        onDone: () -> Unit,
        onError: (String) -> Unit
    )

    // Text generation with JSON schema (non-streaming)
    fun generateJson(prompt: String, jsonSchema: String? = null): String
    fun generateJsonWithContext(
        systemPrompt: String,
        contextBlock: String,
        userPrompt: String,
        jsonSchema: String? = null
    ): String

    // Text generation with JSON schema (streaming)
    fun generateJsonStream(prompt: String, jsonSchema: String? = null, callback: GenStream)
    fun generateJsonStreamWithContext(
        systemPrompt: String,
        contextBlock: String,
        userPrompt: String,
        jsonSchema: String? = null,
        callback: GenStream
    )

    // Model metadata
    fun getModelFinetuneType(): String?               // "general.finetune" GGUF key — e.g. "instruct", "chat"; null means base model

    // Chat template support
    fun getModelChatTemplate(): String?               // returns the chat template embedded in the loaded GGUF, or null
    fun applyChatTemplate(
        messages: List<Pair<String, String>>,         // list of (role, content) pairs
        addAssistantPrefix: Boolean                   // true to append the assistant turn prefix
    ): String?                                        // rendered prompt string, or null if model/template unavailable

    // KV cache session support
    fun sessionReset(): Boolean                       // clear KV state, keep model loaded
    fun sessionSave(path: String): Boolean            // persist KV state to file
    fun sessionLoad(path: String): Boolean            // restore KV state from file
    fun generateContinue(prompt: String): String      // generate using existing KV cache

    // Concurrent sessions — each session owns an isolated KV cache; model weights are shared
    fun createSession(name: String = ""): LlamaSession? // null on WASM (not supported)

    // Generation parameters (applied on next generate call)
    fun updateGenerateParams(
        temperature: Float,       // randomness (0.0–2.0)
        maxTokens: Int,           // max output tokens
        topP: Float,              // nucleus sampling threshold
        topK: Int,                // top-k sampling
        repeatPenalty: Float,     // penalty for repeated tokens
        contextLength: Int,       // KV context window size (requires model reload)
        numThreads: Int,          // CPU threads for inference
        useMmap: Boolean,         // memory-map model weights (requires model reload)
        flashAttention: Boolean,  // enable Flash Attention (requires model reload)
        batchSize: Int,           // token batch size for prompt processing (requires model reload)
    )

    fun nativeCancelGenerate()                        // cancel ongoing generation

    // Multi-Token Prediction (MTP) — speculative drafting
    fun initMtp(modelPath: String, draftLen: Int = 3): Boolean  // enable MTP; same GGUF as generation model
    fun shutdownMtp()                                 // disable MTP; generation continues without it
}

interface GenStream {
    fun onDelta(text: String)
    fun onComplete()
    fun onError(message: String)
}

// Concurrent session handle — created via LlamaBridge.createSession()
expect class LlamaSession {
    val name: String                                  // human-readable label assigned at creation
    fun stream(prompt: String, callback: GenStream)  // run inference in this session's context
    fun cancel()                                     // cancel the in-progress stream
    fun close()                                      // release native KV cache resources
}

// Generate and keep the KV state in memory
LlamaBridge.generate("Tell me about Kotlin.")

// Save the KV state to disk
LlamaBridge.sessionSave("/path/to/session.bin")

// ... later or in a new process ...

// Restore state and continue from where you left off
LlamaBridge.sessionLoad("/path/to/session.bin")
val continuation = LlamaBridge.generateContinue("What about multiplatform support?")

// Reset state without unloading the model
LlamaBridge.sessionReset()

// Load the model once
LlamaBridge.initGenerateModel(modelPath)

// Create two independent sessions with human-readable names
val sessionA = LlamaBridge.createSession(name = "Agent A") ?: error("Session creation failed")
val sessionB = LlamaBridge.createSession(name = "Agent B") ?: error("Session creation failed")

// Run them concurrently (e.g. launch in separate coroutines)
launch { sessionA.stream("Tell me about Kotlin.", callback = agentACallback) }
launch { sessionB.stream("Explain coroutines.", callback = agentBCallback) }

// Cancel an in-progress session
sessionA.cancel()

// Always close sessions when done to free native KV cache memory
sessionA.close()
sessionB.close()

val modelPath = LlamaBridge.getModelPath("Qwen3.5-1.7B-Instruct-Q4_K_M.gguf")

// 1. Load the trunk model as usual
LlamaBridge.initGenerateModel(modelPath)

// 2. Enable MTP — pass the same GGUF, the library loads only the MTP layers
//    draftLen: max speculative tokens per step (1–8 recommended, default 3)
val mtpReady = LlamaBridge.initMtp(modelPath, draftLen = 3)
if (!mtpReady) {
    println("Model does not contain MTP layers — running without speculative drafting.")
}

// 3. Use the API exactly as before — MTP is transparent
LlamaBridge.generateStream(
    prompt = prompt,
    callback = object : GenStream {
        override fun onDelta(text: String) { print(text) }
        override fun onComplete()          { println("\n[done]") }
        override fun onError(msg: String)  { println("Error: $msg") }
    }
)

// 4. Optionally disable MTP at runtime (trunk model stays loaded)
LlamaBridge.shutdownMtp()

LlamaBridge.initGenerateModel(modelPath)

when (LlamaBridge.getModelFinetuneType()?.lowercase()) {
    "instruct", "chat" -> { /* proceed normally */ }
    else -> showWarning("This appears to be a base model. For best results, use an instruction-tuned model.")
}

// Build a multi-turn conversation prompt using the model's own template
val prompt = LlamaBridge.applyChatTemplate(
    messages = listOf(
        "system" to "You are a helpful assistant.",
        "user"   to "What is Kotlin Multiplatform?",
    ),
    addAssistantPrefix = true          // appends the assistant-turn prefix so the model starts generating
)

if (prompt != null) {
    val response = LlamaBridge.generate(prompt)
    println(response)
}

// Inspect the raw template if needed (e.g. for debugging or custom rendering)
val templateString = LlamaBridge.getModelChatTemplate()
println(templateString)

object WhisperBridge {
    /** Returns a platform-specific absolute path for the model filename. */
    fun getModelPath(modelFileName: String): String

    /** Loads the model at [modelPath]. Returns true if loaded. */
    fun initModel(modelPath: String): Boolean

    /**
     * Transcribes a WAV file and returns text.
     * Tip: record WAV as 16 kHz, mono, 16-bit PCM for best compatibility.
     *
     * @param initialPrompt Optional text prepended to the decoder input (up to 224 tokens).
     *   Use it to bias transcription toward domain-specific vocabulary (e.g. medical terms).
     */
    fun transcribeWav(wavPath: String, language: String? = null, initialPrompt: String? = null): String

    /** Frees native resources. */
    fun release()
}

import com.llamatik.library.platform.WhisperBridge

val modelPath = WhisperBridge.getModelPath("ggml-tiny-q8_0.bin")

// 1) Init once (e.g. app start)
WhisperBridge.initModel(modelPath)

// 2) Record to a WAV file (16kHz mono PCM16) using your own recorder
val wavPath: String = "/path/to/recording.wav"

// 3) Transcribe
val text = WhisperBridge.transcribeWav(wavPath, language = null).trim()
println(text)

// 4) Optional: release on app shutdown
WhisperBridge.release()

object StableDiffusionBridge {

    /** Returns absolute model path (copied from assets/bundle if needed). */
    fun getModelPath(modelFileName: String): String

    /**
     * Loads the Stable Diffusion model.
     * @param threads CPU threads to use; -1 lets the backend decide.
     */
    fun initModel(modelPath: String, threads: Int = -1): Boolean

    /**
     * Text-to-image generation. Returns raw RGBA pixels (width * height * 4 bytes).
     * Returns an empty array on failure.
     */
    fun txt2img(
        prompt: String,
        negativePrompt: String? = null,
        width: Int = 512,
        height: Int = 512,
        steps: Int = 20,
        cfgScale: Float = 7.0f,
        seed: Long = -1L,
    ): ByteArray

    /**
     * Image-to-image generation. Starts from [initImageRgba] and steers it with [prompt].
     * [strength] controls how much the source image is preserved (0.0 = unchanged, 1.0 = ignored).
     * Returns raw RGBA pixels (width * height * 4 bytes). Returns an empty array on failure.
     */
    fun img2img(
        initImageRgba: ByteArray,
        initImageW: Int,
        initImageH: Int,
        prompt: String,
        negativePrompt: String? = null,
        width: Int = 512,
        height: Int = 512,
        steps: Int = 20,
        cfgScale: Float = 7.0f,
        strength: Float = 0.75f,
        seed: Long = -1L,
    ): ByteArray

    /** Releases native resources. */
    fun release()
}

import com.llamatik.library.platform.StableDiffusionBridge

val modelPath = StableDiffusionBridge.getModelPath("dreamshaper.safetensors")
StableDiffusionBridge.initModel(modelPath, threads = 4)

val rgba = StableDiffusionBridge.txt2img(
    prompt = "A cyberpunk llama in neon Tokyo",
    negativePrompt = "blurry, low quality",
    width = 512,
    height = 512,
    steps = 20,
    cfgScale = 7.0f,
    seed = 42L,
)
// Convert rgba (ByteArray, width*height*4) to a platform Bitmap / UIImage / BufferedImage

val sourceRgba: ByteArray = /* existing RGBA image bytes */

val rgba = StableDiffusionBridge.img2img(
    initImageRgba = sourceRgba,
    initImageW = 512,
    initImageH = 512,
    prompt = "The same scene as a watercolor painting",
    negativePrompt = "low quality",
    strength = 0.75f,
    seed = 42L,
)

object MultimodalBridge {
    /**
     * Load the vision model and its multimodal projector (mmproj) side-by-side.
     * Both files must be available on disk before calling this.
     *
     * @param modelPath  Absolute path to the GGUF vision model.
     * @param mmprojPath Absolute path to the GGUF mmproj file.
     * @return true on success.
     */
    fun initModel(modelPath: String, mmprojPath: String): Boolean

    /**
     * Analyze an image given as raw bytes (JPEG/PNG/BMP), streaming the response
     * token by token via [callback].
     *
     * Must be called from a background thread/coroutine; blocks until generation completes.
     */
    fun analyzeImageBytesStream(imageBytes: ByteArray, prompt: String, callback: GenStream)

    /** Cancel an in-progress analyzeImageBytesStream call. */
    fun cancelAnalysis()

    /** Free all native resources (model, mmproj context, llama context). */
    fun release()
}

import com.llamatik.library.platform.MultimodalBridge

// 1) Init once — both model and mmproj must be downloaded first
val loaded = MultimodalBridge.initModel(
    modelPath  = "/path/to/SmolVLM-256M-Instruct-Q8_0.gguf",
    mmprojPath = "/path/to/mmproj-SmolVLM-256M-Instruct-f16.gguf"
)

// 2) Analyze an image (e.g. loaded from disk or camera)
val imageBytes: ByteArray = File("/path/to/photo.jpg").readBytes()

MultimodalBridge.analyzeImageBytesStream(
    imageBytes = imageBytes,
    prompt     = "Describe what you see in this image.",
    callback   = object : GenStream {
        override fun onDelta(text: String)   { print(text) }
        override fun onComplete()            { println("\n[done]") }
        override fun onError(message: String){ println("Error: $message") }
    }
)

// 3) Optional: cancel mid-stream
MultimodalBridge.cancelAnalysis()

// 4) Optional: release on app shutdown
MultimodalBridge.release()