import { useAuth } from '@clerk/nextjs';
import { autorun } from 'mobx';
import {
  createContext,
  useCallback,
  useContext,
  useEffect,
  useMemo,
  useRef,
} from 'react';
import storageAvailable from 'storage-available';

import { useStores } from '@/app/(root)/AppProviders';
import { useContextSelector } from '@/hooks/useContextSelector';
import logWebUserEvent from '@/logging/logWebUserEvent';
import { getClipTitle } from '@/utils/clip';

import CreateFormContext from '../create/v2/CreateFormContext';
import resetCreateInputs from '../create/v2/actions/resetCreateInputs';
import { useOrpheusReferenceTracking } from './hooks/useOrpheusReferenceTracking';
import { useChatMessagesStore, useChatStore } from './stores';
import { ChatSession } from './types';
import {
  WORKSPACES_QUERY_KEY,
  workspaceListQueryClient,
} from './useWorkspaces';
import { MODAL_SRV, getModalBaseUrl, getUUID } from './utils';

// Enable MapSet support for immer

const getNewSessionUUID = () => {
  const randomUUID = getUUID();
  if (storageAvailable('localStorage')) {
    localStorage.setItem('chat-uuid', randomUUID);
    localStorage.setItem('chat-uuid-saved-at', Date.now().toString());
  }
  return randomUUID;
};

export const useChat = (isEnabled: boolean = true) => {
  const {
    variant,
    setVariant,

    // utils (combine with messages)
    isChatOwner,
    setIsChatOwner,
    toolArgsCache,
    addToolArgsToCache,
    setAbly,

    // references
    references,
    addReference,
    removeReference,
    clearReferences,
    clearReferenceType,
    referencedClipExtendFromSeconds,
    setReferencedClipExtendFromSeconds,
    isSelectingLyrics,
    setIsSelectingLyrics,

    // sessions
    chatUUID,
    setChatUUID,
    addSessionId,
    prependSessionId,
    setSessionMapValue,
    isSessionsDataLoaded,
    setIsSessionsDataLoaded,
    isLoadingChatHistory,
    setIsLoadingChatHistory,
    isCreatingNewSession,
    setIsCreatingNewSession,
    createdSessionWorkspaceIds,
    addCreatedSessionWorkspaceId,
    workspaceId,
    setWorkspaceId,
    isMaximumSessionLength,
    setIsMaximumSessionLength,

    // UI
    isSidebarDismissed,
    setIsSidebarDismissed,
    isChatDismissed,
    setIsChatDismissed,
    isChatHistoryOpen,
    setIsChatHistoryOpen,
    isShareModalOpen,
    setIsShareModalOpen,
    isEditingModalOpen,
    setIsEditingModalOpen,
    isLibraryModalOpen,
    setIsLibraryModalOpen,
    lastPlayedClipId,
    setLastPlayedClipId,
    isSessionListShown,
    setIsSessionListShown,
    sessionIds,
    sessionMap,
    editingContext,
    setEditingContext,
    triggerWorkspaceRefetch,
    workspaceRefetchedAt,
    isOutOfCreditsDismissed,
    setIsOutOfCreditsDismissed,
  } = useChatStore();
  const resetMsgData = useChatMessagesStore((state) => state.resetMsgData);
  const setLoadedMessages = useChatMessagesStore(
    (state) => state.setLoadedMessages
  );

  const setState = useContextSelector(
    CreateFormContext,
    (context) => context.setState
  );

  const rootStore = useStores();
  const lastRunningRequests = useRef<Set<string>>(new Set());
  const { getToken } = useAuth();
  const {
    addReference: trackedAddReference,
    removeReference: trackedRemoveReference,
    clearReferences: trackedClearReferences,
    clearReferenceType: trackedClearReferenceType,
  } = useOrpheusReferenceTracking();

  const sessions: ChatSession[] = useMemo(
    () =>
      [
        ...sessionIds
          .map((sessionId: string) => sessionMap.get(sessionId))
          .filter((session: ChatSession | undefined) => !!session),
      ] as ChatSession[],
    [sessionIds, sessionMap]
  );

  const currentSession = useMemo(() => {
    return sessions.find(
      (session: ChatSession) => session.session_id === chatUUID
    );
  }, [sessions, chatUUID]);

  useEffect(() => {
    setWorkspaceId(currentSession?.workspace_id ?? null);
  }, [currentSession?.workspace_id]);

  useEffect(() => {
    return () => {
      if (fetchChatHistoryAbortControllerRef.current) {
        fetchChatHistoryAbortControllerRef.current.abort();
        fetchChatHistoryAbortControllerRef.current = null;
      }
    };
  }, []);

  const attemptRegisterSession = useCallback(
    async (sessionId: string, workspaceId?: string) => {
      const workspaceIdToRegister =
        workspaceId ??
        (await rootStore.project.createProject('Untitled', '', true))?.id;

      if (!workspaceIdToRegister) {
        console.error('Failed to register workspace for session');
        return {};
      }
      if (!workspaceId) {
        addCreatedSessionWorkspaceId(workspaceIdToRegister);
      }
      const modalSrv = getModalBaseUrl();
      const token = await getToken();
      const resp = await fetch(`${modalSrv}/register-session`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          Authorization: `Bearer ${token}`,
        },
        body: JSON.stringify({
          session_id: sessionId,
          workspace_id: workspaceIdToRegister,
          creation_source: !workspaceId ? 'orpheus' : null,
        }),
      });
      if (resp.ok) {
        if (rootStore.project.currentProjectId !== workspaceIdToRegister) {
          rootStore.project.checkAndUpdateProjectState(workspaceIdToRegister);
        }
        const existingSession = sessionMap.get(sessionId) ?? {};
        setSessionMapValue(sessionId, {
          ...existingSession,
          workspace_id: workspaceIdToRegister,
        });
        if (storageAvailable('localStorage') && rootStore.session?.userId) {
          localStorage.setItem(
            `orpheus_session_last_touched:${rootStore.session?.userId}`,
            Date.now().toString()
          );
          localStorage.setItem(
            `orpheus_last_workspace_id:${rootStore.session?.userId}`,
            workspaceIdToRegister
          );
        }

        // Log workspace creation if it was created from Orpheus
        if (!workspaceId) {
          logWebUserEvent(
            {
              actionName: 'OrpheusWorkspaceCreated',
              context: {
                workspaceId: workspaceIdToRegister,
                workspaceName: 'Untitled',
                source: 'chat',
              },
            },
            rootStore.session
          );
        }
      } else {
        console.error('Failed to register session');
        return {};
      }
      await workspaceListQueryClient.invalidateQueries({
        queryKey: [WORKSPACES_QUERY_KEY],
        refetchType: 'all',
      });
      return {
        sessionId,
        workspaceId: workspaceIdToRegister,
      };
    },
    [
      rootStore.project,
      currentSession,
      workspaceId,
      setSessionMapValue,
      sessionMap,
      rootStore.session,
    ]
  );

  const registerClipStatusChange = useCallback(
    async ({
      clipId,
      status,
      duration,
      title,
      task,
      lyrics,
      tags,
    }: {
      clipId: string;
      status: string;
      duration: number;
      title: string;
      task?: string | null;
      lyrics?: string;
      tags?: string;
    }) => {
      const token = await getToken();
      if (!chatUUID) return;
      await fetch(`${MODAL_SRV}/context-state`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          Authorization: `Bearer ${token}`,
        },
        body: JSON.stringify({
          session_id: chatUUID,
          type: 'clip_status',
          metadata: {
            clip_id: clipId,
            status,
            duration,
            task,
            title,
            lyrics,
            tags,
          },
        }),
      });
    },
    [getToken, chatUUID]
  );

  const registerModelChange = useCallback(
    async (model: string) => {
      const token = await getToken();
      if (!chatUUID) return;
      await fetch(`${MODAL_SRV}/context-state`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          Authorization: `Bearer ${token}`,
        },
        body: JSON.stringify({
          session_id: chatUUID,
          type: 'model_status',
          metadata: {
            model_name: model,
          },
        }),
      });
    },
    [getToken, chatUUID]
  );

  const createChatSession = useCallback(
    async (workspaceId?: string) => {
      useChatMessagesStore.getState().resetMsgData();
      useChatMessagesStore.getState().setCurrentAssistantPrompt(null);
      useChatMessagesStore.getState().setIsLoading(false);
      const newChatUUID = getNewSessionUUID();
      setChatUUID(newChatUUID);
      setLoadedMessages([]);
      addSessionId(newChatUUID);
      setSessionMapValue(newChatUUID, { session_id: newChatUUID });
      setLastPlayedClipId(null);
      setIsChatHistoryOpen(false);
      setIsSidebarDismissed(true);
      setState(resetCreateInputs);
      setIsMaximumSessionLength(false);
      setIsLoadingChatHistory(false);
      clearReferences();
      if (workspaceId) {
        await attemptRegisterSession(newChatUUID, workspaceId);
        setWorkspaceId(workspaceId);
      } else {
        setIsCreatingNewSession(true);
      }
      logWebUserEvent(
        {
          actionName: 'OrpheusChatPageImpression',
          context: {
            sessionId: newChatUUID,
            workspaceId: workspaceId,
            isNewSession: true,
            hasExistingMessages: false,
          },
        },
        rootStore.session
      );
      return newChatUUID;
    },
    [setChatUUID, getUUID, workspaceId, rootStore.session]
  );

  const fetchChatHistoryAbortControllerRef = useRef<AbortController | null>(
    null
  );

  const fetchChatHistory = useCallback(
    async (sessionId: string, workspaceId: string) => {
      // Cancel any in-flight request
      if (fetchChatHistoryAbortControllerRef.current) {
        fetchChatHistoryAbortControllerRef.current.abort();
      }

      const abortController = new AbortController();
      fetchChatHistoryAbortControllerRef.current = abortController;

      const token = await getToken();
      const modalBaseUrl = getModalBaseUrl();

      setIsLoadingChatHistory(true);

      try {
        const response = await fetch(
          `${modalBaseUrl}/chat-history/${sessionId}`,
          {
            method: 'GET',
            headers: {
              'Content-Type': 'application/json',
              Authorization: `Bearer ${token}`,
            },
            signal: abortController.signal, // Add abort signal
          }
        );
        if (response.ok) {
          const chatHistory = await response.json();
          setLoadedMessages(
            chatHistory?.messages?.map((msg: any) => ({
              ...msg,
              role: msg.role,
              message: msg.content,
            })) ?? null
          );
          const lastAssistantPromptMsg = chatHistory?.messages?.findLast(
            (msg: any) => msg.role === 'assistant' && !!msg.current_prompt
          );
          if (lastAssistantPromptMsg) {
            useChatMessagesStore
              .getState()
              .setCurrentAssistantPrompt(lastAssistantPromptMsg.current_prompt);
          }
          resetMsgData();
          if (chatHistory?.messages?.length) {
            setIsCreatingNewSession(false);
          }
          setIsMaximumSessionLength(false);

          if (storageAvailable('localStorage') && rootStore.session?.userId) {
            localStorage.setItem(
              `orpheus_session_last_touched:${rootStore.session?.userId}`,
              Date.now().toString()
            );
            localStorage.setItem(
              `orpheus_last_workspace_id:${rootStore.session?.userId}`,
              workspaceId
            );
          }
          // Log chat page impression
          logWebUserEvent(
            {
              actionName: 'OrpheusChatPageImpression',
              context: {
                sessionId: sessionId,
                workspaceId: workspaceId,
                isNewSession:
                  !chatHistory?.messages || chatHistory.messages.length === 0,
                hasExistingMessages: !!(
                  chatHistory?.messages && chatHistory.messages.length > 0
                ),
              },
            },
            rootStore.session
          );

          chatHistory?.messages
            ?.filter((historyItem: any) => !!historyItem.tool_calls)
            ?.forEach((historyItem: any) => {
              historyItem.tool_calls.forEach((toolCall: any) => {
                try {
                  const args = JSON.parse(toolCall.function.arguments);
                  addToolArgsToCache(toolCall.id, args);
                } catch (error) {
                  console.log(error);
                }
              });
            });
        }
        // TODO: Add error state here
      } catch (error) {
        // Ignore abort errors - they're expected when cancelling requests
        if (error instanceof Error && error.name === 'AbortError') {
          return;
        }
        // Handle other errors
        console.error('Error fetching chat history:', error);
      } finally {
        // Only set loading to false if this request wasn't aborted
        if (!abortController.signal.aborted) {
          setIsLoadingChatHistory(false);
        }
        // Clear the ref if this was the current request
        if (fetchChatHistoryAbortControllerRef.current === abortController) {
          fetchChatHistoryAbortControllerRef.current = null;
        }
      }
    },
    [
      getToken,
      setLoadedMessages,
      resetMsgData,
      addToolArgsToCache,
      rootStore.session,
    ]
  );

  useEffect(() => {
    if (!isEnabled) return;
    const dispose = autorun(() => {
      const currentRequests = rootStore.clips.runningRequests;
      if (lastRunningRequests.current) {
        const setDiff = lastRunningRequests.current.difference(currentRequests);
        setDiff.forEach((clipId: string) => {
          const clip = rootStore.clips.clipById[clipId];
          const NOTIFICATION_STATES = ['error', 'complete', 'streaming'];
          if (
            clip?.status &&
            NOTIFICATION_STATES.includes(clip?.status ?? '')
          ) {
            const clipTitle =
              getClipTitle(clip) +
              ` (${(clip.batch_index ?? 0) + (clip.metadata?.batch_offset ?? 0) + 1})`;
            registerClipStatusChange({
              clipId: clip.id,
              status: clip.status,
              duration: clip.metadata?.duration ?? 0,
              title: clipTitle,
              task: clip.metadata?.task,
              lyrics: clip.metadata?.prompt ?? undefined,
              tags: clip.metadata?.tags ?? undefined,
            });
          }
        });
      }
      lastRunningRequests.current = new Set([...currentRequests]);
    });
    return dispose;
  }, [isEnabled, registerClipStatusChange]);

  const handleSetChatUUID = useCallback(
    (chatUUID: string) => {
      setChatUUID(chatUUID);
      if (storageAvailable('localStorage')) {
        localStorage.setItem('chat-uuid', chatUUID);
        localStorage.setItem('chat-uuid-saved-at', Date.now().toString());
      }
    },
    [setChatUUID]
  );

  return useMemo(() => {
    return {
      createChatSession,
      chatUUID,
      setChatUUID: handleSetChatUUID,
      sessions,
      currentSession,
      references,
      addReference: trackedAddReference,
      removeReference: trackedRemoveReference,
      clearReferences: trackedClearReferences,
      clearReferenceType: trackedClearReferenceType,
      isSelectingLyrics,
      setIsSelectingLyrics,
      referencedClipExtendFromSeconds,
      setReferencedClipExtendFromSeconds,
      variant,
      setVariant,
      toolArgsCache,
      setAbly,
      isShareModalOpen,
      setIsShareModalOpen,
      isEditingModalOpen,
      setIsEditingModalOpen,
      isLibraryModalOpen,
      setIsLibraryModalOpen,
      addToolArgsToCache,
      isSidebarDismissed,
      setIsSidebarDismissed,
      isChatDismissed,
      setIsChatDismissed,
      isChatHistoryOpen,
      setIsChatHistoryOpen,
      addSessionId,
      prependSessionId,
      setSessionMapValue,
      registerClipStatusChange,
      registerModelChange,
      isChatOwner,
      setIsChatOwner,
      lastPlayedClipId,
      setLastPlayedClipId,
      isSessionListShown,
      setIsSessionListShown,
      isSessionsDataLoaded,
      setIsSessionsDataLoaded,
      workspaceId,
      setWorkspaceId,
      attemptRegisterSession,
      isCreatingNewSession,
      setIsCreatingNewSession,
      createdSessionWorkspaceIds,
      addCreatedSessionWorkspaceId,
      editingContext,
      setEditingContext,
      isMaximumSessionLength,
      setIsMaximumSessionLength,
      triggerWorkspaceRefetch,
      workspaceRefetchedAt,
      fetchChatHistory,
      isLoadingChatHistory,
      setIsLoadingChatHistory,
      isOutOfCreditsDismissed,
      setIsOutOfCreditsDismissed,
    };
  }, [
    createChatSession,
    chatUUID,
    handleSetChatUUID,
    sessions,
    currentSession,
    references,
    addReference,
    removeReference,
    clearReferenceType,
    clearReferences,
    isSelectingLyrics,
    setIsSelectingLyrics,
    referencedClipExtendFromSeconds,
    setReferencedClipExtendFromSeconds,
    variant,
    setVariant,
    toolArgsCache,
    setAbly,
    isShareModalOpen,
    setIsShareModalOpen,
    isEditingModalOpen,
    setIsEditingModalOpen,
    isLibraryModalOpen,
    setIsLibraryModalOpen,
    addToolArgsToCache,
    isSidebarDismissed,
    setIsSidebarDismissed,
    isChatDismissed,
    setIsChatDismissed,
    isChatHistoryOpen,
    setIsChatHistoryOpen,
    addSessionId,
    prependSessionId,
    setSessionMapValue,
    registerClipStatusChange,
    registerModelChange,
    isChatOwner,
    setIsChatOwner,
    lastPlayedClipId,
    setLastPlayedClipId,
    isSessionListShown,
    setIsSessionListShown,
    isSessionsDataLoaded,
    setIsSessionsDataLoaded,
    workspaceId,
    setWorkspaceId,
    attemptRegisterSession,
    isCreatingNewSession,
    setIsCreatingNewSession,
    createdSessionWorkspaceIds,
    addCreatedSessionWorkspaceId,
    editingContext,
    setEditingContext,
    isMaximumSessionLength,
    setIsMaximumSessionLength,
    triggerWorkspaceRefetch,
    workspaceRefetchedAt,
    fetchChatHistory,
    isLoadingChatHistory,
    setIsLoadingChatHistory,
    isOutOfCreditsDismissed,
    setIsOutOfCreditsDismissed,
  ]);
};

const OrpheusChatContext = createContext<ReturnType<typeof useChat> | null>(
  null
);
export const OrpheusChatProvider = OrpheusChatContext.Provider;
export const useChatContext = () => {
  const context = useContext(OrpheusChatContext);
  if (!context) {
    throw new Error('useChatContext must be used within a OrpheusChatProvider');
  }
  return context;
};