path: root/src/client/views/nodes/imageEditor/ImageEditor.tsx

/* eslint-disable no-use-before-define */
import { Checkbox, FormControlLabel, Slider, TextField } from '@mui/material';
import { Button, IconButton, Type } from '@dash/components';
import * as React from 'react';
import { useEffect, useRef, useState } from 'react';
import { CgClose } from 'react-icons/cg';
import { IoMdRedo, IoMdUndo } from 'react-icons/io';
import { ClientUtils } from '../../../../ClientUtils';
import { Doc, DocListCast } from '../../../../fields/Doc';
import { List } from '../../../../fields/List';
import { NumCast } from '../../../../fields/Types';
import { Networking } from '../../../Network';
import { DocUtils } from '../../../documents/DocUtils';
import { Docs } from '../../../documents/Documents';
import { CollectionDockingView } from '../../collections/CollectionDockingView';
import { CollectionFreeFormView } from '../../collections/collectionFreeForm';
import { ImageEditorData } from '../ImageBox';
import { OpenWhereMod } from '../OpenWhere';
import './ImageEditor.scss';
import { ApplyFuncButtons, ImageToolButton } from './ImageEditorButtons';
import { BrushHandler } from './imageEditorUtils/BrushHandler';
import { APISuccess, ImageUtility } from './imageEditorUtils/ImageHandler';
import { PointerHandler } from './imageEditorUtils/PointerHandler';
import { activeColor, bgColor, brushWidthOffset, canvasSize, eraserColor, freeformRenderSize, newCollectionSize, offsetDistanceY, offsetX } from './imageEditorUtils/imageEditorConstants';
import { CutMode, CursorData, ImageDimensions, ImageEditTool, ImageToolType, Point } from './imageEditorUtils/imageEditorInterfaces';
import { DocumentView } from '../DocumentView';
import { SettingsManager } from '../../../util/SettingsManager';
import { Upload } from '../../../../server/SharedMediaTypes';

interface GenerativeFillProps {
    imageEditorOpen: boolean;
    imageEditorSource: string;
    imageRootDoc: Doc | undefined;
    addDoc: ((doc: Doc | Doc[], annotationKey?: string) => boolean) | undefined;
}

// Added field on image doc: gen_fill_children: List of children Docs

/**
 * The image editor interface can be accessed by opening a document's context menu, then going to Options --> Open Image Editor.
 * The image editor supports various operations on images. Currently, there is a Generative Fill feature that allows users to erase
 * part of an image, add an optional prompt, and send this to GPT. GPT then returns a newly generated image that replaces the erased
 * portion based on the optional prompt. There is also an image cutting tool that allows users to cut images in different ways to
 * reshape the images, take out portions of images, and overall use them more creatively (see the header comment for cutImage() for more information).
 */
const ImageEditor = ({ imageEditorOpen, imageEditorSource, imageRootDoc, addDoc }: GenerativeFillProps) => {
    const canvasRef = useRef<HTMLCanvasElement>(null);
    const canvasBackgroundRef = useRef<HTMLCanvasElement>(null);
    const drawingAreaRef = useRef<HTMLDivElement>(null);
    const [cursorData, setCursorData] = useState<CursorData>({
        x: 0,
        y: 0,
        width: 150,
    });
    const [isBrushing, setIsBrushing] = useState(false);
    const [canvasScale, setCanvasScale] = useState(0.5);
    // format: array of [image source, corresponding image Doc]
    const [edits, setEdits] = useState<{ url: string; saveRes: Doc | undefined }[]>([]);
    const [edited, setEdited] = useState(false);
    const [isFirstDoc, setIsFirstDoc] = useState<boolean>(true);
    const [input, setInput] = useState('');
    const [loading, setLoading] = useState(false);
    const [canvasDims, setCanvasDims] = useState<ImageDimensions>({
        width: canvasSize,
        height: canvasSize,
    });
    const [cutType, setCutType] = useState<CutMode>(CutMode.IN);
    // whether to create a new collection or not
    const [isNewCollection, setIsNewCollection] = useState(true);
    // the current image in the main canvas
    const currImg = useRef<HTMLImageElement | null>(null);
    // the unedited version of each generation (parent)
    const originalImg = useRef<HTMLImageElement | null>(null);
    const originalDoc = useRef<Doc | null>(null);
    // stores history of data urls
    const undoStack = useRef<string[]>([]);
    // stores redo stack
    const redoStack = useRef<string[]>([]);

    // references to keep track of tree structure
    const newCollectionRef = useRef<Doc | null>(null);
    const parentDoc = useRef<Doc | null>(null);
    const childrenDocs = useRef<Doc[]>([]);

    // constants for image cutting
    const cutPts = useRef<Point[]>([]);

    /**
     *
     * @param type The new tool type we are changing to
     */
    const changeTool = (type: ImageToolType) => {
        setCurrToolType(type);
        setCursorData(prev => ({ ...prev, width: currTool().sliderDefault as number }));
    };
    // Undo and Redo
    const handleUndo = () => {
        const ctx = ImageUtility.getCanvasContext(canvasRef);
        if (!ctx || !currImg.current || !canvasRef.current) return;

        const target = undoStack.current[undoStack.current.length - 1];
        if (!target) {
            ImageUtility.drawImgToCanvas(currImg.current, canvasRef, canvasDims.width, canvasDims.height);
        } else {
            redoStack.current = [...redoStack.current, canvasRef.current.toDataURL()];
            const img = new Image();
            img.src = target;
            ImageUtility.drawImgToCanvas(img, canvasRef, canvasDims.width, canvasDims.height);
            undoStack.current = undoStack.current.slice(0, -1);
        }
    };

    const handleRedo = () => {
        const ctx = ImageUtility.getCanvasContext(canvasRef);
        if (!ctx || !currImg.current || !canvasRef.current) return;

        const target = redoStack.current[redoStack.current.length - 1];
        if (target) {
            undoStack.current = [...undoStack.current, canvasRef.current?.toDataURL()];
            const img = new Image();
            img.src = target;
            ImageUtility.drawImgToCanvas(img, canvasRef, canvasDims.width, canvasDims.height);
            redoStack.current = redoStack.current.slice(0, -1);
        }
    };

    // resets any erase strokes
    const handleReset = () => {
        if (!canvasRef.current || !currImg.current) return;
        const ctx = ImageUtility.getCanvasContext(canvasRef);
        if (!ctx) return;
        ctx.clearRect(0, 0, canvasSize, canvasSize);
        undoStack.current = [];
        redoStack.current = [];
        cutPts.current.length = 0;
        ImageUtility.drawImgToCanvas(currImg.current, canvasRef, canvasDims.width, canvasDims.height);
    };

    // initiate brushing
    const handlePointerDown = (e: React.PointerEvent) => {
        const canvas = canvasRef.current;
        if (!canvas) return;
        const ctx = ImageUtility.getCanvasContext(canvasRef);
        if (!ctx) return;

        undoStack.current = [...undoStack.current, canvasRef.current.toDataURL()];
        redoStack.current = [];

        setIsBrushing(true);
        const { x, y } = PointerHandler.getPointRelativeToElement(canvas, e, canvasScale);
        BrushHandler.brushCircleOverlay(x, y, cursorData.width / 2 / canvasScale, ctx, eraserColor /* , brushStyle === BrushStyle.SUBTRACT */);
    };

    // stop brushing, push to undo stack
    const handlePointerUp = () => {
        const ctx = ImageUtility.getCanvasContext(canvasBackgroundRef);
        if (!ctx) return;
        if (!isBrushing) return;
        setIsBrushing(false);
    };

    // handles brushing on pointer movement
    useEffect(() => {
        const canvas = canvasRef.current;
        if (!isBrushing || !canvas) return undefined;
        const ctx = ImageUtility.getCanvasContext(canvasRef);
        if (!ctx) return undefined;

        const handlePointerMove = (e: PointerEvent) => {
            const currPoint = PointerHandler.getPointRelativeToElement(canvas, e, canvasScale);
            const lastPoint: Point = {
                x: currPoint.x - e.movementX / canvasScale,
                y: currPoint.y - e.movementY / canvasScale,
            };
            const pts = BrushHandler.createBrushPathOverlay(lastPoint, currPoint, cursorData.width / 2 / canvasScale, ctx, eraserColor);
            cutPts.current.push(...pts);
        };

        drawingAreaRef.current?.addEventListener('pointermove', handlePointerMove);
        return () => drawingAreaRef.current?.removeEventListener('pointermove', handlePointerMove);
    }, [isBrushing]);

    // first load
    useEffect(() => {
        if (imageEditorSource && imageEditorSource) {
            ImageUtility.urlToBase64(imageEditorSource).then(res => {
                if (res) {
                    const img = new Image();
                    img.src = `data:image/png;base64,${res}`;
                    img.onload = () => {
                        currImg.current = img;
                        originalImg.current = img;
                        const imgWidth = img.naturalWidth;
                        const imgHeight = img.naturalHeight;
                        const scale = Math.min(canvasSize / imgWidth, canvasSize / imgHeight);
                        const width = imgWidth * scale;
                        const height = imgHeight * scale;
                        setCanvasDims({ width, height });
                    };
                }
            });
        }

        // cleanup
        return () => {
            setInput('');
            setEdited(false);
            newCollectionRef.current = null;
            parentDoc.current = null;
            childrenDocs.current = [];
            currImg.current = null;
            originalImg.current = null;
            originalDoc.current = null;
            undoStack.current = [];
            redoStack.current = [];
            ImageUtility.clearCanvas(canvasRef);
        };
    }, [canvasRef, imageEditorSource]);

    // once the appropriate dimensions are set, draw the image to the canvas
    useEffect(() => {
        if (!currImg.current) return;
        ImageUtility.drawImgToCanvas(currImg.current, canvasRef, canvasDims.width, canvasDims.height);
    }, [canvasDims]);

    // handles brush sizing
    useEffect(() => {
        const handleKeyPress = (e: KeyboardEvent) => {
            if (e.key === 'ArrowUp') {
                e.preventDefault();
                e.stopPropagation();
                setCursorData(data => ({ ...data, width: data.width + 5 }));
            } else if (e.key === 'ArrowDown') {
                e.preventDefault();
                e.stopPropagation();
                setCursorData(data => (data.width >= 20 ? { ...data, width: data.width - 5 } : data));
            }
        };
        window.addEventListener('keydown', handleKeyPress);
        return () => window.removeEventListener('keydown', handleKeyPress);
    }, []);

    // handle pinch zoom
    useEffect(() => {
        const handlePinch = (e: WheelEvent) => {
            e.preventDefault();
            e.stopPropagation();
            const delta = e.deltaY;
            const scaleFactor = delta > 0 ? 0.98 : 1.02;
            setCanvasScale(prevScale => prevScale * scaleFactor);
        };

        drawingAreaRef.current?.addEventListener('wheel', handlePinch, {
            passive: false,
        });
        return () => drawingAreaRef.current?.removeEventListener('wheel', handlePinch);
    }, [drawingAreaRef]);

    // updates the current position of the cursor
    const updateCursorData = (e: React.PointerEvent) => {
        const drawingArea = drawingAreaRef.current;
        if (!drawingArea) return;
        const { x, y } = PointerHandler.getPointRelativeToElement(drawingArea, e, 1);
        setCursorData(data => ({
            ...data,
            x,
            y,
        }));
    };

    // Get AI Edit for Generative Fill
    const getEdit = async () => {
        const img = currImg.current;
        if (!img) return;
        const canvas = canvasRef.current;
        if (!canvas) return;
        const ctx = ImageUtility.getCanvasContext(canvasRef);
        if (!ctx) return;
        setLoading(true);
        setEdited(true);
        try {
            const canvasOriginalImg = ImageUtility.getCanvasImg(img);
            if (!canvasOriginalImg) return;
            const canvasMask = ImageUtility.getCanvasMask(canvas, canvas);
            if (!canvasMask) return;
            const maskBlob = await ImageUtility.canvasToBlob(canvasMask);
            const imgBlob = await ImageUtility.canvasToBlob(canvasOriginalImg);
            const res = await ImageUtility.getEdit(imgBlob, maskBlob, input || 'Fill in the image in the same style', 2);
            if ((res as any).status == 'error') {
                alert((res as any).message);
            }

            // create first image
            if (!newCollectionRef.current) {
                createNewCollection();
            } else {
                childrenDocs.current = [];
            }
            if (!(originalImg.current && imageRootDoc)) return;
            // add the doc to the main freeform
            await createNewImgDoc(originalImg.current, true);
            originalImg.current = currImg.current;
            originalDoc.current = parentDoc.current;
            const { urls } = res as APISuccess;
            if (res.status !== 'error') {
                const imgUrls = await Promise.all(urls.map(url => ImageUtility.convertImgToCanvasUrl(url, canvasDims.width, canvasDims.height)));
                const imgRes = await Promise.all(
                    imgUrls.map(async url => {
                        const saveRes = await onSave(url);
                        return { url, saveRes };
                    })
                );
                setEdits(imgRes);
                const image = new Image();
                image.src = imgUrls[0];
                ImageUtility.drawImgToCanvas(image, canvasRef, canvasDims.width, canvasDims.height);
                currImg.current = image;
                parentDoc.current = imgRes[0].saveRes ?? null;
            }
        } catch (err) {
            console.log(err);
        }
        setLoading(false);
    };

    /**
     * This function performs image cutting based on the inputted BrushMode. There are currently four ways to cut images:
     *      1. By outlining the area that should be kept (BrushMode.IN)
     *      2. By outlining the area that should be removed (BrushMode.OUT)
     *      3. By drawing in the area that should be kept (where the image is brushed, the image will remain and everything else will be removed) (BrushMode.DRAW_IN)
     *      4. By drawing the area that she be removed, so this operates as an eraser (BrushMode.ERASE)
     * @param currCutType BrushMode enum that determines what kind of cutting operation to perform
     * @param firstDoc boolean for whether it's the first edited image. This is for positioning of the edited images when they render on the canvas.
     */
    const cutImage = async (currCutType: CutMode, brushWidth: number, prevEdits: { url: string; saveRes: Doc | undefined }[], firstDoc: boolean) => {
        const img = currImg.current;
        const canvas = canvasRef.current;
        if (!canvas || !img) return;
        const ctx = ImageUtility.getCanvasContext(canvasRef);
        if (!ctx) return;
        // get the original image
        const canvasOriginalImg = ImageUtility.getCanvasImg(img);
        if (!canvasOriginalImg) return;
        setLoading(true);
        const currPts = [...cutPts.current];
        if (currCutType !== CutMode.ERASE) handleReset(); // gets rid of the visible brush strokes (mostly needed for line_in) unless it's erasing (which depends on the brush strokes)
        let minX = img.width;
        let maxX = 0;
        let minY = img.height;
        let maxY = 0;
        // currPts is populated by the brush strokes' points, so this code is drawing a path along the points
        if (currPts.length) {
            ctx.beginPath();
            ctx.moveTo(currPts[0].x, currPts[0].y);
            for (let i = 0; i < currPts.length; i++) {
                ctx.lineTo(currPts[i].x, currPts[i].y);
                minX = Math.min(currPts[i].x, minX);
                minY = Math.min(currPts[i].y, minY);
                maxX = Math.max(currPts[i].x, maxX);
                maxY = Math.max(currPts[i].y, maxY);
            }
            switch (
                currCutType // use different canvas operations depending on the type of cutting we're applying
            ) {
                case CutMode.IN:
                    ctx.closePath();
                    ctx.globalCompositeOperation = 'destination-in';
                    ctx.fill();
                    break;
                case CutMode.OUT:
                    ctx.closePath();
                    ctx.globalCompositeOperation = 'destination-out';
                    ctx.fill();
                    break;
                case CutMode.DRAW_IN:
                    ctx.globalCompositeOperation = 'destination-in';
                    ctx.lineWidth = brushWidth + brushWidthOffset; // added offset because width gets cut off a little bit
                    ctx.stroke();
                    break;
            }
        }

        const url = canvas.toDataURL();
        if (!newCollectionRef.current) {
            createNewCollection();
        }

        const image = new Image();
        image.src = url;
        image.onload = async () => {
            let finalImg: HTMLImageElement | undefined = image;
            let finalImgURL: string = url;
            // crop the image for these brush modes to remove excess blank space around the image contents
            if (currCutType == CutMode.IN || currCutType == CutMode.DRAW_IN) {
                const croppedData = cropImage(image, Math.max(minX, 0), Math.min(maxX, image.width), Math.max(minY, 0), Math.min(maxY, image.height));
                finalImg = croppedData;
                finalImgURL = croppedData.src;
            }
            currImg.current = finalImg;
            const newImgDoc = await createNewImgDoc(finalImg, firstDoc);
            if (newImgDoc) {
                // set the image to transparent to remove the background / brushstrokes
                newImgDoc.$backgroundColor = 'transparent';
                newImgDoc.$disableMixBlend = true;
                if (firstDoc) setIsFirstDoc(false);
                setEdits([...prevEdits, { url: finalImgURL, saveRes: undefined }]);
            }
            setLoading(false);
            cutPts.current.length = 0;
        };
    };

    /**
     * Creates a new collection to put the image edits on. Adds to a new tab on the right if "Create New Collection" is checked.
     * @returns
     */
    const createNewCollection = () => {
        if (!isNewCollection && imageRootDoc) {
            // if the parent hasn't been set yet
            if (!parentDoc.current) parentDoc.current = imageRootDoc;
        } else {
            if (!(originalImg.current && imageRootDoc)) return;
            // create new collection and add it to the view
            newCollectionRef.current = Docs.Create.FreeformDocument([], {
                x: NumCast(imageRootDoc.x) + NumCast(imageRootDoc._width) + offsetX,
                y: NumCast(imageRootDoc.y),
                _width: newCollectionSize,
                _height: newCollectionSize,
                title: 'Image edit collection',
            });
            DocUtils.MakeLink(imageRootDoc, newCollectionRef.current, { link_relationship: 'Image Edit Version History' });
            // opening new tab
            CollectionDockingView.AddSplit(newCollectionRef.current, OpenWhereMod.right);
        }
    };

    /**
     * This function crops an image based on the inputted dimensions. This is used to automatically adjust the images that are
     * edited to be smaller than the original (i.e. for cutting into a small part of the image)
     */
    const cropImage = (image: HTMLImageElement, minX: number, maxX: number, minY: number, maxY: number) => {
        const croppedCanvas = document.createElement('canvas');
        const croppedCtx = croppedCanvas.getContext('2d');
        if (!croppedCtx) return image;
        const cropWidth = Math.abs(maxX - minX);
        const cropHeight = Math.abs(maxY - minY);
        croppedCanvas.width = cropWidth;
        croppedCanvas.height = cropHeight;
        croppedCtx.globalCompositeOperation = 'source-over';
        croppedCtx.clearRect(0, 0, cropWidth, cropHeight);
        croppedCtx.drawImage(image, minX, minY, cropWidth, cropHeight, 0, 0, cropWidth, cropHeight);
        const croppedURL = croppedCanvas.toDataURL();
        const croppedImage = new Image();
        croppedImage.src = croppedURL;
        return croppedImage;
    };

    // adjusts all the img positions to be aligned
    const adjustImgPositions = () => {
        if (!parentDoc.current) return;
        const startY = NumCast(parentDoc.current.y);
        const children = DocListCast(parentDoc.current.gen_fill_children);
        const len = children.length;
        const initialYPositions: number[] = [];
        for (let i = 0; i < len; i++) {
            initialYPositions.push(startY + i * offsetDistanceY);
        }
        children.forEach((doc, i) => {
            if (len % 2 === 1) {
                doc.y = initialYPositions[i] - Math.floor(len / 2) * offsetDistanceY;
            } else {
                doc.y = initialYPositions[i] - (len / 2 - 1 / 2) * offsetDistanceY;
            }
        });
    };

    // creates a new image document and returns its reference
    const createNewImgDoc = async (img: HTMLImageElement, firstDoc: boolean /*, parent?: Doc */): Promise<Doc | undefined> => {
        if (!imageRootDoc) return undefined;
        const { src } = img;
        const [result] = (await Networking.PostToServer('/uploadRemoteImage', { sources: [src] })) as Upload.ImageInformation[];
        const source = ClientUtils.prepend(result.accessPaths.agnostic.client);

        if (firstDoc) {
            const x = 0;
            const initialY = 0;
            const newImg = Docs.Create.ImageDocument(source, {
                x: x,
                y: initialY,
                _height: freeformRenderSize,
                _width: freeformRenderSize,
                data_nativeWidth: result.nativeWidth,
                data_nativeHeight: result.nativeHeight,
            });
            if (isNewCollection && newCollectionRef.current) {
                Doc.AddDocToList(newCollectionRef.current, undefined, newImg);
            } else {
                addDoc?.(newImg);
            }
            parentDoc.current = newImg;
            return newImg;
        }
        if (!parentDoc.current) return undefined;
        const x = NumCast(parentDoc.current.x) + freeformRenderSize + offsetX;
        const initialY = 0;

        const newImg = Docs.Create.ImageDocument(source, {
            x: x,
            y: initialY,
            _height: freeformRenderSize,
            _width: freeformRenderSize,
            data_nativeWidth: result.nativeWidth,
            data_nativeHeight: result.nativeHeight,
        });

        const parentList = DocListCast(parentDoc.current.gen_fill_children);
        if (parentList.length > 0) {
            parentList.push(newImg);
            parentDoc.current.gen_fill_children = new List<Doc>(parentList);
        } else {
            parentDoc.current.gen_fill_children = new List<Doc>([newImg]);
        }

        DocUtils.MakeLink(parentDoc.current, newImg, { link_relationship: `Image edit; Prompt: ${input}` });
        adjustImgPositions();

        if (isNewCollection && newCollectionRef.current) {
            Doc.AddDocToList(newCollectionRef.current, undefined, newImg);
        } else {
            addDoc?.(newImg);
        }
        return newImg;
    };

    // Saves an image to the collection
    const onSave = async (src: string) => {
        const img = new Image();
        img.src = src;
        if (!currImg.current || !originalImg.current || !imageRootDoc) return undefined;
        try {
            return await createNewImgDoc(img, false);
        } catch (err) {
            console.log(err);
        }
        return undefined;
    };

    // Closes the editor view
    const handleViewClose = () => {
        ImageEditorData.Open = false;
        ImageEditorData.Source = '';
        if (newCollectionRef.current) {
            DocumentView.addViewRenderedCb(newCollectionRef.current, dv => (dv.ComponentView as CollectionFreeFormView)?.fitContentOnce());
        }
        setEdits([]);
        setIsFirstDoc(true);
    };

    function currTool() {
        return imageEditTools.find(tool => tool.type === currToolType) ?? genFillTool;
    }

    // defines the tools and sets current tool
    const genFillTool: ImageEditTool = { type: ImageToolType.GenerativeFill, btnText: 'GET EDITS', icon: 'fill', applyFunc: getEdit, sliderMin: 25, sliderMax: 500, sliderDefault: 150 };
    const cutTool: ImageEditTool = { type: ImageToolType.Cut, btnText: 'CUT IMAGE', icon: 'scissors', applyFunc: cutImage, sliderMin: 1, sliderMax: 50, sliderDefault: 5 };
    const imageEditTools: ImageEditTool[] = [genFillTool, cutTool];
    const [currToolType, setCurrToolType] = useState<ImageToolType>(ImageToolType.GenerativeFill);

    // the top controls for making a new collection, resetting, and applying edits,
    function renderControls() {
        return (
            <div className="imageEditorTopBar">
                <h1>Image Editor</h1>
                {/* <IconButton text="Cut out" icon={<FontAwesomeIcon icon="scissors" />} /> */}
                <div className="imageEditorControls">
                    <FormControlLabel
                        control={
                            <Checkbox
                                // disable once edited has been clicked (doesn't make sense to change after first edit)
                                disabled={edited}
                                checked={isNewCollection}
                                onChange={() => setIsNewCollection(prev => !prev)}
                            />
                        }
                        label="Create New Collection"
                        labelPlacement="end"
                        sx={{ whiteSpace: 'nowrap' }}
                    />
                    <ApplyFuncButtons onClick={() => currTool().applyFunc(cutType, cursorData.width, edits, isFirstDoc)} loading={loading} onReset={handleReset} btnText={currTool().btnText} />
                    <IconButton color={activeColor} tooltip="close" icon={<CgClose size="16px" />} onClick={handleViewClose} />
                </div>
            </div>
        );
    }

    // the side icons including tool type, the slider, and undo/redo
    function renderSideIcons() {
        return (
            <div className="sideControlsContainer" style={{ backgroundColor: bgColor }}>
                <div className="sideControls">
                    <div className="imageToolsContainer">{imageEditTools.map(tool => ImageToolButton(tool, tool.type === currTool().type, changeTool))}</div>
                    {currTool().type == ImageToolType.Cut && (
                        <div className="cutToolsContainer">
                            <Button style={{ width: '100%' }} text="Keep in" type={Type.TERT} color={cutType == CutMode.IN ? SettingsManager.userColor : bgColor} onClick={() => setCutType(CutMode.IN)} />
                            <Button style={{ width: '100%' }} text="Keep out" type={Type.TERT} color={cutType == CutMode.OUT ? SettingsManager.userColor : bgColor} onClick={() => setCutType(CutMode.OUT)} />
                            <Button style={{ width: '100%' }} text="Draw in" type={Type.TERT} color={cutType == CutMode.DRAW_IN ? SettingsManager.userColor : bgColor} onClick={() => setCutType(CutMode.DRAW_IN)} />
                            <Button style={{ width: '100%' }} text="Erase" type={Type.TERT} color={cutType == CutMode.ERASE ? SettingsManager.userColor : bgColor} onClick={() => setCutType(CutMode.ERASE)} />
                        </div>
                    )}
                    <div className="sliderContainer" onPointerDown={e => e.stopPropagation()}>
                        {currTool().type === ImageToolType.GenerativeFill && (
                            <Slider
                                sx={{
                                    '& input[type="range"]': {
                                        writingMode: 'vertical-lr',
                                        direction: 'rtl',
                                        // WebkitAppearance: 'slider-vertical',
                                    },
                                }}
                                orientation="vertical"
                                min={genFillTool.sliderMin}
                                max={genFillTool.sliderMax}
                                defaultValue={genFillTool.sliderDefault}
                                size="small"
                                valueLabelDisplay="auto"
                                onChange={(e, val) => setCursorData(prev => ({ ...prev, width: val as number }))}
                            />
                        )}
                        {currTool().type === ImageToolType.Cut && (
                            <Slider
                                sx={{
                                    '& input[type="range"]': {
                                        writingMode: 'vertical-lr',
                                        direction: 'rtl',
                                        // WebkitAppearance: 'slider-vertical',
                                    },
                                }}
                                orientation="vertical"
                                min={cutTool.sliderMin}
                                max={cutTool.sliderMax}
                                defaultValue={cutTool.sliderDefault}
                                size="small"
                                valueLabelDisplay="auto"
                                onChange={(e, val) => setCursorData(prev => ({ ...prev, width: val as number }))}
                            />
                        )}
                    </div>
                    {/* Undo and Redo */}
                    <div className="undoRedoContainer">
                        <IconButton
                            style={{ cursor: 'pointer' }}
                            onPointerDown={e => {
                                e.stopPropagation();
                                handleUndo();
                            }}
                            onPointerUp={e => e.stopPropagation()}
                            color={activeColor}
                            tooltip="Undo"
                            icon={<IoMdUndo />}
                        />
                        <IconButton
                            style={{ cursor: 'pointer' }}
                            onPointerDown={e => {
                                e.stopPropagation();
                                handleRedo();
                            }}
                            onPointerUp={e => e.stopPropagation()}
                            color={activeColor}
                            tooltip="Redo"
                            icon={<IoMdRedo />}
                        />
                    </div>
                </div>
            </div>
        );
    }

    // circular pointer for drawing/erasing
    function renderPointer() {
        return (
            <div
                className="pointer"
                style={{
                    left: cursorData.x,
                    top: cursorData.y,
                    width: cursorData.width,
                    height: cursorData.width,
                }}>
                <div className="innerPointer" />
            </div>
        );
    }

    // the previews for each edit
    function renderEditThumbnails() {
        return (
            <div className="editsBox">
                {edits.map(edit => (
                    <img
                        key={edit.url}
                        alt="image edits"
                        width={75}
                        src={edit.url}
                        onClick={async () => {
                            const img = new Image();
                            img.src = edit.url;
                            ImageUtility.drawImgToCanvas(img, canvasRef, img.width, img.height);
                            currImg.current = img;
                            parentDoc.current = edit.saveRes ?? null;
                        }}
                    />
                ))}
                {/* Original img thumbnail */}
                {edits.length > 0 && (
                    <div style={{ position: 'relative' }}>
                        <label className="originalImageLabel">Original</label>
                        <img
                            alt="image stuff"
                            width={75}
                            src={originalImg.current?.src}
                            onClick={() => {
                                if (!originalImg.current) return;
                                const img = new Image();
                                img.src = originalImg.current.src;
                                ImageUtility.drawImgToCanvas(img, canvasRef, canvasDims.width, canvasDims.height);
                                currImg.current = img;
                                if (!parentDoc.current) parentDoc.current = originalDoc.current;
                            }}
                        />
                    </div>
                )}
            </div>
        );
    }

    // the prompt box for generative fill
    function renderPromptBox() {
        return (
            <div>
                <TextField
                    value={input}
                    onChange={e => setInput(e.target.value)}
                    disabled={isBrushing}
                    type="text"
                    label="Prompt"
                    placeholder="Prompt..."
                    InputLabelProps={{ style: { fontSize: '16px' } }}
                    inputProps={{ style: { fontSize: '16px' } }}
                    sx={{
                        backgroundColor: '#ffffff',
                        position: 'absolute',
                        bottom: '16px',
                        transform: 'translateX(calc(50vw - 50%))',
                        width: 'calc(100vw - 64px)',
                    }}
                />
            </div>
        );
    }

    return (
        <div className="imageEditorContainer" style={{ display: imageEditorOpen ? 'flex' : 'none' }}>
            {renderControls()}
            {/* Main canvas for editing */}
            <div
                className="drawingArea" // this only works if pointerevents: none is set on the custom pointer
                ref={drawingAreaRef}
                onPointerOver={updateCursorData}
                onPointerMove={updateCursorData}
                onPointerDown={handlePointerDown}
                onPointerUp={handlePointerUp}>
                <canvas ref={canvasRef} width={canvasDims.width} height={canvasDims.height} style={{ transform: `scale(${canvasScale})` }} />
                <canvas ref={canvasBackgroundRef} width={canvasDims.width} height={canvasDims.height} style={{ transform: `scale(${canvasScale})` }} />
                {renderPointer()}
                {renderSideIcons()}
                {renderEditThumbnails()}
            </div>
            {currTool().type === ImageToolType.GenerativeFill && renderPromptBox()}
        </div>
    );
};

export default ImageEditor;