import {
  component$,
  useSignal,
  useStore,
  useVisibleTask$,
} from "@builder.io/qwik";
import {
  UPDATE_SCROLL_EVENT,
  ACTION_ITEMS_LENGTH_CHANGE,
  createVirtualStore,
  UPDATE_VIRTUAL_STATE,
  UPDATE_SCROLL_END_EVENT,
  getScrollSize,
  ACTION_START_OFFSET_CHANGE,
  createScroller,
  createResizer,
  microtask,
  sort,
} from "../core";
import { ListItem } from "./ListItem";
import { getKey } from "./utils";
import type {
  VirtualizerProps,
  VirtualizerHandle,
  JSXNode,
  CustomContainerComponent,
} from "./types";

/**
 * Hook to get children elements and count
 */
const useChildren = (
  children: JSXNode | ((index: number) => JSXNode),
  renderCountProp?: number
): [(index: number) => JSXNode, number] => {
  if (typeof children === "function") {
    return [children, renderCountProp || 0];
  }

  // If children is an array, return function to access by index
  if (Array.isArray(children)) {
    return [(index: number) => children[index], children.length];
  }

  // Single child
  return [(index: number) => (index === 0 ? children : null), 1];
};

/**
 * Customizable list virtualizer for advanced usage. See {@link VirtualizerProps} and {@link VirtualizerHandle}.
 */
export const Virtualizer = component$<VirtualizerProps>((props: VirtualizerProps) => {
  const {
    children,
    count: renderCountProp,
    overscan,
    itemSize,
    shift,
    horizontal: horizontalProp,
    keepMounted,
    cache,
    startMargin = 0,
    ssrCount,
    as: Element = "div",
    item: ItemElement = "div",
    scrollRef,
    onScroll$: onScrollProp,
    onScrollEnd$: onScrollEndProp,
    ref,
  } = props;

  const [getElement, count] = useChildren(children, renderCountProp);

  const containerRef = useSignal<HTMLDivElement>();
  const isSSR = useSignal(!!ssrCount);

  // Create virtual store and related objects
  const virtualStore = useStore(() => {
    const _isHorizontal = !!horizontalProp;
    const _store = createVirtualStore(
      count,
      itemSize,
      overscan,
      ssrCount,
      cache,
      !itemSize
    );
    return {
      store: _store,
      resizer: createResizer(_store, _isHorizontal),
      scroller: createScroller(_store, _isHorizontal),
      isHorizontal: _isHorizontal,
    };
  });

  const { store, resizer, scroller, isHorizontal } = virtualStore;

  // State for re-rendering
  const stateVersion = useSignal(store.$getStateVersion());

  // Update store when count changes
  if (count !== store.$getItemsLength()) {
    store.$update(ACTION_ITEMS_LENGTH_CHANGE, [count, shift]);
  }
  if (startMargin !== store.$getStartSpacerSize()) {
    store.$update(ACTION_START_OFFSET_CHANGE, startMargin);
  }

  const [startIndex, endIndex] = store.$getRange();
  const isScrolling = store.$isScrolling();
  const totalSize = store.$getTotalSize();

  // Setup effects
  useVisibleTask$(({ cleanup }) => {
    isSSR.value = false;

    // Subscribe to store updates
    const unsubscribeStore = store.$subscribe(UPDATE_VIRTUAL_STATE, (sync: boolean) => {
      stateVersion.value = store.$getStateVersion();
    });

    const unsubscribeOnScroll = store.$subscribe(UPDATE_SCROLL_EVENT, () => {
      if (onScrollProp) {
        onScrollProp(store.$getScrollOffset());
      }
    });

    const unsubscribeOnScrollEnd = store.$subscribe(
      UPDATE_SCROLL_END_EVENT,
      () => {
        if (onScrollEndProp) {
          onScrollEndProp();
        }
      }
    );

    const assignScrollableElement = (e: HTMLElement) => {
      resizer.$observeRoot(e);
      scroller.$observe(e);
    };

    if (scrollRef?.value) {
      microtask(() => assignScrollableElement(scrollRef.value as HTMLElement));
    } else if (containerRef.value) {
      assignScrollableElement(containerRef.value.parentElement!);
    }

    cleanup(() => {
      unsubscribeStore();
      unsubscribeOnScroll();
      unsubscribeOnScrollEnd();
      resizer.$dispose();
      scroller.$dispose();
    });
  });

  // Setup scroll jump fix effect
  useVisibleTask$(({ track }) => {
    track(() => stateVersion.value);
    scroller.$fixScrollJump();
  });

  // Setup imperative handle
  useVisibleTask$(() => {
    if (ref) {
      const handle: VirtualizerHandle = {
        get cache() {
          return store.$getCacheSnapshot();
        },
        get scrollOffset() {
          return store.$getScrollOffset();
        },
        get scrollSize() {
          return getScrollSize(store);
        },
        get viewportSize() {
          return store.$getViewportSize();
        },
        findStartIndex: store.$findStartIndex,
        findEndIndex: store.$findEndIndex,
        getItemOffset: store.$getItemOffset,
        getItemSize: store.$getItemSize,
        scrollToIndex: scroller.$scrollToIndex,
        scrollTo: scroller.$scrollTo,
        scrollBy: scroller.$scrollBy,
      };
      ref.value = handle;
    }
  });

  const getListItem = (index: number): JSXNode => {
    const e = getElement(index);

    return (
      <ListItem
        key={getKey(e, index)}
        children={e}
        resizer={resizer.$observeItem}
        index={index}
        offset={store.$getItemOffset(index)}
        hide={store.$isUnmeasuredItem(index)}
        as={ItemElement as "div"}
        isHorizontal={isHorizontal}
        isSSR={isSSR.value}
      />
    );
  };

  // Build items list
  const items: JSXNode[] = [];
  for (let i = startIndex, j = endIndex; i <= j; i++) {
    items.push(getListItem(i));
  }

  if (keepMounted) {
    const startItems: JSXNode[] = [];
    const endItems: JSXNode[] = [];
    sort(keepMounted).forEach((index) => {
      if (index < startIndex) {
        startItems.push(getListItem(index));
      }
      if (index > endIndex) {
        endItems.push(getListItem(index));
      }
    });

    items.unshift(...startItems);
    items.push(...endItems);
  }

  // Container styles
  const containerStyle = {
    // contain: "content",
    "overflow-anchor": "none", // opt out browser's scroll anchoring
    flex: "none", // flex style can break layout
    position: "relative",
    visibility: "hidden", // TODO replace with other optimization methods
    width: isHorizontal ? totalSize + "px" : "100%",
    height: isHorizontal ? "100%" : totalSize + "px",
    "pointer-events": isScrolling ? "none" : undefined,
  };

  // Handle different container element types
  if (typeof Element === "string") {
    switch (Element) {
      case "div":
        return (
          <div ref={containerRef} style={containerStyle}>
            {items}
          </div>
        );
      case "ul":
        return (
          <ul ref={containerRef} style={containerStyle}>
            {items}
          </ul>
        );
      case "ol":
        return (
          <ol ref={containerRef} style={containerStyle}>
            {items}
          </ol>
        );
      default:
        return (
          <div ref={containerRef} style={containerStyle}>
            {items}
          </div>
        );
    }
  } else {
    // Custom component
    const Component = Element as CustomContainerComponent;
    return (
      <Component ref={containerRef} style={containerStyle}>
        {items}
      </Component>
    );
  }
});