import { useEffect, useRef, useState } from 'react';
import { APISuccess, ImageUtility } from './generativeFillUtils/ImageHandler';
import { BrushHandler } from './generativeFillUtils/BrushHandler';
import { IconButton, TextField } from '@mui/material';
import { CursorData, Point } from './generativeFillUtils/generativeFillInterfaces';
import { activeColor, canvasSize, eraserColor, freeformRenderSize } from './generativeFillUtils/generativeFillConstants';
import { PointerHandler } from './generativeFillUtils/PointerHandler';
import { BsBrush, BsEraser, BsX } from 'react-icons/bs';
import { AiOutlineUpload } from 'react-icons/ai';
import { CiUndo, CiRedo } from 'react-icons/ci';
import Buttons from './GenerativeFillButtons';
import React = require('react');
import './GenerativeFill.scss';
import { EditableText } from 'browndash-components';
import { MainView } from '../../MainView';
import { Doc } from '../../../../fields/Doc';
import { Networking } from '../../../Network';
import { Utils } from '../../../../Utils';
import { DocUtils, Docs } from '../../../documents/Documents';
import { NumCast } from '../../../../fields/Types';

/**
 * For images not 1024x1024 fill in the rest in solid black, or a
 * reflected version of the image.
 */

/**
 * TODO: Look into img onload, sometimes the canvas doesn't update properly
 */

enum BrushStyle {
    ADD,
    SUBTRACT,
    MARQUEE,
}

interface ImageEdit {
    imgElement: HTMLImageElement;
    parent: ImageEdit;
    children: ImageEdit[];
}

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

const GenerativeFill = ({ imageEditorOpen, imageEditorSource, imageRootDoc, addDoc }: GenerativeFillProps) => {
    const canvasRef = useRef<HTMLCanvasElement>(null);
    const canvasBackgroundRef = useRef<HTMLCanvasElement>(null);
    const drawingAreaRef = useRef<HTMLDivElement>(null);
    const fileRef = useRef<HTMLInputElement>(null);
    const [cursorData, setCursorData] = useState<CursorData>({
        x: 0,
        y: 0,
        width: 150,
    });
    const [isBrushing, setIsBrushing] = useState(false);
    const [canvasScale, setCanvasScale] = useState(0.5);
    const [edits, setEdits] = useState<string[]>([]);
    const [brushStyle, setBrushStyle] = useState<BrushStyle>(BrushStyle.ADD);
    const [input, setInput] = useState('');
    const [loading, setLoading] = useState(false);
    // used to store the current image loaded to the main canvas
    const currImg = useRef<HTMLImageElement | null>(null);
    const originalImg = useRef<HTMLImageElement | null>(null);
    // stores history of data urls
    const undoStack = useRef<string[]>([]);
    // stores redo stack
    const redoStack = useRef<string[]>([]);
    // will change later, for now, stores an array [<col (generation)>, <row (edit in generation)>]
    const freeformPosition = useRef<number[]>([0, 0]);

    // 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);
        } else {
            redoStack.current = [...redoStack.current, canvasRef.current.toDataURL()];
            const img = new Image();
            img.src = target;
            ImageUtility.drawImgToCanvas(img, canvasRef);
            undoStack.current = undoStack.current.slice(0, -1);
        }
    };

    const handleRedo = () => {
        // TODO: handle undo as well
        const target = redoStack.current[redoStack.current.length - 1];
        if (!target) {
        } else {
            const img = new Image();
            img.src = target;
            ImageUtility.drawImgToCanvas(img, canvasRef);
            redoStack.current = redoStack.current.slice(0, -1);
        }
    };

    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 = [];
        ImageUtility.drawImgToCanvas(currImg.current, canvasRef, true);
    };

    // 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 = (e: React.PointerEvent) => {
        const ctx = ImageUtility.getCanvasContext(canvasBackgroundRef);
        if (!ctx) return;
        if (!isBrushing) return;
        setIsBrushing(false);
    };

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

        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,
            };
            BrushHandler.createBrushPathOverlay(lastPoint, currPoint, cursorData.width / 2 / canvasScale, ctx, eraserColor, brushStyle === BrushStyle.SUBTRACT);
        };

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

    // first load
    useEffect(() => {
        if (!imageEditorSource || imageEditorSource === '') return;
        const img = new Image();
        img.src = imageEditorSource;
        ImageUtility.drawImgToCanvas(img, canvasRef);
        currImg.current = img;
        originalImg.current = img;
    }, [canvasRef, imageEditorSource]);

    // 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; // Adjust the scale factor as per your requirement
            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,
        }));
    };

    // File upload
    const uploadImg = (e: React.ChangeEvent<HTMLInputElement>) => {
        if (e.target.files) {
            const file = e.target.files[0];
            const image = new Image();
            const imgUrl = URL.createObjectURL(file);
            image.src = imgUrl;
            ImageUtility.drawImgToCanvas(image, canvasRef);
            currImg.current = image;
        }
    };

    // Get AI Edit
    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);
        try {
            const maskBlob = await ImageUtility.canvasToBlob(canvas);
            const imgBlob = await ImageUtility.canvasToBlob(ImageUtility.getCanvasImg(img));

            // const res = await ImageUtility.getEdit(imgBlob, maskBlob, input !== '' ? input + ' in the same style' : 'Fill in the image in the same style', 2);

            const res = await ImageUtility.mockGetEdit(img.src);
            const { urls } = res as APISuccess;
            const image = new Image();
            image.src = urls[0];
            setLoading(false);
            setEdits(urls);
            ImageUtility.drawImgToCanvas(image, canvasRef);
            currImg.current = image;
        } catch (err) {
            console.log(err);
        }
    };

    const onSave = async () => {
        if (!currImg.current || !imageRootDoc) return;
        try {
            const src = currImg.current.src;
            console.log(src);
            const [result] = await Networking.PostToServer('/uploadRemoteImage', { sources: [src] });
            const source = Utils.prepend(result.accessPaths.agnostic.client);
            console.log(source);
            const newImg = Docs.Create.ImageDocument(source, {
                x: NumCast(imageRootDoc.x) + NumCast(imageRootDoc._width) + 20,
                y: NumCast(imageRootDoc.y),
                _height: freeformRenderSize,
                _width: freeformRenderSize,
                data_nativeWidth: result.nativeWidth,
                data_nativeHeight: result.nativeHeight,
            });

            addDoc?.(newImg);
            // Create link between prompt and image
            DocUtils.MakeLink(imageRootDoc, newImg, { link_relationship: 'Image Edit' });
            console.log('done');
        } catch (err) {
            console.log(err);
        }
    };

    return (
        <div className="generativeFillContainer" style={{ display: imageEditorOpen ? 'flex' : 'none' }}>
            <div className="generativeFillControls">
                <h1>Generative Fill</h1>
                <div style={{ display: 'flex', alignItems: 'center', gap: '1.5rem' }}>
                    <Buttons canvasRef={canvasRef} currImg={currImg} getEdit={getEdit} loading={loading} onSave={onSave} onReset={handleReset} />
                    <IconButton
                        onClick={() => {
                            MainView.Instance.setImageEditorOpen(false);
                            MainView.Instance.setImageEditorSource('');
                            setEdits([]);
                        }}>
                        <BsX color={activeColor} />
                    </IconButton>
                </div>
            </div>
            {/* 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={canvasSize} height={canvasSize} style={{ transform: `scale(${canvasScale})` }} />
                <canvas ref={canvasBackgroundRef} width={canvasSize} height={canvasSize} style={{ transform: `scale(${canvasScale})` }} />
                <div
                    className="pointer"
                    style={{
                        left: cursorData.x,
                        top: cursorData.y,
                        width: cursorData.width,
                        height: cursorData.width,
                    }}>
                    <div className="innerPointer"></div>
                </div>
                {/* Icons */}
                <div className="iconContainer">
                    <input ref={fileRef} type="file" accept="image/*" onChange={uploadImg} style={{ display: 'none' }} />
                    <IconButton
                        onClick={() => {
                            if (fileRef.current) {
                                fileRef.current.click();
                            }
                        }}>
                        <AiOutlineUpload />
                    </IconButton>
                    <IconButton
                        onClick={() => {
                            setBrushStyle(BrushStyle.ADD);
                        }}>
                        <BsBrush color={brushStyle === BrushStyle.ADD ? activeColor : 'inherit'} />
                    </IconButton>
                    {/* <IconButton
                        onClick={() => {
                            setBrushStyle(BrushStyle.SUBTRACT);
                        }}>
                        <BsEraser color={brushStyle === BrushStyle.SUBTRACT ? activeColor : 'inherit'} />
                    </IconButton> */}
                    {/* Undo and Redo */}
                    <IconButton
                        onPointerDown={e => {
                            e.stopPropagation();
                            handleUndo();
                        }}
                        onPointerUp={e => {
                            e.stopPropagation();
                        }}>
                        <CiUndo />
                    </IconButton>
                    <IconButton
                        onPointerDown={e => {
                            e.stopPropagation();
                            handleRedo();
                        }}
                        onPointerUp={e => {
                            e.stopPropagation();
                        }}>
                        <CiRedo />
                    </IconButton>
                </div>
                {/* Edits box */}
                <div className="editsBox">
                    {edits.map(edit => (
                        <img
                            key={edit}
                            width={100}
                            height={100}
                            src={edit}
                            onClick={() => {
                                const img = new Image();
                                img.src = edit;
                                ImageUtility.drawImgToCanvas(img, canvasRef);
                                currImg.current = img;
                            }}
                        />
                    ))}
                    {edits.length > 0 && (
                        <div style={{ position: 'relative' }}>
                            <label
                                style={{
                                    position: 'absolute',
                                    bottom: 10,
                                    left: 10,
                                    color: '#ffffff',
                                    fontSize: '0.8rem',
                                    letterSpacing: '1px',
                                    textTransform: 'uppercase',
                                }}>
                                Original
                            </label>
                            <img
                                width={100}
                                height={100}
                                src={originalImg.current?.src}
                                style={{ cursor: 'pointer' }}
                                onClick={() => {
                                    if (!originalImg.current) return;
                                    const img = new Image();
                                    img.src = originalImg.current.src;
                                    ImageUtility.drawImgToCanvas(img, canvasRef);
                                    currImg.current = img;
                                }}
                            />
                        </div>
                    )}
                </div>
            </div>
            <div>
                <TextField
                    value={input}
                    onChange={e => setInput(e.target.value)}
                    disabled={isBrushing}
                    type="text"
                    label="Prompt"
                    placeholder="Prompt..."
                    InputLabelProps={{ style: { fontSize: '1.5rem' } }}
                    inputProps={{ style: { fontSize: '1.5rem' } }}
                    sx={{
                        backgroundColor: '#ffffff',
                        position: 'absolute',
                        bottom: '1rem',
                        transform: 'translateX(calc(50vw - 50%))',
                        width: 'calc(100vw - 4rem)',
                        scale: 1.2,
                    }}
                />
            </div>
        </div>
    );
};

export default GenerativeFill;