import clsx from "clsx";
import * as React from "react";
import { AwaitProps, Await as RRAwait } from "react-router";
import { SuspenseContext } from "../../contexts";
import { BaseCollection, useCollection, useShallowRender } from "../Collection";
import * as styles from "./styles.module.css";

export function SwitchContainer({
  children,
  ...props
}: React.ComponentProps<"div"> & Pick<AwaitProps<any>, "resolve">) {
  const { portal, collection } = useCollection({ children });
  return (
    <>
      {portal}
      <SwitchImpl collection={collection} {...props} />
    </>
  );
}

function SwitchImpl({
  collection,
  ...props
}: React.ComponentProps<"div"> &
  Pick<AwaitProps<any>, "resolve"> & {
    collection: BaseCollection<object>;
  }) {
  const [fallback, children] = React.useMemo(() => {
    let fallback: React.ReactNode = null,
      children: React.ReactNode = null;
    for (const node of collection) {
      switch (node.type) {
        case "fallback":
          fallback = node.props.children;
          break;
        case "content":
          children = node.props.children;
          break;
      }
    }
    return [fallback, children];
  }, [collection]);
  return (
    <SwitchContainerComponent
      {...props}
      fallback={fallback}
      children={children}
    />
  );
}

function SwitchContainerComponent({
  fallback,
  children,
  resolve,
  className,
  preview,
  ...props
}: React.ComponentProps<"div"> &
  Pick<AwaitProps<any>, "resolve"> & {
    fallback: React.ReactNode;
    preview?: true;
  } & Pick<AwaitProps<any>, "children">) {
  const child = React.useMemo(() => {
    if (resolve instanceof Promise) {
      return <RRAwait resolve={resolve} children={children} />;
    }
    return preview
      ? fallback
      : resolve
        ? typeof children === "function"
          ? children(resolve)
          : children
        : fallback;
  }, [resolve, children, fallback, preview]);
  return (
    <React.Suspense fallback={fallback}>
      <SuspenseContext.Provider value={true}>
        <div className={clsx(className, styles.container)} {...props}>
          {child}
        </div>
      </SuspenseContext.Provider>
    </React.Suspense>
  );
}

SwitchContainer.Fallback = ({ ref, ...props }: React.ComponentProps<"div">) => {
  return useShallowRender("fallback", props, ref as any);
};

SwitchContainer.Content = ({
  ref,
  ...props
}: React.ComponentProps<"div"> & Pick<AwaitProps<any>, "children">) => {
  return useShallowRender("content", props, ref as any);
};