// src/model.js
import { AutoModelForImageTextToText, AutoProcessor, env, TextStreamer, RawImage } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers/dist/transformers.min.js';

// --- Configuration ---
const MODEL_ID = 'mistralai/Ministral-3-3B-Instruct-2512-ONNX';
const DEVICE = 'webgpu';

// --- State ---
let model;
let processor;

// --- Public API ---
/**
 * Loads the model and processor, and sets up the environment.
 * @param {function(object): void} progressCallback - A function to call with progress updates.
 */
export async function loadModel(progressCallback) {
    env.allowLocalModels = false;

    // The wasm files are bundled with the library, but we need to specify the path
    // to them when using a CDN.
    env.backends.onnx.wasm.wasmPaths = 'https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/';
    
    // Set device
    env.backends.onnx.device = DEVICE;

    try {
        console.log('Loading model...');
        model = await AutoModelForImageTextToText.from_pretrained(MODEL_ID, {
            progress_callback: progressCallback,
            device: DEVICE,
        });
        console.log('Model loaded.');

        console.log('Loading processor...');
        processor = await AutoProcessor.from_pretrained(MODEL_ID, {
            progress_callback: progressCallback,
        });
        console.log('Processor loaded.');

    } catch (error) {
        console.error('Error loading model:', error);
        throw error; // Propagate the error to the caller
    }
}

/**
 * Runs inference with the loaded model.
 * @param {string} inputText - The user's input text.
 * @param {function(string): void} streamCallback - A function to call with each generated token.
 */
export async function runInference(inputText, streamCallback) {
    if (!model || !processor) {
        throw new Error('Model not loaded. Please call loadModel() first.');
    }

    const messages = [
        {
            role: 'user',
            content: [
                { type: 'image' },
                { type: 'text', text: inputText }
            ],
        },
    ];
    const prompt = processor.apply_chat_template(messages, { add_generation_prompt: true });

    // Create a 1x1 black RawImage
    const blackPixel = new Uint8ClampedArray([0, 0, 0, 255]); // RGBA
    const dummyImage = new RawImage(blackPixel, 1, 1, 4);

    // Pass the dummy image to the processor
    const inputs = await processor(dummyImage, prompt);

    let fullOutput = '';
    const streamer = new TextStreamer(processor.tokenizer, {
        skip_prompt: true,
        skip_special_tokens: true,
        callback_function: (text) => {
            fullOutput += text;
            streamCallback(fullOutput);
        },
    });

    await model.generate({
        ...inputs,
        max_new_tokens: 512,
        streamer: streamer,
        do_sample: true,
        temperature: 0.7,
        top_p: 0.95,
        top_k: 50,
        eos_token_id: processor.tokenizer.eos_token_id,
        pad_token_id: processor.tokenizer.pad_token_id,
    });
}