import {create} from 'zustand';
import {useShallow} from 'zustand/react/shallow';
import {Asset} from '@modules/chats/interfaces';

export enum FileDownloadStatus {
  NOT_INITIALIZED = 'not_initialized',
  BACKEND_NOT_READY = 'backend_not_ready',
  IDLE = 'idle',
  PROGRESS = 'progress',
  DONE = 'done',
  ERROR = 'error',
  NOT_AUTOMATIC_DOWNLOAD = 'not_automatic_download',
}

export interface StoredFile {
  assetId: number;
  uri: Nullable<string>;
  remoteUrl: string;
  filename: string;
  state: FileDownloadStatus;
  progress: number;
  asset: Asset;
}

export interface DownloadAsset {
  assetId: number;
  uri: Nullable<string>;
  messageId: string;
  remoteUrl: string;
  filename: string;
  asset: Asset;
}

type DownloadsByMessage = Map<string, Record<string, StoredFile>>;

export type UpdateFileParams = {
  assetId: number;
  messageId: string;
  state: FileDownloadStatus;
};

type FileDownloadStore = {
  downloadsByMessage: DownloadsByMessage;
  fileAutomaticDownload: boolean;
  setFileAutomaticDownload: (value: boolean) => void;
  subscribeFile: (params: DownloadAsset) => void;
  subscribeFileAsDone: (params: DownloadAsset) => void;
  subscribeFileAsDoneBatch: (data: DownloadAsset[]) => void;
  subscribeFilesBatch: (data: DownloadAsset[]) => void;
  setProgress: (data: {
    assetId: number;
    messageId: string;
    progress: number;
  }) => void;
  updateFileState: (data: UpdateFileParams) => void;
  updateFileStateBatch: (data: UpdateFileParams[]) => void;
  clear: () => void;
  updateAssetsBatch: (
    data: {assetId: number; messageId: string; asset: Asset}[],
  ) => void;
};

export const useAttachments = create<FileDownloadStore>((set, get) => ({
  downloadsByMessage: new Map(),
  assets: {},
  fileAutomaticDownload: false,
  subscribeFile: ({assetId, uri, messageId, remoteUrl, filename, asset}) =>
    set(state => {
      const downloadsByMessage = new Map(state.downloadsByMessage);
      const messageData = downloadsByMessage.get(messageId) || {};

      const currentItem = messageData[assetId];

      // ignore if the file already downloaded
      if (currentItem?.state === FileDownloadStatus.DONE) {
        return state;
      }

      const newItem: StoredFile = {
        assetId,
        uri,
        remoteUrl,
        filename,
        state: FileDownloadStatus.IDLE,
        progress: 0,
        asset,
      };

      downloadsByMessage.set(messageId, {
        ...messageData,
        [assetId]: newItem,
      });

      return {
        downloadsByMessage,
      };
    }),
  subscribeFileAsDone: ({
    assetId,
    uri,
    messageId,
    remoteUrl,
    filename,
    asset,
  }) =>
    set(state => {
      const downloadsByMessage = new Map(state.downloadsByMessage);
      const messageData = downloadsByMessage.get(messageId) || {};
      const newItem: StoredFile = {
        assetId,
        uri,
        remoteUrl,
        filename,
        state: FileDownloadStatus.DONE,
        progress: 1,
        asset,
      };

      downloadsByMessage.set(messageId, {
        ...messageData,
        [assetId]: newItem,
      });

      return {
        downloadsByMessage,
      };
    }),

  subscribeFileAsDoneBatch: data => {
    const downloadsByMessage = new Map(get().downloadsByMessage);

    for (const {assetId, uri, messageId, remoteUrl, filename, asset} of data) {
      const messageData = downloadsByMessage.get(messageId) || {};
      const newItem: StoredFile = {
        assetId,
        uri,
        remoteUrl,
        filename,
        state: FileDownloadStatus.DONE,
        progress: 1,
        asset,
      };

      downloadsByMessage.set(messageId, {
        ...messageData,
        [assetId]: newItem,
      });
    }

    set({downloadsByMessage});
  },

  subscribeFilesBatch: data => {
    const downloadsByMessage = new Map(get().downloadsByMessage);

    for (const {assetId, uri, messageId, remoteUrl, filename, asset} of data) {
      const messageData = downloadsByMessage.get(messageId) || {};

      const currentItem = messageData[assetId];

      // ignore if the file already downloaded
      if (currentItem?.state === FileDownloadStatus.DONE) {
        continue;
      }

      const newItem: StoredFile = {
        assetId,
        uri,
        remoteUrl,
        filename,
        state: FileDownloadStatus.IDLE,
        progress: 0,
        asset,
      };

      downloadsByMessage.set(messageId, {
        ...messageData,
        [assetId]: newItem,
      });
    }

    set(() => {
      return {downloadsByMessage};
    });
  },

  // TEMPORARILY DISABLED: nothing in the UI currently consumes per-download
  // progress (the renderers all use `isLoading` boolean + a static spinner,
  // not a progress bar). Every `setProgress` call previously triggered a
  // Zustand state mutation, which re-ran every `useAttachmentData`
  // selector, forcing re-renders of every chat row holding that asset —
  // tens or hundreds of re-renders per download. Leaving the whole
  // implementation in place (commented-out branch below) so we can turn it
  // back on the moment a real progress indicator is reintroduced.
  //
  // Note: the existing logic also guarded against downgrading a DONE entry
  // back to PROGRESS (late progress callback after completion) — that
  // protection is implicitly preserved by this no-op, since we don't touch
  // state at all.
  setProgress: () => {
    /* Re-enable the block below to resume progress tracking:
    set(state => {
      const downloadsByMessage = new Map(state.downloadsByMessage);
      const messageData = downloadsByMessage.get(messageId);

      if (!messageData || !messageData[assetId]) {
        return state;
      }

      const existing = messageData[assetId];
      // Never downgrade a DONE entry — a late progress callback from a
      // finished download must not revert the state to PROGRESS and cause
      // the UI to flash back to loading.
      if (existing.state === FileDownloadStatus.DONE) {
        return state;
      }

      const targetFile: StoredFile = {
        ...existing,
        state: FileDownloadStatus.PROGRESS,
        progress,
      };

      downloadsByMessage.set(messageId, {
        ...messageData,
        [assetId]: targetFile,
      });

      return {
        downloadsByMessage,
      };
    });
    */
  },
  setFileAutomaticDownload: (value: boolean) =>
    set(state => ({
      ...state,
      fileAutomaticDownload: value,
    })),
  updateFileState: ({assetId, messageId, state: newState}) => {
    const item = get().downloadsByMessage.get(messageId)?.[assetId];

    // ignore if the file already downloaded
    if (item?.state === FileDownloadStatus.DONE) {
      return;
    }

    set(state => {
      const downloadsByMessage = new Map(state.downloadsByMessage);

      const messageData = downloadsByMessage.get(messageId);
      if (!messageData || !messageData[assetId]) {
        return state;
      }
      const targetFile: StoredFile = {
        ...(messageData || {})[assetId],
        state: newState,
      };

      downloadsByMessage.set(messageId, {
        ...messageData,
        [assetId]: targetFile,
      });

      return {
        downloadsByMessage,
      };
    });
  },

  updateFileStateBatch: data => {
    const downloadsByMessage = new Map(get().downloadsByMessage);

    for (const {assetId, messageId, state} of data) {
      const messageData = downloadsByMessage.get(messageId);
      if (
        !messageData ||
        !messageData[assetId] ||
        messageData[assetId].state === FileDownloadStatus.DONE // ignore if the file already downloaded
      ) {
        continue;
      }

      const targetFile: StoredFile = {
        ...(messageData || {})[assetId],
        state,
        progress: state === FileDownloadStatus.DONE ? 1 : 0,
      };
      downloadsByMessage.set(messageId, {
        ...messageData,
        [assetId]: targetFile,
      });
    }

    set(() => ({downloadsByMessage}));
  },
  clear: () => {
    set({downloadsByMessage: new Map()});
  },
  updateAssetsBatch: data => {
    const downloadsByMessage = new Map(get().downloadsByMessage);
    for (const {assetId, messageId, asset} of data) {
      const messageData = downloadsByMessage.get(messageId);
      if (!messageData) {
        continue;
      }
      messageData[assetId] = {
        ...messageData[assetId],
        asset,
      };
    }
    set(() => ({downloadsByMessage}));
  },
}));

export const subscribeFilesBatch = (
  data: Parameters<FileDownloadStore['subscribeFilesBatch']>[0],
) => {
  useAttachments.getState().subscribeFilesBatch(data);
};

export const subscribeFileAsDone = (
  params: Parameters<FileDownloadStore['subscribeFile']>[0],
) => {
  useAttachments.getState().subscribeFileAsDone(params);
};

export const subscribeFileAsDoneBatch = (
  params: Parameters<FileDownloadStore['subscribeFileAsDoneBatch']>[0],
) => {
  useAttachments.getState().subscribeFileAsDoneBatch(params);
};

export const getDownloadsByMessageId = (messageId: string) => {
  return useAttachments.getState().downloadsByMessage.get(messageId) || {};
};

export const getDownloadProgress = (assetId: number, messageId: string) => {
  const downloads = getDownloadsByMessageId(messageId);
  return downloads[assetId]?.progress || 0;
};

export const getFileAutomaticDownload = () => {
  return useAttachments.getState().fileAutomaticDownload;
};

export const setDownloadProgress = (
  data: Parameters<FileDownloadStore['setProgress']>[0],
) => {
  useAttachments.getState().setProgress(data);
};

export const updateFileState = (
  data: Parameters<FileDownloadStore['updateFileState']>[0],
) => {
  useAttachments.getState().updateFileState(data);
};

export const updateFileStateBatch = (
  data: Parameters<FileDownloadStore['updateFileStateBatch']>[0],
) => {
  useAttachments.getState().updateFileStateBatch(data);
};

export const useAttachmentData = (messageId?: string, assetId?: number) => {
  const asset = useAttachments(
    useShallow(state =>
      messageId && assetId
        ? state.downloadsByMessage.get(messageId)?.[assetId]
        : undefined,
    ),
  );
  return asset;
};

export const useAttachmentsIds = (messageId: string) =>
  useAttachments(
    useShallow(state =>
      Object.values(state.downloadsByMessage.get(messageId) || {}).map(
        file => file.assetId,
      ),
    ),
  );

export const getAttachmentData = (messageId: string, assetId: number) => {
  const downloadsByMessage = useAttachments.getState().downloadsByMessage;
  const messageData = downloadsByMessage.get(messageId);
  return messageData?.[assetId];
};

export const updateAssetsBatch = (
  data: Parameters<FileDownloadStore['updateAssetsBatch']>[0],
) => {
  useAttachments.getState().updateAssetsBatch(data);
};
