// Background Removal Worker using briaai/RMBG-1.4
// Using @huggingface/transformers v3

import { AutoModel, AutoProcessor, RawImage, env } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.5.1';

console.log('[BG-Processor Worker] Script loaded from /public');

// Configuration
env.allowLocalModels = false;
env.useBrowserCache = true;
// Force WASM backend (more compatible)
env.backends.onnx.wasm.proxy = false;

let model = null;
let processor = null;

/**
 * Load the RMBG-1.4 model (WASM only for compatibility)
 */
const loadModel = async () => {
    if (!model) {
        console.log('[BG-Processor Worker] Loading briaai/RMBG-1.4 model (WASM)...');
        model = await AutoModel.from_pretrained('briaai/RMBG-1.4', {
            device: 'wasm',
            dtype: 'fp32',
        });
        processor = await AutoProcessor.from_pretrained('briaai/RMBG-1.4');
        console.log('[BG-Processor Worker] Model loaded successfully.');
    }
    return { model, processor };
};

/**
 * Apply the alpha mask to the original image
 */
const applyMask = async (originalBlob, maskData, width, height) => {
    const bitmap = await createImageBitmap(originalBlob);

    const canvas = new OffscreenCanvas(bitmap.width, bitmap.height);
    const ctx = canvas.getContext('2d');
    if (!ctx) throw new Error("No Canvas context");

    // Draw original image
    ctx.drawImage(bitmap, 0, 0);

    // Get image data
    const imageData = ctx.getImageData(0, 0, bitmap.width, bitmap.height);
    const data = imageData.data;

    // Create mask canvas at model output size
    const maskCanvas = new OffscreenCanvas(width, height);
    const maskCtx = maskCanvas.getContext('2d');
    const maskImageData = maskCtx.createImageData(width, height);

    // Convert model output to grayscale image
    for (let i = 0; i < maskData.length; i++) {
        const val = Math.round(Math.max(0, Math.min(1, maskData[i])) * 255);
        maskImageData.data[i * 4] = val;
        maskImageData.data[i * 4 + 1] = val;
        maskImageData.data[i * 4 + 2] = val;
        maskImageData.data[i * 4 + 3] = 255;
    }
    maskCtx.putImageData(maskImageData, 0, 0);

    // Scale mask to original size
    const scaledMaskCanvas = new OffscreenCanvas(bitmap.width, bitmap.height);
    const scaledMaskCtx = scaledMaskCanvas.getContext('2d');
    scaledMaskCtx.drawImage(maskCanvas, 0, 0, bitmap.width, bitmap.height);
    const scaledMaskData = scaledMaskCtx.getImageData(0, 0, bitmap.width, bitmap.height);

    // Apply mask as alpha
    for (let i = 0; i < data.length; i += 4) {
        data[i + 3] = scaledMaskData.data[i]; // Use R channel as alpha
    }

    ctx.putImageData(imageData, 0, 0);
    return await canvas.convertToBlob({ type: 'image/png' });
};

self.onmessage = async (e) => {
    const { type, id, imageBlob } = e.data;

    if (type === 'ping') {
        self.postMessage({ type: 'pong' });
        return;
    }

    if (!imageBlob) return;

    console.log(`[BG-Processor Worker] Received request for ${id}`);

    try {
        const { model, processor } = await loadModel();

        // Convert blob to RawImage
        const url = URL.createObjectURL(imageBlob);
        const image = await RawImage.fromURL(url);
        URL.revokeObjectURL(url);

        console.log('[BG-Processor Worker] Running inference...');

        // Process image
        const { pixel_values } = await processor(image);

        // Run model
        const { output } = await model({ input: pixel_values });

        // Get mask data - output is a Tensor
        const maskData = output.data;
        const [batch, channels, height, width] = output.dims;

        console.log(`[BG-Processor Worker] Mask dims: ${width}x${height}`);
        console.log('[BG-Processor Worker] Applying mask...');

        const processedBlob = await applyMask(imageBlob, maskData, width, height);

        self.postMessage({ id, status: 'success', blob: processedBlob });
        console.log(`[BG-Processor Worker] Successfully processed ${id}`);
    } catch (err) {
        console.error(`[BG-Processor Worker] Processing Error (${id}):`, err);
        self.postMessage({ id, status: 'error', error: err.message });
    }
};