/*
Part of the implementation is borrowed and modified from multidiffusion-upscaler-for-automatic1111,
publicly available at https://github.com/pkuliyi2015/multidiffusion-upscaler-for-automatic1111
*/

const BBOX_MAX_NUM = 16;
const BBOX_WARNING_SIZE = 1280;
const DEFAULT_X = 0.4;
const DEFAULT_Y = 0.4;
const DEFAULT_H = 0.2;
const DEFAULT_W = 0.2;

// ref: https://html-color.codes/
const COLOR_MAP = [
 ['#ff0000', 'rgba(255, 0, 0, 0.3)'], // red
 ['#ff9900', 'rgba(255, 153, 0, 0.3)'], // orange
 ['#ffff00', 'rgba(255, 255, 0, 0.3)'], // yellow
 ['#33cc33', 'rgba(51, 204, 51, 0.3)'], // green
 ['#33cccc', 'rgba(51, 204, 204, 0.3)'], // indigo
 ['#0066ff', 'rgba(0, 102, 255, 0.3)'], // blue
 ['#6600ff', 'rgba(102, 0, 255, 0.3)'], // purple
 ['#cc00cc', 'rgba(204, 0, 204, 0.3)'], // dark pink
 ['#ff6666', 'rgba(255, 102, 102, 0.3)'], // light red
 ['#ffcc66', 'rgba(255, 204, 102, 0.3)'], // light orange
 ['#99cc00', 'rgba(153, 204, 0, 0.3)'], // lime green
 ['#00cc99', 'rgba(0, 204, 153, 0.3)'], // teal
 ['#0099cc', 'rgba(0, 153, 204, 0.3)'], // steel blue
 ['#9933cc', 'rgba(153, 51, 204, 0.3)'], // lavender
 ['#ff3399', 'rgba(255, 51, 153, 0.3)'], // hot pink
 ['#996633', 'rgba(153, 102, 51, 0.3)'], // brown
];

const RESIZE_BORDER = 5;
const MOVE_BORDER = 5;

const t2i_bboxes = new Array(BBOX_MAX_NUM).fill(null);
const i2i_bboxes = new Array(BBOX_MAX_NUM).fill(null);

function gradioApp() {
 const elems = document.getElementsByTagName('gradio-app')
 const gradioShadowRoot = elems.length == 0 ? null : elems[0].shadowRoot
 return !!gradioShadowRoot ? gradioShadowRoot : document;
}

// ↓↓↓ called from gradio ↓↓↓

function onCreateT2IRefClick(overwrite) {
 let width, height;
 if (overwrite) {
 const overwriteInputs = gradioApp().querySelectorAll('#MD-overwrite-width-t2i input, #MD-overwrite-height-t2i input');
 width = parseInt(overwriteInputs[0].value);
 height = parseInt(overwriteInputs[2].value);
 } else {
 const sizeInputs = gradioApp().querySelectorAll('#txt2img_width input, #txt2img_height input');
 width = parseInt(sizeInputs[0].value);
 height = parseInt(sizeInputs[2].value);
 }

 if (isNaN(width)) width = 512;
 if (isNaN(height)) height = 512;

 // Concat it to string to bypass the gradio bug
 // 向黑恶势力低头
 return width.toString() + 'x' + height.toString();
}

function onCreateI2IRefClick() {
 const canvas = gradioApp().querySelector('#img2img_image img');
 return canvas.src;
}

function onBoxEnableClick(is_t2i, idx, enable) {
 let canvas = null;
 let bboxes = null;
 let locator = null;
 if (is_t2i) {
 // locator = () => gradioApp().querySelector('#MD-bbox-ref-t2i');
 locator = () => gradioApp().querySelector('#MD-bbox-rect-t2i');
 bboxes = t2i_bboxes;
 
 } else {
 locator = () => gradioApp().querySelector('#MD-bbox-ref-i2i');
 bboxes = i2i_bboxes;
 }
 ref_div = locator();
 canvas = ref_div.querySelector('img');
 if (!canvas) { return false; }

 if (enable) {
 // Check if the bounding box already exists
 if (!bboxes[idx]) {
 // Initialize bounding box
 const bbox = [DEFAULT_X, DEFAULT_Y, DEFAULT_W, DEFAULT_H];
 const colorMap = COLOR_MAP[idx % COLOR_MAP.length];
 const div = document.createElement('div');
 div.id = 'MD-bbox-' + (is_t2i ? 't2i-' : 'i2i-') + idx;
 div.style.left = '0px';
 div.style.top = '0px';
 div.style.width = '0px';
 div.style.height = '0px';
 div.style.position = 'absolute';
 div.style.border = '2px solid ' + colorMap[0];
 div.style.background = colorMap[1];
 div.style.zIndex = '900';
 div.style.display = 'none';
 // A text tip to warn the user if bbox is too large
 const tip = document.createElement('span');
 tip.id = 'MD-tip-' + (is_t2i ? 't2i-' : 'i2i-') + idx;
 tip.style.left = '50%';
 tip.style.top = '50%';
 tip.style.position = 'absolute';
 tip.style.transform = 'translate(-50%, -50%)';
 tip.style.fontSize = '12px';
 tip.style.fontWeight = 'bold';
 tip.style.textAlign = 'center';
 tip.style.color = colorMap[0];
 tip.style.zIndex = '901';
 tip.style.display = 'none';
 tip.innerHTML = 'Warning: Region very large!<br>Take care of VRAM usage!';
 div.appendChild(tip);
 div.addEventListener('mousedown', function (e) {
 if (e.button === 0) { onBoxMouseDown(e, is_t2i, idx); }
 });
 div.addEventListener('mousemove', function (e) {
 updateCursorStyle(e, is_t2i, idx);
 });

 const shower = function() { // insert to DOM if necessary
 if (!gradioApp().querySelector('#' + div.id)) {
 locator().appendChild(div);
 }
 }
 bboxes[idx] = [div, bbox, shower];
 }

 // Show the bounding box
 displayBox(canvas, is_t2i, bboxes[idx]);
 return true;
 } else {
 if (!bboxes[idx]) { return false; }
 const [div, bbox, shower] = bboxes[idx];
 div.style.display = 'none';
 }
 return false;
}

function onBoxChange(is_t2i, idx, what, v) {
 // This function handles all the changes of the bounding box
 // Including the rendering and python slider update
 let bboxes = null;
 let canvas = null;
 if (is_t2i) {
 bboxes = t2i_bboxes;
 canvas = gradioApp().querySelector('#MD-bbox-rect-t2i img');
 } else {
 bboxes = i2i_bboxes;
 canvas = gradioApp().querySelector('#MD-bbox-ref-i2i img');
 }
 if (!bboxes[idx] || !canvas) {
 switch (what) {
 case 'x': return DEFAULT_X;
 case 'y': return DEFAULT_Y;
 case 'w': return DEFAULT_W;
 case 'h': return DEFAULT_H;
 }
 }
 const [div, bbox, shower] = bboxes[idx];
 if (div.style.display === 'none') { return v; }

 // parse trigger
 switch (what) {
 case 'x': bbox[0] = v; break;
 case 'y': bbox[1] = v; break;
 case 'w': bbox[2] = v; break;
 case 'h': bbox[3] = v; break;
 }
 displayBox(canvas, is_t2i, bboxes[idx]);
 return v;
}

// ↓↓↓ called from js ↓↓↓

function getSeedInfo(is_t2i, id, current_seed) {
 const info_id = is_t2i ? '#html_info_txt2img' : '#html_info_img2img';
 const info_div = gradioApp().querySelector(info_id);
 try{
 current_seed = parseInt(current_seed);
 } catch(e) {
 current_seed = -1;
 }
 if (!info_div) return current_seed;
 let info = info_div.innerHTML;
 if (!info) return current_seed;
 // remove all html tags
 info = info.replace(/<[^>]*>/g, '');
 // Find a json string 'region control:' in the info
 // get its index
 idx = info.indexOf('Region control');
 if (idx == -1) return current_seed;
 // get the json string (detect the bracket)
 // find the first '{'
 let start_idx = info.indexOf('{', idx);
 let bracket = 1;
 let end_idx = start_idx + 1;
 while (bracket > 0 && end_idx < info.length) {
 if (info[end_idx] == '{') bracket++;
 if (info[end_idx] == '}') bracket--;
 end_idx++;
 }
 if (bracket > 0) {
 return current_seed;
 }
 // get the json string
 let json_str = info.substring(start_idx, end_idx);
 // replace the single quote to double quote
 json_str = json_str.replace(/'/g, '"');
 // replace python True to javascript true, False to false
 json_str = json_str.replace(/True/g, 'true');
 // parse the json string
 let json = JSON.parse(json_str);
 // get the seed if the region id is in the json
 const region_id = 'Region ' + id.toString();
 if (!(region_id in json)) return current_seed;
 const region = json[region_id];
 if (!('seed' in region)) return current_seed;
 let seed = region['seed'];
 try{
 seed = parseInt(seed);
 } catch(e) {
 return current_seed;
 }
 return seed;
}

function displayBox(canvas, is_t2i, bbox_info) {
 // check null input
 const [div, bbox, shower] = bbox_info;
 const [x, y, w, h] = bbox;
 if (!canvas || !div || x == null || y == null || w == null || h == null) { return; }

 // client: canvas widget display size
 // natural: content image real size
 let vpScale = Math.min(canvas.clientWidth / canvas.naturalWidth, canvas.clientHeight / canvas.naturalHeight);
 let canvasCenterX = canvas.clientWidth / 2;
 let canvasCenterY = canvas.clientHeight / 2;
 let scaledX = canvas.naturalWidth * vpScale;
 let scaledY = canvas.naturalHeight * vpScale;
 let viewRectLeft = canvasCenterX - scaledX / 2;
 let viewRectRight = canvasCenterX + scaledX / 2;
 let viewRectTop = canvasCenterY - scaledY / 2;
 let viewRectDown = canvasCenterY + scaledY / 2;

 let xDiv = viewRectLeft + scaledX * x;
 let yDiv = viewRectTop + scaledY * y;
 let wDiv = Math.min(scaledX * w, viewRectRight - xDiv);
 let hDiv = Math.min(scaledY * h, viewRectDown - yDiv);

 // Calculate warning bbox size
 let upscalerFactor = 1.0;
 if (!is_t2i) {
 const upscalerInput = parseFloat(gradioApp().querySelector('#MD-i2i-upscaler-factor input').value);
 if (!isNaN(upscalerInput)) upscalerFactor = upscalerInput;
 }
 let maxSize = BBOX_WARNING_SIZE / upscalerFactor * vpScale;
 let maxW = maxSize / scaledX;
 let maxH = maxSize / scaledY;
 if (w > maxW || h > maxH) {
 div.querySelector('span').style.display = 'block';
 } else {
 div.querySelector('span').style.display = 'none';
 }

 // update <div> when not equal
 div.style.left = xDiv + 'px';
 div.style.top = yDiv + 'px';
 div.style.width = wDiv + 'px';
 div.style.height = hDiv + 'px';
 div.style.display = 'block';

 // insert it to DOM if not appear
 shower();
}

function onBoxMouseDown(e, is_t2i, idx) {
 let bboxes = null;
 let canvas = null;
 if (is_t2i) {
 bboxes = t2i_bboxes;
 canvas = gradioApp().querySelector('#MD-bbox-rect-t2i img');
 } else {
 bboxes = i2i_bboxes;
 canvas = gradioApp().querySelector('#MD-bbox-ref-i2i img');
 }
 // Get the bounding box
 if (!canvas || !bboxes[idx]) { return; }
 const [div, bbox, shower] = bboxes[idx];

 // Check if the click is inside the bounding box
 const boxRect = div.getBoundingClientRect();
 let mouseX = e.clientX;
 let mouseY = e.clientY;

 const resizeLeft = mouseX >= boxRect.left && mouseX <= boxRect.left + RESIZE_BORDER;
 const resizeRight = mouseX >= boxRect.right - RESIZE_BORDER && mouseX <= boxRect.right;
 const resizeTop = mouseY >= boxRect.top && mouseY <= boxRect.top + RESIZE_BORDER;
 const resizeBottom = mouseY >= boxRect.bottom - RESIZE_BORDER && mouseY <= boxRect.bottom;

 const moveHorizontal = mouseX >= boxRect.left + MOVE_BORDER && mouseX <= boxRect.right - MOVE_BORDER;
 const moveVertical = mouseY >= boxRect.top + MOVE_BORDER && mouseY <= boxRect.bottom - MOVE_BORDER;

 if (!resizeLeft && !resizeRight && !resizeTop && !resizeBottom && !moveHorizontal && !moveVertical) { return; }

 const horizontalPivot = resizeLeft ? bbox[0] + bbox[2] : bbox[0];
 const verticalPivot = resizeTop ? bbox[1] + bbox[3] : bbox[1];

 // Canvas can be regarded as invariant during the drag operation
 // Calculate in advance to reduce overhead

 // Calculate viewport scale based on the current canvas size and the natural image size
 let vpScale = Math.min(canvas.clientWidth / canvas.naturalWidth, canvas.clientHeight / canvas.naturalHeight);
 let vpOffset = canvas.getBoundingClientRect();

 // Calculate scaled dimensions of the canvas
 let scaledX = canvas.naturalWidth * vpScale;
 let scaledY = canvas.naturalHeight * vpScale;

 // Calculate the canvas center and view rectangle coordinates
 let canvasCenterX = (vpOffset.left + window.scrollX) + canvas.clientWidth / 2;
 let canvasCenterY = (vpOffset.top + window.scrollY) + canvas.clientHeight / 2;
 let viewRectLeft = canvasCenterX - scaledX / 2 - window.scrollX;
 let viewRectRight = canvasCenterX + scaledX / 2 - window.scrollX;
 let viewRectTop = canvasCenterY - scaledY / 2 - window.scrollY;
 let viewRectDown = canvasCenterY + scaledY / 2 - window.scrollY;

 mouseX = Math.min(Math.max(mouseX, viewRectLeft), viewRectRight);
 mouseY = Math.min(Math.max(mouseY, viewRectTop), viewRectDown);
 
 //const accordion = gradioApp().querySelector(`#MD-accordion-${is_t2i ? 't2i' : 'i2i'}-${idx}`);
 const accordion = gradioApp().querySelector('#MD-tab-t2i');

 // Move or resize the bounding box on mousemove
 function onMouseMove(e) {
 // Prevent selecting anything irrelevant
 e.preventDefault();

 // Get the new mouse position
 let newMouseX = e.clientX;
 let newMouseY = e.clientY;

 // clamp the mouse position to the view rectangle
 newMouseX = Math.min(Math.max(newMouseX, viewRectLeft), viewRectRight);
 newMouseY = Math.min(Math.max(newMouseY, viewRectTop), viewRectDown);

 // Calculate the mouse movement delta
 const dx = (newMouseX - mouseX) / scaledX;
 const dy = (newMouseY - mouseY) / scaledY;

 // Update the mouse position
 mouseX = newMouseX;
 mouseY = newMouseY;

 // if no move just return
 if (dx === 0 && dy === 0) { return; }

 // Update the mouse position
 let [x, y, w, h] = bbox;
 if (moveHorizontal && moveVertical) {
 // If moving the bounding box
 x = Math.min(Math.max(x + dx, 0), 1 - w);
 y = Math.min(Math.max(y + dy, 0), 1 - h);
 } else {
 // If resizing the bounding box
 if (resizeLeft || resizeRight) {
 if (x < horizontalPivot) {
 if (dx <= w) {
 // If still within the left side of the pivot
 x = x + dx;
 w = w - dx;
 } else {
 // If crossing the pivot
 w = dx - w;
 x = horizontalPivot;
 }
 } else {
 if (w + dx < 0) {
 // If still within the right side of the pivot
 x = horizontalPivot + w + dx;
 w = - dx - w;
 } else {
 // If crossing the pivot
 x = horizontalPivot;
 w = w + dx;
 }
 }

 // Clamp the bounding box to the image
 if (x < 0) {
 w = w + x;
 x = 0;
 } else if (x + w > 1) {
 w = 1 - x;
 }
 }
 // Same as above, but for the vertical axis
 if (resizeTop || resizeBottom) {
 if (y < verticalPivot) {
 if (dy <= h) {
 y = y + dy;
 h = h - dy;
 } else {
 h = dy - h;
 y = verticalPivot;
 }
 } else {
 if (h + dy < 0) {
 y = verticalPivot + h + dy;
 h = - dy - h;
 } else {
 y = verticalPivot;
 h = h + dy;
 }
 }
 if (y < 0) {
 h = h + y;
 y = 0;
 } else if (y + h > 1) {
 h = 1 - y;
 }
 }
 }
 const [div, old_bbox, _] = bboxes[idx];

 // If all the values are the same, just return
 if (old_bbox[0] === x && old_bbox[1] === y && old_bbox[2] === w && old_bbox[3] === h) { return; }
 // else update the bbox
 const event = new Event('input');
 const coords = [x, y, w, h];
 // <del>The querySelector is not very efficient, so we query it once and reuse it</del>
 // caching will result gradio bugs that stucks bbox and cannot move & drag
 const sliderIds = ['x', 'y', 'w', 'h'];
 // We try to select the input sliders
 const sliderSelectors = sliderIds.map(id => `#MD-${is_t2i ? 't2i' : 'i2i'}-${idx}-${id} input`).join(', ');
 let sliderInputs = accordion.querySelectorAll(sliderSelectors);
 // alert(sliderInputs.length)
 if (sliderInputs.length == 0) { 
 // If we failed, the accordion is probably closed and sliders are removed in the dom, so we open it
 accordion.querySelector('.label-wrap').click();
 // and try again
 sliderInputs = accordion.querySelectorAll(sliderSelectors);
 // If we still failed, we just return
 if (sliderInputs.length == 0) { return; }
 }
 for (let i = 0; i < 4; i++) {
 if (old_bbox[i] !== coords[i]) {
 sliderInputs[2*i].value = coords[i];
 sliderInputs[2*i].dispatchEvent(event);
 }
 }
 }

 // Remove the mousemove and mouseup event listeners
 function onMouseUp() {
 document.removeEventListener('mousemove', onMouseMove);
 document.removeEventListener('mouseup', onMouseUp);
 }

 // Add the event listeners
 document.addEventListener('mousemove', onMouseMove);
 document.addEventListener('mouseup', onMouseUp);
}

function updateCursorStyle(e, is_t2i, idx) {
 // This function changes the cursor style when hovering over the bounding box
 const bboxes = is_t2i ? t2i_bboxes : i2i_bboxes;
 if (!bboxes[idx]) return;

 const div = bboxes[idx][0];
 const boxRect = div.getBoundingClientRect();
 const mouseX = e.clientX;
 const mouseY = e.clientY;

 const resizeLeft = mouseX >= boxRect.left && mouseX <= boxRect.left + RESIZE_BORDER;
 const resizeRight = mouseX >= boxRect.right - RESIZE_BORDER && mouseX <= boxRect.right;
 const resizeTop = mouseY >= boxRect.top && mouseY <= boxRect.top + RESIZE_BORDER;
 const resizeBottom = mouseY >= boxRect.bottom - RESIZE_BORDER && mouseY <= boxRect.bottom;

 if ((resizeLeft && resizeTop) || (resizeRight && resizeBottom)) {
 div.style.cursor = 'nwse-resize';
 } else if ((resizeLeft && resizeBottom) || (resizeRight && resizeTop)) {
 div.style.cursor = 'nesw-resize';
 } else if (resizeLeft || resizeRight) {
 div.style.cursor = 'ew-resize';
 } else if (resizeTop || resizeBottom) {
 div.style.cursor = 'ns-resize';
 } else {
 div.style.cursor = 'move';
 }
}

// ↓↓↓ auto called event listeners ↓↓↓

function updateBoxes(is_t2i) {
 // This function redraw all bounding boxes
 let bboxes = null;
 let canvas = null;
 if (is_t2i) {
 bboxes = t2i_bboxes;
 canvas = gradioApp().querySelector('#MD-bbox-rect-t2i img');
 } else {
 bboxes = i2i_bboxes;
 canvas = gradioApp().querySelector('#MD-bbox-ref-i2i img');
 }
 if (!canvas) return;

 for (let idx = 0; idx < bboxes.length; idx++) {
 if (!bboxes[idx]) continue;
 const [div, bbox, shower] = bboxes[idx];
 if (div.style.display === 'none') { return; }

 displayBox(canvas, is_t2i, bboxes[idx]);
 }
}

window.addEventListener('resize', _ => {
 updateBoxes(true);
 updateBoxes(false);
});

// ======== Gradio Bug Fix ========
// For Gradio versions > 3.16.0 and < 3.29.0, the accordion DOM will be deleted when it is closed.
// We need to judge the versions and listen to the accordion open event, rerender the bbox at that time.
// This silly bug fix is only for compatibility, we recommend to update the gradio version to 3.29.0 or higher.
try {
 const GRADIO_VERSIONS = window.gradio_config["version"].split(".");
 const gradio_major_version = parseInt(GRADIO_VERSIONS[0]);
 const gradio_minor_version = parseInt(GRADIO_VERSIONS[1]);
 if (gradio_major_version == 3 && gradio_minor_version > 16 && gradio_minor_version < 29) {
 let listener = e => {
 if (!e) { return; }
 if (!e.target) { return; }
 if (!e.target.classList) { return; }
 if (!e.target.classList.contains('label-wrap')) { return; }
 for (let tab of ['t2i', 'i2i']) {
 const div = gradioApp().querySelector('#MD-bbox-control-' + tab +' div.label-wrap');
 if (!div) { continue; }
 updateBoxes(tab === 't2i');
 }
 };
 window.addEventListener('DOMNodeInserted', listener);
 }
} catch (ignored) {
 // If the above code failed, the gradio version shouldn't be in the range of 3.16.0 to 3.29.0, so we just return.
}
// ======== Gradio Bug Fix ========