import { useState, useCallback, useRef } from 'react';
import { AutoModel, AutoProcessor, RawImage, env } from '@huggingface/transformers';

export function useBackgroundRemoval() {
  const [loading, setLoading] = useState(false);
  const [progress, setProgress] = useState(0);
  const [hasModel, setHasModel] = useState(false);
  const modelRef = useRef(null);
  const processorRef = useRef(null);

  const loadModel = useCallback(async () => {
    if (modelRef.current && processorRef.current) return true;

    setLoading(true);
    setProgress(0);

    try {
      // Reduce ONNX Runtime Web console noise (node assignment warnings, etc.).
      // Note: depending on ONNX Runtime build/version, some warnings may still appear.
      env.backends.onnx.logLevel = 'error';

      modelRef.current = await AutoModel.from_pretrained('briaai/RMBG-1.4', {
        dtype: 'q8',
        device: navigator.gpu ? 'webgpu' : 'wasm',
        progress_callback: (p) => {
          if (p.progress != null) setProgress(Math.round(p.progress * 50));
        },
      });

      processorRef.current = await AutoProcessor.from_pretrained('briaai/RMBG-1.4');

      setProgress(50);
      setLoading(false);
      setHasModel(true);
      return true;
    } catch (error) {
      console.error('Failed to load background removal model:', error);
      setLoading(false);
      return false;
    }
  }, []);

  const removeBackground = useCallback(async (imageSrc) => {
    if (!modelRef.current || !processorRef.current) {
      const loaded = await loadModel();
      if (!loaded) return null;
    }

    setLoading(true);
    setProgress(50);

    try {
      const image = await RawImage.fromURL(imageSrc);
      const { pixel_values } = await processorRef.current(image);

      setProgress(70);

      const { output } = await modelRef.current({ input: pixel_values });

      setProgress(90);

      const mask = await RawImage.fromTensor(
        output[0].mul(255).to('uint8')
      ).resize(image.width, image.height);

      const canvas = document.createElement('canvas');
      canvas.width = image.width;
      canvas.height = image.height;
      const ctx = canvas.getContext('2d');

      ctx.drawImage(await image.toCanvas(), 0, 0);

      const imageData = ctx.getImageData(0, 0, canvas.width, canvas.height);
      for (let i = 0; i < mask.data.length; i++) {
        imageData.data[i * 4 + 3] = mask.data[i];
      }
      ctx.putImageData(imageData, 0, 0);

      // Tighten the result to its visible (non-transparent) pixels so the
      // returned image's bounding box is the subject, not the original
      // photo's mostly-transparent canvas. Callers use the returned `bbox`
      // to shrink the element's transform handles around the remaining
      // content. ALPHA_THRESHOLD = 8 (≈3% opacity) skips near-invisible
      // model noise at edges without clipping visible details.
      const W = canvas.width;
      const H = canvas.height;
      const ALPHA_THRESHOLD = 8;
      const pixels = imageData.data;
      let minX = W;
      let minY = H;
      let maxX = -1;
      let maxY = -1;
      for (let y = 0; y < H; y++) {
        for (let x = 0; x < W; x++) {
          if (pixels[(y * W + x) * 4 + 3] > ALPHA_THRESHOLD) {
            if (x < minX) minX = x;
            if (x > maxX) maxX = x;
            if (y < minY) minY = y;
            if (y > maxY) maxY = y;
          }
        }
      }

      let resultCanvas = canvas;
      let bbox;
      if (maxX < 0) {
        // Entirely transparent — the model decided the whole image was
        // background. Return the (empty) full canvas rather than a 0×0
        // crop, so the caller can still produce a sensible element
        // update if it wants to.
        bbox = { sx: 0, sy: 0, sw: W, sh: H, srcWidth: W, srcHeight: H };
      } else {
        const cropW = maxX - minX + 1;
        const cropH = maxY - minY + 1;
        if (cropW !== W || cropH !== H) {
          resultCanvas = document.createElement('canvas');
          resultCanvas.width = cropW;
          resultCanvas.height = cropH;
          resultCanvas.getContext('2d').drawImage(canvas, -minX, -minY);
        }
        bbox = { sx: minX, sy: minY, sw: cropW, sh: cropH, srcWidth: W, srcHeight: H };
      }

      setProgress(100);
      setLoading(false);
      return {
        url: resultCanvas.toDataURL('image/png'),
        bbox,
      };
    } catch (error) {
      console.error('Background removal failed:', error);
      setLoading(false);
      return null;
    }
  }, [loadModel]);

  return {
    loading,
    progress,
    hasModel,
    loadModel,
    removeBackground,
  };
}
