import {
  type FC,
  type RefObject,
  useCallback,
  useRef,
  useEffect,
  useLayoutEffect,
  memo,
  type PropsWithChildren,
  type ComponentType,
  useMemo,
  Fragment,
} from "react";
import { type UseBoundStore, type StoreApi, create } from "zustand";
import { useAui } from "@assistant-ui/store";
import { ThreadListItemRuntimeProvider } from "../providers/ThreadListItemRuntimeProvider";
import type { ThreadRuntimeCore } from "../../runtime/interfaces/thread-runtime-core";
import type { ThreadListRuntimeCore } from "../../runtime/interfaces/thread-list-runtime-core";
import type { AssistantRuntime } from "../../runtime/api/assistant-runtime";
import { BaseSubscribable } from "../../subscribable/subscribable";
import type { ThreadRuntimeImpl } from "../../runtime/api/thread-runtime";
import { ThreadListRuntimeImpl } from "../../runtime/api/thread-list-runtime";

type RemoteThreadListHook = () => AssistantRuntime;

type RemoteThreadListHookInstance = {
  runtime?: ThreadRuntimeCore;
};

const ProviderRenderDetector: FC<{
  detectorRef: RefObject<boolean>;
}> = ({ detectorRef }) => {
  useLayoutEffect(() => {
    detectorRef.current = true;
  }, [detectorRef]);
  return null;
};
export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
  private useRuntimeHook: UseBoundStore<
    StoreApi<{ useRuntime: RemoteThreadListHook }>
  >;
  private instances = new Map<string, RemoteThreadListHookInstance>();
  private useAliveThreadsKeysChanged = create(() => ({}));
  private parent: ThreadListRuntimeCore;

  constructor(
    runtimeHook: RemoteThreadListHook,
    parent: ThreadListRuntimeCore,
  ) {
    super();
    this.parent = parent;
    this.useRuntimeHook = create(() => ({ useRuntime: runtimeHook }));
  }

  public startThreadRuntime(threadId: string) {
    if (!this.instances.has(threadId)) {
      this.instances.set(threadId, {});
      this.useAliveThreadsKeysChanged.setState({}, true);
    }

    return new Promise<ThreadRuntimeCore>((resolve, reject) => {
      const callback = () => {
        const instance = this.instances.get(threadId);
        if (!instance) {
          dispose();
          reject(new Error("Thread was deleted before runtime was started"));
        } else if (!instance.runtime) {
          return; // misc update
        } else {
          dispose();
          resolve(instance.runtime);
        }
      };
      const dispose = this.subscribe(callback);
      callback();
    });
  }

  public getThreadRuntimeCore(threadId: string) {
    const instance = this.instances.get(threadId);
    if (!instance) return undefined;
    return instance.runtime;
  }

  public stopThreadRuntime(threadId: string) {
    this.instances.delete(threadId);
    this.useAliveThreadsKeysChanged.setState({}, true);
  }

  public setRuntimeHook(newRuntimeHook: RemoteThreadListHook) {
    const prevRuntimeHook = this.useRuntimeHook.getState().useRuntime;
    if (prevRuntimeHook !== newRuntimeHook) {
      this.useRuntimeHook.setState({ useRuntime: newRuntimeHook }, true);
    }
  }

  // Rendered outside the user's Provider so deferred `children` cannot strand the binding.
  private _RuntimeBinder: FC<PropsWithChildren<{ threadId: string }>> = ({
    threadId,
    children,
  }) => {
    const { useRuntime } = this.useRuntimeHook();
    const runtime = useRuntime();

    const threadBinding = (runtime.thread as ThreadRuntimeImpl)
      .__internal_threadBinding;

    const updateRuntime = useCallback(() => {
      const aliveThread = this.instances.get(threadId);
      if (!aliveThread)
        throw new Error("Thread not found. This is a bug in assistant-ui.");

      aliveThread.runtime = threadBinding.getState();
      this._notifySubscribers();
    }, [threadId, threadBinding]);

    const isMounted = useRef(false);
    if (!isMounted.current) {
      updateRuntime();
    }

    useEffect(() => {
      isMounted.current = true;
      updateRuntime();
      return threadBinding.outerSubscribe(updateRuntime);
    }, [threadBinding, updateRuntime]);

    const aui = useAui();
    const initPromiseRef = useRef<Promise<unknown> | undefined>(undefined);

    useEffect(() => {
      const runtimeCore = threadBinding.getState();
      const setGetInitializePromise = (runtimeCore as Record<string, unknown>)
        .__internal_setGetInitializePromise;
      if (typeof setGetInitializePromise === "function") {
        setGetInitializePromise.call(runtimeCore, () => initPromiseRef.current);
      }
    }, [threadBinding]);

    useEffect(() => {
      return runtime.threads.main.unstable_on("initialize", () => {
        const state = aui.threadListItem().getState();
        if (state.status === "new") {
          initPromiseRef.current = aui.threadListItem().initialize();

          const dispose = runtime.thread.unstable_on("runEnd", () => {
            dispose();

            aui.threadListItem().generateTitle();
          });
        }
      });
    }, [runtime, aui]);

    return <>{children}</>;
  };

  private _OuterActiveThreadProvider: FC<{
    threadId: string;
    provider: ComponentType<PropsWithChildren>;
  }> = memo(({ threadId, provider: Provider }) => {
    const runtime = useMemo(
      () => new ThreadListRuntimeImpl(this.parent).getItemById(threadId),
      [threadId],
    );

    const detectorRef = useRef(false);
    useEffect(() => {
      if (process.env.NODE_ENV !== "production" && Provider !== Fragment) {
        const id = setTimeout(() => {
          if (!detectorRef.current) {
            console.warn(
              "RemoteThreadListAdapter.unstable_Provider did not render its `children` synchronously. " +
                "Render `children` on first commit; deferring them behind a loading state, Suspense boundary, " +
                "or `useEffect` gate leaves thread context unavailable to downstream consumers.",
            );
          }
        }, 100);
        return () => clearTimeout(id);
      }
      return undefined;
    }, [Provider]);

    return (
      <ThreadListItemRuntimeProvider runtime={runtime}>
        <this._RuntimeBinder threadId={threadId}>
          <Provider>
            <ProviderRenderDetector detectorRef={detectorRef} />
          </Provider>
        </this._RuntimeBinder>
      </ThreadListItemRuntimeProvider>
    );
  });

  public __internal_RenderThreadRuntimes: FC<{
    provider: ComponentType<PropsWithChildren>;
  }> = ({ provider }) => {
    this.useAliveThreadsKeysChanged(); // trigger re-render on alive threads change

    return Array.from(this.instances.keys()).map((threadId) => (
      <this._OuterActiveThreadProvider
        key={threadId}
        threadId={threadId}
        provider={provider}
      />
    ));
  };
}