import type { CSSProperties, ReactNode } from 'react';
import React, { useCallback, useEffect, useState } from 'react';
import { defaultRangeExtractor, useVirtualizer } from '@tanstack/react-virtual';
import type { Range } from '@tanstack/react-virtual';

import { Combobox as CoreCombobox, Text as CoreText } from '@shoptet/ui-core-web';
import type { ComboboxOptionState } from '@shoptet/ui-core-web';

import { getDataListCellProps } from '../../DataList/DataListItemCell';
import { getDataListSectionHeaderProps } from '../../DataList/DataListSectionHeader';
import { getTypographyProps } from '../../Typography/getTypographyProps';

import type { SelectKey } from './SingleSelect';

/** A flattened listbox row: a group heading, or a single option. */
export type SelectRow<T extends object, V extends SelectKey> =
  | { kind: 'heading'; key: string; label: string }
  | { kind: 'option'; key: string; option: T; label: string; value: V; disabled: boolean };

/** Estimated row heights and inner spacing for the virtualizer. */
export interface SelectVirtualizerLayout {
  rowHeight: number;
  headingHeight: number;
  gap: number;
  padding: number;
}

const absoluteRowStyle = (start: number): CSSProperties => ({
  position: 'absolute',
  top: 0,
  left: 0,
  width: '100%',
  transform: `translateY(${start}px)`,
});

export interface SingleSelectVirtualizedListboxProps<T extends object, V extends SelectKey> {
  rows: Array<SelectRow<T, V>>;
  layout: SelectVirtualizerLayout;
  className: string;
  optionClassName: string;
  style?: CSSProperties;
  /**
   * CSS selector of the listbox's scrolling ancestor (the popover on desktop, the dialog body on
   * small screens). Resolved via `listbox.closest(scrollSelector)`.
   */
  scrollSelector: string;
  /** Index (into `rows`) of the active option, or `null`; force-mounted so its id stays in the DOM. */
  activeRowIndex: number | null;
  /** Index (into `rows`) of the selected option to reveal on open, or `null`. */
  selectedRowIndex: number | null;
  /** Renders a single option row's inner content (indicators + `renderOption`) from the option state. */
  renderOptionRow: (option: T, state: ComboboxOptionState<V>) => ReactNode;
}

/**
 * The virtualized `Combobox.Listbox`: a `role="listbox"` holding an inner `role="presentation"` sizing
 * element (its height is the full list size) with the windowed rows positioned absolutely on top of it.
 * Positioning + the measure ref go directly on each `Combobox.Option` / heading element, so no
 * non-presentation wrapper sits between the listbox and its options. Shared by the desktop popover and
 * the small-screen dialog search paths.
 */
export function SingleSelectVirtualizedListbox<T extends object, V extends SelectKey>({
  rows,
  layout,
  className,
  optionClassName,
  style,
  scrollSelector,
  activeRowIndex,
  selectedRowIndex,
  renderOptionRow,
  ...rest
}: SingleSelectVirtualizedListboxProps<T, V>) {
  rest satisfies Record<PropertyKey, never>;

  const [listboxElement, setListboxElement] = useState<HTMLDivElement | null>(null);
  const getScrollElement = useCallback(
    () => listboxElement?.closest<HTMLElement>(scrollSelector) ?? null,
    [listboxElement, scrollSelector]
  );

  const estimateSize = useCallback(
    (index: number) => (rows[index]?.kind === 'heading' ? layout.headingHeight : layout.rowHeight),
    [rows, layout.headingHeight, layout.rowHeight]
  );

  // Always mount the active row so `aria-activedescendant` references a live element even when the
  // active option has scrolled outside the rendered window.
  const rangeExtractor = useCallback(
    (range: Range) => {
      const indexes = new Set(defaultRangeExtractor(range));
      if (activeRowIndex !== null) indexes.add(activeRowIndex);
      return [...indexes].toSorted((a, b) => a - b);
    },
    [activeRowIndex]
  );

  const virtualizer = useVirtualizer({
    count: rows.length,
    getScrollElement,
    estimateSize,
    overscan: 8,
    gap: layout.gap,
    paddingStart: layout.padding,
    paddingEnd: layout.padding,
    rangeExtractor,
  });

  // Reveal the selected option on open. Keyed on `listboxElement` so it re-runs once the listbox
  // mounts (the callback ref is null on the first render): only then does the scroll element exist
  // and floating-ui has sized the popover. Skip until the scroll element has layout, then scroll —
  // and scroll again on the next frame, once dynamic row measurement has settled, since the first
  // `scrollToIndex` runs against estimated sizes and can land short. `align: 'start'` keeps the
  // revealed option near the field instead of at the far scroll edge.
  useEffect(() => {
    if (selectedRowIndex === null || listboxElement === null) return;
    const scrollElement = listboxElement.closest<HTMLElement>(scrollSelector);
    if (scrollElement === null || scrollElement.clientHeight === 0) return;

    const scroll = () => virtualizer.scrollToIndex(selectedRowIndex, { align: 'start' });
    scroll();
    // The first pass runs against estimated row sizes and can land short; re-issue on the next frame,
    // once dynamic measurement has settled.
    const frame = requestAnimationFrame(scroll);
    return () => cancelAnimationFrame(frame);
  }, [selectedRowIndex, listboxElement, scrollSelector, virtualizer]);

  // Keep the active option in view as keyboard navigation moves it (including outside the window).
  useEffect(() => {
    if (activeRowIndex !== null) {
      virtualizer.scrollToIndex(activeRowIndex, { align: 'auto' });
    }
  }, [activeRowIndex, virtualizer]);

  const virtualItems = virtualizer.getVirtualItems();

  // `measureElement` reads the row index off `data-index` on the measured node. `Combobox.Option`
  // exposes only `ref`, so set the attribute in a ref callback (before the ResizeObserver fires),
  // then hand the node to the virtualizer for dynamic measurement.
  const measureRef = useCallback(
    (index: number) => (node: HTMLElement | null) => {
      if (node) node.dataset.index = String(index);
      virtualizer.measureElement(node);
    },
    [virtualizer]
  );

  return (
    <CoreCombobox.Listbox ref={setListboxElement} className={className} style={style}>
      <div role='presentation' style={{ position: 'relative', width: '100%', height: virtualizer.getTotalSize() }}>
        {virtualItems.map(virtualItem => {
          const row = rows[virtualItem.index];
          if (row === undefined) return null;

          if (row.kind === 'heading') {
            return (
              <div
                key={row.key}
                role='presentation'
                ref={measureRef(virtualItem.index)}
                {...getDataListSectionHeaderProps({})}
                style={absoluteRowStyle(virtualItem.start)}
              >
                <CoreText {...getTypographyProps({ appearance: 'sm', ...getDataListCellProps({}) })}>
                  {row.label}
                </CoreText>
              </div>
            );
          }

          return (
            <CoreCombobox.Option<V>
              key={row.key}
              value={row.value}
              label={row.label}
              disabled={row.disabled}
              ref={measureRef(virtualItem.index)}
              className={optionClassName}
              style={absoluteRowStyle(virtualItem.start)}
            >
              {({ state }) => renderOptionRow(row.option, state)}
            </CoreCombobox.Option>
          );
        })}
      </div>
    </CoreCombobox.Listbox>
  );
}