import { useAuth } from '@clerk/nextjs';
import clsx from 'clsx';
import { observer } from 'mobx-react-lite';
import { createContext, useCallback, useMemo } from 'react';

import Button, {
  ButtonShape,
  ButtonSize,
  ButtonVariant,
} from '@/components/button/Button';
import { MultiSelectContextProvider } from '@/components/clipBrowser/MultiSelectContext';
import ImageWithFallback from '@/components/image/ImageWithFallback';
import { SkeletonBone } from '@/components/layout/Skeleton';
import SpinnerSVG from '@/components/svg/SpinnerSVG';
import { CheckIcon } from '@/icons';
import logWebUserEvent from '@/logging/logWebUserEvent';

import { useStores } from '../AppProviders';
import { AssistantMessage } from './components/messages/AssistantMessage';
import { ClipsMessage } from './components/messages/ClipsMessage';
import { DownloadLinkMessage } from './components/messages/DownloadLinkMessage';
import { ErrorMessage } from './components/messages/ErrorMessage';
import { IncomingMessageLoading } from './components/messages/IncomingMessageLoading';
import { PlaylistMessage } from './components/messages/PlaylistMessage';
import { ShareLinkMessage } from './components/messages/ShareLinkMessage';
import { StylesLyricsCard } from './components/messages/StylesLyricsCard';
import { SuccessMessage } from './components/messages/SuccessMessage';
import { UserMessage } from './components/messages/UserMessage';
import { useChatMessagesStore } from './stores';
import { useChatContext } from './useChat';
import { ABTestVariant, useChatMessageActions } from './useChatMessageActions';
import { useSyncQueueWithMessages } from './useSyncQueueWithMessages';
import {
  consolidateMessages,
  getModalBaseUrl,
  parsePartialMessage,
} from './utils';

export const OrpheusPlaybackContext = createContext<{
  clipIds: string[];
} | null>(null);

export const ABMessageContext = createContext<{ abTestGroupId: string } | null>(
  null
);

export const MessageList = observer(
  ({ onAllowChildScroll }: { onAllowChildScroll: () => void }) => {
    const {
      msgMap,
      msgFeedbackMap,
      msgIds,
      loadedMessages,
      isStreaming,
      isLoading,
    } = useChatMessagesStore();

    const { session } = useStores();
    const { chatUUID } = useChatContext();
    const { handleSend: sendMessage, selectVariantForAB } =
      useChatMessageActions();
    const { getToken } = useAuth();
    const handleSelectVariant = useCallback(
      async (abTestGroupId: string, variant: ABTestVariant) => {
        const modalBaseUrl = getModalBaseUrl();
        const token = await getToken();
        const data = await fetch(`${modalBaseUrl}/select-ab-response`, {
          method: 'POST',
          body: JSON.stringify({
            session_id: chatUUID,
            ab_test_group_id: abTestGroupId,
            preferred_ab_variant: variant.toUpperCase(),
          }),
          headers: {
            'Content-Type': 'application/json',
            Authorization: `Bearer ${token}`,
          },
        });
        if (data.ok) {
          logWebUserEvent(
            {
              actionName: 'OrpheusABMessageSelected',
              context: {
                sessionId: chatUUID,
                abTestGroupId,
                abVariant: variant,
              },
            },
            session
          );
          selectVariantForAB(abTestGroupId, variant);
        }
      },
      [chatUUID, selectVariantForAB, getToken, session]
    );

    const messagesWithIds = useMemo(() => {
      return consolidateMessages([
        ...(loadedMessages || []).map((msg) => ({
          id: msg.message_id,
          data: msg,
        })),
        ...msgIds.map((msgId) => ({ id: msgId, data: msgMap.get(msgId) })),
      ]);
    }, [msgIds, msgMap, loadedMessages]);
    const lastMessage = parsePartialMessage(
      messagesWithIds[messagesWithIds.length - 1]?.data
    );

    const isLastMessageAB = useMemo(() => {
      return Boolean(lastMessage?.a) || Boolean(lastMessage?.b);
    }, [lastMessage]);

    // raw messages used only for syncing playback queue
    const rawMessages = useMemo(() => {
      return messagesWithIds
        .map((msg) => {
          const rawMsg = msg.data;
          if (rawMsg.a && rawMsg.b) {
            return [...rawMsg.a, ...rawMsg.b].map((msgId: string) =>
              msgMap.get(msgId)
            );
          } else {
            return [rawMsg];
          }
        })
        .flat();
    }, [messagesWithIds, msgMap]);

    const { messageClipIds, clipsIndexes } =
      useSyncQueueWithMessages(rawMessages);

    const renderMessage = useCallback(
      (
        index: number,
        messageId: any,
        parsedMessage: any,
        previousParsedMessage: any
      ) => {
        if (
          (parsedMessage.type === 'tool_call' &&
            parsedMessage?.name !== 'simple_message') ||
          parsedMessage.isHiddenMessage
        ) {
          return null;
        }
        if (
          !!parsedMessage.parsedMessage?.success_message ||
          !!parsedMessage.parsed_content?.success_message
        ) {
          return (
            <SuccessMessage
              key={index}
              message={
                parsedMessage.parsedMessage?.success_message ??
                parsedMessage.parsed_content?.success_message
              }
            />
          );
        }
        if (
          !!parsedMessage.parsedMessage?.error_message ||
          !!parsedMessage.parsed_content?.error_message
        ) {
          return (
            <ErrorMessage
              key={index}
              message={
                parsedMessage.parsedMessage?.error_message ??
                parsedMessage.parsed_content?.error_message
              }
            />
          );
        }
        if (
          !!parsedMessage.parsedMessage?.share_link ||
          !!parsedMessage.parsed_content?.share_link
        ) {
          return (
            <ShareLinkMessage
              key={index}
              clipId={
                parsedMessage.parsedMessage?.clip_id ??
                parsedMessage.parsed_content?.clip_id
              }
              link={
                parsedMessage.parsedMessage?.share_link ??
                parsedMessage.parsed_content?.share_link
              }
            />
          );
        }
        if (
          !!parsedMessage.parsedMessage?.download_link ||
          !!parsedMessage.parsed_content?.download_link
        ) {
          return (
            <DownloadLinkMessage
              key={index}
              clipId={
                parsedMessage.parsedMessage?.clip_id ??
                parsedMessage.parsed_content?.clip_id
              }
              link={
                parsedMessage.parsedMessage?.download_link ??
                parsedMessage.parsed_content?.download_link
              }
            />
          );
        }
        if (!!parsedMessage?.parsed_content?.result) {
          // this is a search result loaded from chat history
          const clipIds = (
            parsedMessage.parsed_content.result?.['']?.result ?? []
          ).map((result: any) => result.id);
          if (clipIds.length > 0) {
            return (
              <ClipsMessage
                key={index}
                clipIds={clipIds}
                toolCallId={parsedMessage.tool_call_id || ''}
                role={parsedMessage.role}
                showCreateMore={false}
                showStylesAndLyrics={false}
                startingIndex={clipsIndexes[index] || 0}
                onAllowChildScroll={onAllowChildScroll}
              />
            );
          } else {
            return null;
          }
        }
        if (!!parsedMessage.parsed_content?.playlist?.id) {
          return (
            <PlaylistMessage
              key={index}
              playlistIds={[parsedMessage.parsed_content?.playlist?.id]}
            />
          );
        }
        if (
          !!parsedMessage?.clipIds ||
          !!parsedMessage?.parsed_content?.clipIds
        ) {
          return (
            <ClipsMessage
              key={index}
              clipIds={
                parsedMessage.clipIds ?? parsedMessage?.parsed_content?.clipIds
              }
              stylesVersion={
                parsedMessage?.parsedMessage?.styles_version ??
                parsedMessage?.parsed_content?.styles_version
              }
              lyricsVersion={
                parsedMessage?.parsedMessage?.lyrics_version ??
                parsedMessage?.parsed_content?.lyrics_version
              }
              toolCallId={parsedMessage.tool_call_id || ''}
              role={parsedMessage.role}
              showCreateMore={
                !(parsedMessage?.tool_name ?? '').startsWith('search_')
              }
              showStylesAndLyrics={
                // !(parsedMessage?.tool_name ?? '').startsWith('search_')
                false
              }
              startingIndex={clipsIndexes[index] || 0}
              onAllowChildScroll={onAllowChildScroll}
            />
          );
        }
        if (
          !!parsedMessage?.playlistIds ||
          !!parsedMessage?.parsed_content?.playlistIds
        ) {
          return (
            <PlaylistMessage
              key={index}
              playlistIds={
                parsedMessage.playlistIds ||
                parsedMessage?.parsed_content?.playlistIds
              }
            />
          );
        }

        if ((parsedMessage as any)?.parsedMessage?.type === 'image') {
          return (
            <div
              key={index}
              className={clsx('flex w-full py-4', {
                'flex-row':
                  (parsedMessage as any)?.role !== 'user' ||
                  (parsedMessage as any)?.isCollabUserMessage,
                'flex-row-reverse':
                  (parsedMessage as any)?.role === 'user' &&
                  !(parsedMessage as any)?.isCollabUserMessage,
              })}
            >
              <div className='flex flex-col gap-4'>
                <div className='relative aspect-square min-h-[min(70vw,400px)] w-full max-w-[400px] overflow-hidden rounded-lg'>
                  <SkeletonBone className='absolute inset-0 rounded-lg' />
                  <ImageWithFallback
                    src={(parsedMessage as any)?.parsedMessage?.url}
                    alt='Chat image'
                    className='relative z-10 h-auto w-full object-cover'
                  />
                </div>
              </div>
            </div>
          );
        }

        const isLyricsToolMessage =
          ((parsedMessage as any)?.role === 'tool' ||
            ((parsedMessage as any)?.role === 'assistant' &&
              (parsedMessage as any)?.tool_call_id)) &&
          (parsedMessage as any)?.type !== 'simple_message' &&
          !(
            !!(parsedMessage as any)?.parsed_content?.message &&
            !!(parsedMessage as any)?.parsed_content?.choices
          );

        if (isLyricsToolMessage) {
          const lyricsMessage = (parsedMessage as any)?.message;
          const lyricsVersion = (parsedMessage as any)?.lyrics_version;
          const isLastMessage = index === messagesWithIds.length - 1;
          const cardIsStreaming = isStreaming && isLastMessage;
          const cardIsLoading = isLoading && isLastMessage;

          return (
            <div key={messageId || index}>
              <StylesLyricsCard
                toolCallId={(parsedMessage as any)?.tool_call_id}
                lyrics={lyricsMessage}
                lyricsVersion={lyricsVersion}
                isStreaming={cardIsStreaming}
                isLoading={cardIsLoading}
                onAllowChildScroll={onAllowChildScroll}
              />
            </div>
          );
        }

        return ((parsedMessage as any)?.parsedArguments?.message ??
          (parsedMessage as any)?.message !== '') ? (
          <div
            key={index}
            className={clsx('flex w-full', {
              'flex-row':
                (parsedMessage as any)?.role !== 'user' ||
                (parsedMessage as any)?.isCollabUserMessage,
              'flex-row-reverse':
                (parsedMessage as any)?.role === 'user' &&
                !(parsedMessage as any)?.isCollabUserMessage,
            })}
          >
            <div className='flex flex-col gap-4'>
              <div className='w-auto max-w-full rounded-lg whitespace-pre-wrap'>
                {(parsedMessage as any)?.name !== 'simple_message' &&
                (parsedMessage as any)?.type !== 'simple_message' &&
                !(
                  !!(parsedMessage as any)?.parsed_content?.message &&
                  !!(parsedMessage as any)?.parsed_content?.choices
                ) ? (
                  ((parsedMessage as any)?.parsedArguments?.message ??
                  typeof (parsedMessage as any)?.message === 'string') ? (
                    (parsedMessage as any)?.role === 'user' ? (
                      <UserMessage
                        text={(parsedMessage as any)?.message}
                        referenceClipId={
                          (parsedMessage as any)?.referenced_clip_id
                        }
                        referenceLyrics={
                          (parsedMessage as any)?.referenced_lyrics
                        }
                        referenceStyles={
                          (parsedMessage as any)?.referenced_styles
                        }
                      />
                    ) : (
                      <>
                        {/* Pre-loaded assistant message */}
                        <AssistantMessage
                          messageId={messageId}
                          text={
                            (parsedMessage as any)?.parsedArguments?.message ??
                            (parsedMessage as any)?.parsedMessage?.message ??
                            (parsedMessage as any)?.message ??
                            ''
                          }
                          shouldAnimate={!parsedMessage?.is_history}
                          onAllowChildScroll={onAllowChildScroll}
                          lyrics={
                            !!previousParsedMessage?.referenced_lyrics
                              ? ((parsedMessage as any)?.parsedArguments
                                  ?.current_prompt?.lyrics ??
                                (parsedMessage as any)?.parsedMessage
                                  ?.current_prompt?.lyrics ??
                                (parsedMessage as any)?.current_prompt
                                  ?.lyrics ??
                                undefined)
                              : undefined
                          }
                          lyricsVersion={
                            (parsedMessage as any)?.parsedArguments
                              ?.current_prompt?.lyrics_version ??
                            (parsedMessage as any)?.parsedMessage
                              ?.current_prompt?.lyrics_version ??
                            (parsedMessage as any)?.current_prompt
                              ?.lyrics_version ??
                            undefined
                          }
                          styles={
                            !!previousParsedMessage?.referenced_styles
                              ? ((parsedMessage as any)?.parsedArguments
                                  ?.current_prompt?.tags ??
                                (parsedMessage as any)?.parsedMessage
                                  ?.current_prompt?.tags ??
                                (parsedMessage as any)?.current_prompt?.tags ??
                                undefined)
                              : undefined
                          }
                          stylesVersion={
                            (parsedMessage as any)?.parsedArguments
                              ?.current_prompt?.tags_version ??
                            (parsedMessage as any)?.parsedMessage
                              ?.current_prompt?.tags_version ??
                            (parsedMessage as any)?.current_prompt
                              ?.tags_version ??
                            undefined
                          }
                          isUpsell={
                            (parsedMessage as any)?.parsedArguments
                              ?.is_upsell ??
                            (parsedMessage as any)?.parsedMessage?.is_upsell ??
                            (parsedMessage as any)?.is_upsell ??
                            undefined
                          }
                          likeStatus={
                            msgFeedbackMap.has(messageId)
                              ? msgFeedbackMap.get(messageId)
                              : (parsedMessage as any)?.like_status
                          }
                        />
                      </>
                    )
                  ) : (
                    JSON.stringify((parsedMessage as any)?.message)
                  )
                ) : (
                  <AssistantMessage
                    messageId={messageId}
                    text={
                      (parsedMessage as any)?.parsedArguments?.message ??
                      (parsedMessage as any)?.parsedMessage?.message ??
                      (parsedMessage as any)?.message ??
                      ''
                    }
                    shouldAnimate={!parsedMessage?.is_history}
                    onAllowChildScroll={onAllowChildScroll}
                    lyrics={
                      !!previousParsedMessage?.referenced_lyrics
                        ? ((parsedMessage as any)?.parsedArguments
                            ?.current_prompt?.lyrics ??
                          (parsedMessage as any)?.parsedMessage?.current_prompt
                            ?.lyrics ??
                          (parsedMessage as any)?.current_prompt?.lyrics ??
                          undefined)
                        : undefined
                    }
                    styles={
                      !!previousParsedMessage?.referenced_styles
                        ? ((parsedMessage as any)?.parsedArguments
                            ?.current_prompt?.tags ??
                          (parsedMessage as any)?.parsedMessage?.current_prompt
                            ?.tags ??
                          (parsedMessage as any)?.current_prompt?.tags ??
                          undefined)
                        : undefined
                    }
                    lyricsVersion={
                      (parsedMessage as any)?.parsedArguments?.current_prompt
                        ?.lyrics_version ??
                      (parsedMessage as any)?.parsedMessage?.current_prompt
                        ?.lyrics_version ??
                      (parsedMessage as any)?.current_prompt?.lyrics_version ??
                      undefined
                    }
                    stylesVersion={
                      (parsedMessage as any)?.parsedArguments?.current_prompt
                        ?.tags_version ??
                      (parsedMessage as any)?.parsedMessage?.current_prompt
                        ?.tags_version ??
                      (parsedMessage as any)?.current_prompt?.tags_version ??
                      undefined
                    }
                    isUpsell={
                      (parsedMessage as any)?.parsedArguments?.is_upsell ??
                      (parsedMessage as any)?.parsedMessage?.is_upsell ??
                      (parsedMessage as any)?.is_upsell ??
                      undefined
                    }
                    likeStatus={
                      msgFeedbackMap.has(messageId)
                        ? msgFeedbackMap.get(messageId)
                        : (parsedMessage as any)?.like_status
                    }
                  />
                )}
              </div>
              {(parsedMessage as any)?.isSuggestion ? (
                <Button
                  size={ButtonSize.Mini}
                  shape={ButtonShape.Rounded}
                  onClick={() => {
                    sendMessage({
                      message: (parsedMessage as any)?.message,
                    });
                  }}
                  iconStart={CheckIcon}
                  className={clsx(
                    'w-min border border-white/20 text-xs hover:border-white/50 hover:bg-background-fog-thin'
                  )}
                >
                  Approve
                </Button>
              ) : null}
            </div>
          </div>
        ) : null;
      },
      [
        clipsIndexes,
        messagesWithIds,
        msgFeedbackMap,
        isLoading,
        isStreaming,
        sendMessage,
        onAllowChildScroll,
      ]
    );

    const renderABMessageGroup = useCallback(
      (
        abTestGroupId: string,
        messagesOrIds: (string | object)[],
        startingIndex: number,
        previousParsedMessage: any,
        variant: ABTestVariant,
        handleSelectVariant: (
          abTestGroupId: string,
          variant: ABTestVariant
        ) => void
      ) => {
        return (
          <>
            {messagesOrIds.map((msgIdOrMessage: string | object, i: number) => {
              const message =
                typeof msgIdOrMessage === 'string'
                  ? msgMap.get(msgIdOrMessage)
                  : msgIdOrMessage;
              const parsedPartial = parsePartialMessage(message);
              return renderMessage(
                startingIndex + i,
                message?.message_id ?? msgIdOrMessage,
                parsedPartial,
                previousParsedMessage
              );
            })}
            {isLoading ? (
              <div className='flex w-full justify-center py-8'>
                <SpinnerSVG />
              </div>
            ) : null}
            <div className='flex-1' />
            <div className='flex w-full justify-center px-8'>
              <Button
                variant={ButtonVariant.Standard}
                shape={ButtonShape.Pill}
                size={ButtonSize.Small}
                onClick={() => handleSelectVariant(abTestGroupId, variant)}
                className='w-full bg-background-fog-thick'
              >
                Choose
              </Button>
            </div>
          </>
        );
      },
      [renderMessage, msgMap, isLoading]
    );

    return (
      <OrpheusPlaybackContext.Provider value={{ clipIds: messageClipIds }}>
        <MultiSelectContextProvider>
          {messagesWithIds
            .map(
              (
                messageWithId: { id: string; data: string | object },
                i: number
              ) => {
                const message = messageWithId.data;
                if (
                  !message ||
                  (typeof message === 'string' &&
                    message.replace(/\s/g, '').trim().length === 0)
                ) {
                  return null;
                }
                const parsedPartial: any = parsePartialMessage(message);
                const previousParsedPartial: any =
                  i > 0
                    ? parsePartialMessage(messagesWithIds[i - 1]?.data)
                    : null;

                if (parsedPartial.a && parsedPartial.b) {
                  return (
                    <ABMessageContext.Provider
                      value={{ abTestGroupId: messageWithId.id }}
                    >
                      <div className='flex gap-4'>
                        {/* [A] Container */}
                        <div className='flex w-[400px] flex-col rounded-xl bg-background-fog-thin p-4 break-words'>
                          {renderABMessageGroup(
                            messageWithId.id,
                            parsedPartial.a,
                            i,
                            previousParsedPartial,
                            ABTestVariant.A,
                            handleSelectVariant
                          )}
                        </div>
                        {/* [B] Container */}
                        <div className='flex w-[400px] flex-col rounded-xl bg-background-fog-thin p-4 break-words'>
                          {renderABMessageGroup(
                            messageWithId.id,
                            parsedPartial.b,
                            i + parsedPartial.a.length,
                            previousParsedPartial,
                            ABTestVariant.B,
                            handleSelectVariant
                          )}
                        </div>
                      </div>
                    </ABMessageContext.Provider>
                  );
                }

                return renderMessage(
                  i,
                  messageWithId?.id,
                  parsedPartial,
                  previousParsedPartial
                );
              }
            )
            .filter((msg) => !!msg)}
          {isLoading && !isLastMessageAB ? (
            <div className='justify-left flex w-full flex-row items-center py-1'>
              <IncomingMessageLoading />
            </div>
          ) : null}
        </MultiSelectContextProvider>
      </OrpheusPlaybackContext.Provider>
    );
  }
);