/*
 * Copyright 2020 Adobe. All rights reserved.
 * This file is licensed to you under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License. You may obtain a copy
 * of the License at http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under
 * the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
 * OF ANY KIND, either express or implied. See the License for the specific language
 * governing permissions and limitations under the License.
 */

import {createShadowTreeWalker, getOwnerDocument, getOwnerWindow, nodeContains} from '@react-aria/utils';
import {shadowDOM} from '@react-stately/flags';

const supportsInert = typeof HTMLElement !== 'undefined' && 'inert' in HTMLElement.prototype;

interface AriaHideOutsideOptions {
  root?: Element,
  shouldUseInert?: boolean
}

// Keeps a ref count of all hidden elements. Added to when hiding an element, and
// subtracted from when showing it again. When it reaches zero, aria-hidden is removed.
let refCountMap = new WeakMap<Element, number>();
interface ObserverWrapper {
  visibleNodes: Set<Element>,
  hiddenNodes: Set<Element>,
  observe: () => void,
  disconnect: () => void
}
let observerStack: Array<ObserverWrapper> = [];

/**
 * Hides all elements in the DOM outside the given targets from screen readers using aria-hidden,
 * and returns a function to revert these changes. In addition, changes to the DOM are watched
 * and new elements outside the targets are automatically hidden.
 * @param targets - The elements that should remain visible.
 * @param root - Nothing will be hidden above this element.
 * @returns - A function to restore all hidden elements.
 */
export function ariaHideOutside(targets: Element[], options?: AriaHideOutsideOptions | Element) {
  let windowObj = getOwnerWindow(targets?.[0]);
  let opts = options instanceof windowObj.Element ? {root: options} : options;
  let root = opts?.root ?? document.body;
  let shouldUseInert = opts?.shouldUseInert && supportsInert;
  let visibleNodes = new Set<Element>(targets);
  let hiddenNodes = new Set<Element>();

  let getHidden = (element: Element) => {
    return shouldUseInert && element instanceof windowObj.HTMLElement ? element.inert : element.getAttribute('aria-hidden') === 'true';
  };

  let setHidden = (element: Element, hidden: boolean) => {
    if (shouldUseInert && element instanceof windowObj.HTMLElement) {
      element.inert = hidden;
    } else if (hidden) {
      element.setAttribute('aria-hidden', 'true');
    } else {
      element.removeAttribute('aria-hidden');
      if (element instanceof windowObj.HTMLElement) {
        // We only ever call setHidden with hidden = false when the nodeCount is 1 aka
        // we are trying to make the element visible to screen readers again, so remove inert as well
        element.inert = false;
      }
    }
  };

  let shadowRootsToWatch = new Set<ShadowRoot>();
  if (shadowDOM()) {
    // find all shadow roots that are ancestors of the targets
    // traverse upwards until the root is reached
    for (let target of targets) {
      let node = target;
      while (node && node !== root) {
        let root = node.getRootNode();
        if ('shadowRoot' in root) {
          shadowRootsToWatch.add(root.shadowRoot as ShadowRoot);
        }
        node = root.parentNode as Element;
      }
    }
  }

  let walk = (root: Element) => {
    // Keep live announcer and top layer elements (e.g. toasts) visible.
    for (let element of root.querySelectorAll('[data-live-announcer], [data-react-aria-top-layer]')) {
      visibleNodes.add(element);
    }

    let acceptNode = (node: Element) => {
      // Skip this node and its children if it is one of the target nodes, or a live announcer.
      // Also skip children of already hidden nodes, as aria-hidden is recursive. An exception is
      // made for elements with role="row" since VoiceOver on iOS has issues hiding elements with role="row".
      // For that case we want to hide the cells inside as well (https://bugs.webkit.org/show_bug.cgi?id=222623).
      if (
        hiddenNodes.has(node) ||
        visibleNodes.has(node) ||
        (node.parentElement && hiddenNodes.has(node.parentElement) && node.parentElement.getAttribute('role') !== 'row')
      ) {
        return NodeFilter.FILTER_REJECT;
      }

      // Skip this node but continue to children if one of the targets is inside the node.
      for (let target of visibleNodes) {
        if (nodeContains(node, target)) {
          return NodeFilter.FILTER_SKIP;
        }
      }

      return NodeFilter.FILTER_ACCEPT;
    };

    let walker = createShadowTreeWalker(
      getOwnerDocument(root),
      root,
      NodeFilter.SHOW_ELEMENT,
      {acceptNode}
    );

    // TreeWalker does not include the root.
    let acceptRoot = acceptNode(root);
    if (acceptRoot === NodeFilter.FILTER_ACCEPT) {
      hide(root);
    }

    if (acceptRoot !== NodeFilter.FILTER_REJECT) {
      let node = walker.nextNode() as Element;
      while (node != null) {
        hide(node);
        node = walker.nextNode() as Element;
      }
    }
  };

  let hide = (node: Element) => {
    let refCount = refCountMap.get(node) ?? 0;

    // If already aria-hidden, and the ref count is zero, then this element
    // was already hidden and there's nothing for us to do.
    if (getHidden(node) && refCount === 0) {
      return;
    }

    if (refCount === 0) {
      setHidden(node, true);
    }

    hiddenNodes.add(node);
    refCountMap.set(node, refCount + 1);
  };

  // If there is already a MutationObserver listening from a previous call,
  // disconnect it so the new on takes over.
  if (observerStack.length) {
    observerStack[observerStack.length - 1].disconnect();
  }

  walk(root);

  let observer = new MutationObserver(changes => {
    for (let change of changes) {
      if (change.type !== 'childList') {
        continue;
      }

      // If the parent element of the added nodes is not within one of the targets,
      // and not already inside a hidden node, hide all of the new children.
      if (
        change.target.isConnected &&
        ![...visibleNodes, ...hiddenNodes].some((node) =>
          nodeContains(node, change.target)
        )
      ) {
        for (let node of change.addedNodes) {
          if (
            (node instanceof HTMLElement || node instanceof SVGElement) &&
            (node.dataset.liveAnnouncer === 'true' || node.dataset.reactAriaTopLayer === 'true')
          ) {
            visibleNodes.add(node);
          } else if (node instanceof Element) {
            walk(node);
          }
        }
      }

      if (shadowDOM()) {
        // if any of the observed shadow roots were removed, stop observing them
        for (let shadowRoot of shadowRootsToWatch) {
          if (!shadowRoot.isConnected) {
            observer.disconnect();
            break;
          }
        }
      }
    }
  });

  observer.observe(root, {childList: true, subtree: true});
  let shadowObservers = new Set<MutationObserver>();
  if (shadowDOM()) {
    for (let shadowRoot of shadowRootsToWatch) {
      // Disconnect single target instead of all https://github.com/whatwg/dom/issues/126
      let shadowObserver = new MutationObserver(changes => {
        for (let change of changes) {
          if (change.type !== 'childList') {
            continue;
          }

          // If the parent element of the added nodes is not within one of the targets,
          // and not already inside a hidden node, hide all of the new children.
          if (
            change.target.isConnected &&
            ![...visibleNodes, ...hiddenNodes].some((node) =>
              nodeContains(node, change.target)
            )
          ) {
            for (let node of change.addedNodes) {
              if (
                (node instanceof HTMLElement || node instanceof SVGElement) &&
                (node.dataset.liveAnnouncer === 'true' || node.dataset.reactAriaTopLayer === 'true')
              ) {
                visibleNodes.add(node);
              } else if (node instanceof Element) {
                walk(node);
              }
            }
          }

          if (shadowDOM()) {
            // if any of the observed shadow roots were removed, stop observing them
            for (let shadowRoot of shadowRootsToWatch) {
              if (!shadowRoot.isConnected) {
                observer.disconnect();
                break;
              }
            }
          }
        }
      });
      shadowObserver.observe(shadowRoot, {childList: true, subtree: true});
      shadowObservers.add(shadowObserver);
    }
  }

  let observerWrapper: ObserverWrapper = {
    visibleNodes,
    hiddenNodes,
    observe() {
      observer.observe(root, {childList: true, subtree: true});
    },
    disconnect() {
      observer.disconnect();
    }
  };

  observerStack.push(observerWrapper);

  return (): void => {
    observer.disconnect();
    if (shadowDOM()) {
      for (let shadowObserver of shadowObservers) {
        shadowObserver.disconnect();
      }
    }

    for (let node of hiddenNodes) {
      let count = refCountMap.get(node);
      if (count == null) {
        continue;
      }
      if (count === 1) {
        setHidden(node, false);
        refCountMap.delete(node);
      } else {
        refCountMap.set(node, count - 1);
      }
    }

    // Remove this observer from the stack, and start the previous one.
    if (observerWrapper === observerStack[observerStack.length - 1]) {
      observerStack.pop();
      if (observerStack.length) {
        observerStack[observerStack.length - 1].observe();
      }
    } else {
      observerStack.splice(observerStack.indexOf(observerWrapper), 1);
    }
  };
}

export function keepVisible(element: Element): (() => void) | undefined {
  let observer = observerStack[observerStack.length - 1];
  if (observer && !observer.visibleNodes.has(element)) {
    observer.visibleNodes.add(element);
    return () => {
      observer.visibleNodes.delete(element);
    };
  }
}