import { pipeline, env } from '@xenova/transformers';

console.log('[BG-Removal Worker] Script loaded.');

// Configuration
env.allowLocalModels = false;
env.useBrowserCache = true;
env.allowRemoteModels = true;

let segmenter: any = null;

/**
 * Singleton for the RMBG model
 */
const getSegmenter = async () => {
    if (!segmenter) {
        console.log('[BG-Removal Worker] Loading RMBG-1.4 model...');
        try {
            segmenter = await pipeline('image-segmentation', 'Xenova/RMBG-1.4', {
                quantized: true,
                // @ts-ignore
                device: 'webgpu',
            });
            console.log('[BG-Removal Worker] Model loaded successfully.');
        } catch (err) {
            console.warn('[BG-Removal Worker] WebGPU failed, falling back to WASM:', err);
            segmenter = await pipeline('image-segmentation', 'Xenova/RMBG-1.4', {
                quantized: true,
            });
            console.log('[BG-Removal Worker] Model loaded successfully (WASM).');
        }
    }
    return segmenter;
};

/**
 * Apply the alpha mask to the original image
 */
const applyMaskWithFeathering = async (originalBlob: Blob, mask: any): Promise<Blob> => {
    const bitmap = await createImageBitmap(originalBlob);
    const { width, height } = bitmap;

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

    // 1. Convert 1-channel mask (Alpha) to 4-channel RGBA
    const maskCanvas = new OffscreenCanvas(mask.width, mask.height);
    const maskCtx = maskCanvas.getContext('2d');
    const rgbaData = new Uint8ClampedArray(mask.width * mask.height * 4);
    for (let i = 0; i < mask.data.length; ++i) {
        const val = mask.data[i];
        rgbaData[i * 4] = 0;     // R
        rgbaData[i * 4 + 1] = 0; // G
        rgbaData[i * 4 + 2] = 0; // B
        rgbaData[i * 4 + 3] = val; // A
    }
    const maskImageData = new ImageData(rgbaData, mask.width, mask.height);
    maskCtx!.putImageData(maskImageData, 0, 0);

    // 2. Composite
    ctx.drawImage(maskCanvas, 0, 0, width, height);
    ctx.globalCompositeOperation = 'source-in';
    ctx.drawImage(bitmap, 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-Removal Worker] Received request for ${id}`);

    try {
        const model = await getSegmenter();
        const url = URL.createObjectURL(imageBlob);

        console.log('[BG-Removal Worker] Running inference...');
        const output = await model(url);

        console.log('[BG-Removal Worker] Applying mask...');
        const processedBlob = await applyMaskWithFeathering(imageBlob, output);

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