import { defineTrait } from '../define.js';

/**
 * `focus-restore` — capture the previously-focused element on connect,
 * restore focus to it on disconnect.
 *
 * The canonical "modal-close returns focus to the trigger" pattern. Pairs
 * with `focus-trap` (which holds Tab inside a container) + `portal` (which
 * moves the container into a top-layer root). focus-restore closes the
 * loop: when the container goes away, focus returns to where the user was
 * before it appeared.
 *
 * Without this trait, modals / dialogs / drawers leave focus dangling on
 * `<body>` after dismiss — keyboard users have to tab from the start of
 * the document to get back to the trigger they pressed.
 *
 * On connect:
 *   1. Captures `document.activeElement` as the restore target.
 *   2. Sets `data-focus-restore-active` on the host.
 *   3. (Optional) Moves focus per `data-focus-restore-on-mount`:
 *      - absent / `"none"`: capture only, don't move focus.
 *      - `"host"`: focus the host itself.
 *      - `"first-focusable"`: focus the first tabbable descendant.
 *
 * On disconnect:
 *   1. If the captured target is still in the document AND focusable,
 *      `el.focus({ preventScroll: true })` it.
 *   2. Otherwise fall back to `<body>` (or the host's parent if body is
 *      gone for some reason).
 *   3. Removes `data-focus-restore-active`.
 *   4. Dispatches `focus-restored` so consumers can observe.
 */

const FOCUSABLE_SELECTOR =
  'a[href], [role="button"][tabindex], input:not([disabled]), select:not([disabled]), textarea:not([disabled]), button:not([disabled]), [tabindex]:not([tabindex="-1"]), [contenteditable]';

function isFocusable(el) {
  if (!el || !(el instanceof Element)) return false;
  if (typeof el.focus !== 'function') return false;
  if (el.hasAttribute('disabled')) return false;
  // Element must still be connected to the document for focus to land.
  if (!el.isConnected) return false;
  return true;
}

function findFirstFocusable(host) {
  // Try descendants in document order.
  const candidates = host.querySelectorAll(FOCUSABLE_SELECTOR);
  for (const el of candidates) {
    if (!el.hasAttribute('disabled')) return el;
  }
  return null;
}

export const focusRestore = defineTrait({
  name: 'focus-restore',
  category: 'keyboard-navigation',
  description:
    'Capture previously-focused element on connect, restore focus on disconnect',
  attributes: ['data-focus-restore-active'],
  events: ['focus-restored'],
  config: ['data-focus-restore-on-mount'],
  setup({ host }) {
    // Capture BEFORE we move focus — otherwise step 3 below would set the
    // restore target to the host (or first-focusable) and the trait would
    // restore to itself on disconnect, defeating the contract.
    const restoreTarget = document.activeElement;

    host.setAttribute('data-focus-restore-active', '');

    const onMount = host.getAttribute('data-focus-restore-on-mount') || 'none';
    if (onMount === 'host') {
      // Host needs to be focusable for this to land — common pattern is
      // tabindex="-1" on the dialog container.
      try {
        host.focus({ preventScroll: true });
      } catch {
        /* host not focusable — silent no-op, matches focus-trap's tolerance */
      }
    } else if (onMount === 'first-focusable') {
      const first = findFirstFocusable(host);
      if (first) {
        try {
          first.focus({ preventScroll: true });
        } catch {
          /* swallow — graceful degrade */
        }
      }
    }
    // onMount === 'none' (or absent / unrecognized): capture only, don't
    // move focus. Pointer-opened surfaces shouldn't shift focus on mount.

    return () => {
      host.removeAttribute('data-focus-restore-active');

      let restoredTo = null;

      if (isFocusable(restoreTarget)) {
        try {
          restoreTarget.focus({ preventScroll: true });
          restoredTo = restoreTarget;
        } catch {
          /* fall through to fallback */
        }
      }

      if (!restoredTo) {
        // Fallback: focus body (or host's parent if body is somehow gone).
        const fallback =
          (document.body && document.body.isConnected ? document.body : null) ||
          host.parentElement ||
          null;
        if (fallback && typeof fallback.focus === 'function') {
          try {
            // <body> needs a tabindex to receive .focus() in some engines;
            // we don't mutate it (no trait should leave document-level
            // attribute residue). If focus doesn't land, the browser
            // gracefully drops to no-active — same as the no-trait baseline.
            fallback.focus({ preventScroll: true });
            restoredTo = fallback;
          } catch {
            /* nothing more we can do — accept the dangle */
          }
        }
      }

      host.dispatchEvent(
        new CustomEvent('focus-restored', {
          bubbles: true,
          detail: { restoredTo, capturedTarget: restoreTarget },
        }),
      );
    };
  },
});
