/**
 * Give a badge on ControlNet Accordion indicating total number of active 
 * units.
 * Make active unit's tab name green.
 */
const cnetAllUnits = new Map/* <Element, GradioTab> */();
const cnetAllAccordions = new Set();
onUiUpdate(() => {
    function childIndex(element) {
        // Get all child nodes of the parent
        let children = Array.from(element.parentNode.childNodes);

        // Filter out non-element nodes (like text nodes and comments)
        children = children.filter(child => child.nodeType === Node.ELEMENT_NODE);

        return children.indexOf(element);
    }

    class GradioTab {
        constructor(tab) {
            this.enabledCheckbox = tab.querySelector('.cnet-unit-enabled input');
            this.inputImage = tab.querySelector('.cnet-input-image-group .cnet-image input[type="file"]');
            const tabs = tab.parentNode;
            this.tabNav = tabs.querySelector('.tab-nav');
            this.tabIndex = childIndex(tab) - 1; // -1 because tab-nav is also at the same level.            

            this.attachEnabledButtonListener();
            this.attachTabNavChangeObserver();
            this.attachImageUploadListener();
        }

        getTabNavButton() {
            return this.tabNav.querySelector(`:nth-child(${this.tabIndex + 1})`);
        }

        applyActiveState() {
            const tabNavButton = this.getTabNavButton();
            if (!tabNavButton) return;

            if (this.enabledCheckbox.checked) {
                tabNavButton.classList.add('cnet-unit-active');
            } else {
                tabNavButton.classList.remove('cnet-unit-active');
            }
        }

        attachEnabledButtonListener() {
            this.enabledCheckbox.addEventListener('change', () => {
                this.applyActiveState();
            });
        }

        /**
         * Each time the active tab change, all tab nav buttons are cleared and
         * regenerated by gradio. So we need to reapply the active states on 
         * them.
         */
        attachTabNavChangeObserver() {
            const observer = new MutationObserver((mutationsList, observer) => {
                for (const mutation of mutationsList) {
                    if (mutation.type === 'childList') {
                        this.applyActiveState();
                    }
                }
            });
            observer.observe(this.tabNav, { childList: true });
        }

        attachImageUploadListener() {
            this.inputImage.addEventListener('change', (event) => {
                if (!event.target.files) return;
                if (!this.enabledCheckbox.checked)
                    this.enabledCheckbox.click();
            });
        }
    }

    gradioApp().querySelectorAll('.cnet-unit-tab').forEach(tab => {
        if (cnetAllUnits.has(tab)) return;
        cnetAllUnits.set(tab, new GradioTab(tab));
    });

    function getActiveUnitCount(checkboxes) {
        let activeUnitCount = 0;
        for (const checkbox of checkboxes) {
            if (checkbox.checked)
                activeUnitCount++;
        }
        return activeUnitCount;
    }

    gradioApp().querySelectorAll('#controlnet').forEach(accordion => {
        if (cnetAllAccordions.has(accordion)) return;
        const checkboxes = accordion.querySelectorAll('.cnet-unit-enabled input');
        if (!checkboxes) return;

        const span = accordion.querySelector('.label-wrap span');
        checkboxes.forEach(checkbox => {
            checkbox.addEventListener('change', () => {
                // Remove existing badge.
                if (span.childNodes.length !== 1) {
                    span.removeChild(span.lastChild);
                }
                // Add new badge if necessary.
                const activeUnitCount = getActiveUnitCount(checkboxes);
                if (activeUnitCount > 0) {
                    const div = document.createElement('div');
                    div.classList.add('cnet-badge');
                    div.classList.add('primary');
                    div.innerHTML = `${activeUnitCount} unit${activeUnitCount > 1 ? 's' : ''}`;
                    span.appendChild(div);
                }
            });
        });
        cnetAllAccordions.add(accordion);
    });
});