Firelink/src/store/useDownloadStore.ts

import { create } from 'zustand';
import { invokeCommand as invoke } from '../ipc';
import type { DownloadItem } from '../bindings/DownloadItem';
import type { DownloadStatus } from '../bindings/DownloadStatus';
import type { ExtensionDownload } from '../bindings/ExtensionDownload';
import type { Queue } from '../bindings/Queue';
import { useSettingsStore } from './useSettingsStore';

export type { DownloadCategory } from '../utils/downloads';

const getProxyArgs = (settings: ReturnType<typeof useSettingsStore.getState>) => {
  if (settings.proxyMode === 'custom' && settings.proxyHost) {
    return `http://${settings.proxyHost}:${settings.proxyPort}`;
  }
  return null;
};

export const getSiteLogin = (url: string, settings: ReturnType<typeof useSettingsStore.getState>) => {
  try {
    const urlObj = new URL(url);
    const host = urlObj.hostname.toLowerCase();
    for (const login of settings.siteLogins) {
      let pattern = login.urlPattern.toLowerCase().trim();
      if (pattern.startsWith('*.')) {
        const suffix = pattern.substring(2);
        if (host === suffix || host.endsWith('.' + suffix)) return login;
      } else if (pattern.includes('*')) {
        const regex = new RegExp('^' + pattern.replace(/\*/g, '.*') + '$');
        if (regex.test(host)) return login;
      } else if (host === pattern) {
        return login;
      }
    }
  } catch (e) {}
  return null;
};

const syncSystemIntegrations = () => {
  const settings = useSettingsStore.getState();
  const activeCount = useDownloadStore.getState().downloads.filter(d => d.status === 'downloading').length;
  invoke('update_dock_badge', { count: settings.showDockBadge ? activeCount : 0 }).catch(() => {});
  if (settings.preventsSleepWhileDownloading) {
    invoke('set_prevent_sleep', { prevent: activeCount > 0 }).catch(() => {});
  } else {
    invoke('set_prevent_sleep', { prevent: false }).catch(() => {});
  }
};

// Legacy manual speed limit math removed

export type { DownloadStatus };
export const MAIN_QUEUE_ID = '00000000-0000-0000-0000-000000000001';

export type { DownloadItem, Queue };
export type ExtensionDownloadRequest = ExtensionDownload;

export type DeleteModalState = {
  isOpen: boolean;
  downloadId?: string;
};

interface DownloadState {
  downloads: DownloadItem[];
  queues: Queue[];
  isAddModalOpen: boolean;
  pendingAddUrls: string;
  pendingAddReferer: string;
  pendingAddFilename: string;
  selectedPropertiesDownloadId: string | null;
  toggleAddModal: (isOpen: boolean) => void;
  openAddModalWithUrls: (urls: string, referer?: string | null, filename?: string | null) => void;
  handleExtensionDownload: (request: ExtensionDownloadRequest) => void;
  deleteModalState: DeleteModalState;
  openDeleteModal: (downloadId?: string) => void;
  closeDeleteModal: () => void;
  setSelectedPropertiesDownloadId: (id: string | null) => void;
  addDownload: (item: DownloadItem) => void;
  updateDownload: (id: string, updates: Partial<DownloadItem>) => void;
  removeDownload: (id: string, deleteFile?: boolean) => Promise<void>;
  redownload: (id: string) => void;
  processQueue: () => Promise<void>;
  startQueue: (queueId: string) => Promise<number>;
  pauseQueue: (queueId: string) => Promise<number>;
  addQueue: (name: string) => void;
  renameQueue: (id: string, name: string) => void;
  removeQueue: (id: string) => void;
  initDB: () => Promise<void>;
}

let isProcessingQueue = false;

export const useDownloadStore = create<DownloadState>((set, get) => ({
  downloads: [],
  queues: [{ id: MAIN_QUEUE_ID, name: 'Main Queue', isMain: true }],
  isAddModalOpen: false,
  pendingAddUrls: '',
  pendingAddReferer: '',
  pendingAddFilename: '',
  selectedPropertiesDownloadId: null,
  deleteModalState: { isOpen: false },
  openDeleteModal: (downloadId) => set({ deleteModalState: { isOpen: true, downloadId } }),
  closeDeleteModal: () => set({ deleteModalState: { isOpen: false } }),
  toggleAddModal: (isOpen) => set({
    isAddModalOpen: isOpen,
    pendingAddUrls: '',
    pendingAddReferer: '',
    pendingAddFilename: ''
  }),
  openAddModalWithUrls: (urls, referer, filename) => set((state) => {
    const existingUrls = state.isAddModalOpen && state.pendingAddUrls ? state.pendingAddUrls : '';
    const mergedUrls = existingUrls ? `${existingUrls}\n${urls}` : urls;
    return {
      isAddModalOpen: true,
      pendingAddUrls: mergedUrls,
      pendingAddReferer: referer?.trim() || state.pendingAddReferer || '',
      pendingAddFilename: filename?.trim() || state.pendingAddFilename || ''
    };
  }),
  handleExtensionDownload: (request) => {
    const urls = [...new Set(request.urls.map(url => url.trim()).filter(Boolean))];
    if (urls.length === 0) return;

    get().openAddModalWithUrls(
      urls.join('\n'),
      request.referer,
      urls.length === 1 ? request.filename : null
    );
  },
  setSelectedPropertiesDownloadId: (id) => set({ selectedPropertiesDownloadId: id }),
  addDownload: (item) => {
    set((state) => ({ downloads: [...state.downloads, item] }));
    const toSave = { ...item };
    delete toSave.fraction; delete toSave.speed; delete toSave.eta;
    invoke('db_save_download', { id: item.id, status: item.status, queueId: item.queueId, data: JSON.stringify(toSave) }).catch(console.error);
    get().processQueue();
  },
  updateDownload: (id, updates) => {
    let updatedItem: DownloadItem | null = null;
    set((state) => ({
      downloads: state.downloads.map(d => {
        if (d.id === id) {
          const updated = { 
            ...d, 
            ...updates,
            fraction: updates.fraction !== undefined ? updates.fraction : d.fraction
          };
          updatedItem = updated;
          return updated;
        }
        return d;
      })
    }));
    
    if (updatedItem && Object.keys(updates).some(k => !['fraction', 'speed', 'eta'].includes(k))) {
      const toSave = { ...(updatedItem as DownloadItem) };
      delete toSave.fraction; delete toSave.speed; delete toSave.eta;
      invoke('db_save_download', { id: toSave.id, status: toSave.status, queueId: toSave.queueId, data: JSON.stringify(toSave) }).catch(console.error);
    }
    
    // If status changed to something that frees up a slot, process queue
    if (updates.status && ['completed', 'failed', 'paused'].includes(updates.status)) {
      get().processQueue();
      syncSystemIntegrations();
    } else if (updates.status === 'downloading') {
      syncSystemIntegrations();
    }
  },
  removeDownload: async (id, deleteFile = false) => {
    const item = get().downloads.find(d => d.id === id);
    if (item && item.status === 'downloading') {
      try {
        const filepath = item.destination ? `${item.destination}/${item.fileName}` : null;
        await invoke('remove_download', { id, filepath });
      } catch (e) {
        console.error("Failed to terminate download on deletion:", e);
      }
    } else if (item && deleteFile) {
      try {
        const filepath = item.destination ? `${item.destination}/${item.fileName}` : null;
        await invoke('remove_download', { id, filepath });
      } catch (e) {
        console.error("Failed to delete file from disk:", e);
      }
    }
    set((state) => ({
      downloads: state.downloads.filter(d => d.id !== id)
    }));
    invoke('db_delete_download', { id }).catch(console.error);
    get().processQueue();
    syncSystemIntegrations();
  },
  redownload: (id) => {
    let updatedItem: DownloadItem | null = null;
    set((state) => ({
      downloads: state.downloads.map(d => {
        if (d.id === id) {
          const updated: DownloadItem = { ...d, status: 'queued', _dispatched: false, fraction: 0, speed: '-', eta: '-' };
          updatedItem = updated;
          return updated;
        }
        return d;
      })
    }));
    if (updatedItem) {
      const toSave = { ...(updatedItem as DownloadItem) };
      delete toSave.fraction; delete toSave.speed; delete toSave.eta;
      invoke('db_save_download', { id: toSave.id, status: toSave.status, queueId: toSave.queueId, data: JSON.stringify(toSave) }).catch(console.error);
    }
    get().processQueue();
  },
  startQueue: async (queueId) => {
    const runnableIds = get().downloads
      .filter(item => item.queueId === queueId && (item.status === 'queued' || item.status === 'paused' || item.status === 'failed'))
      .map(item => item.id);

    if (runnableIds.length === 0) return 0;

    set((state) => ({
      downloads: state.downloads.map(item =>
        runnableIds.includes(item.id)
          ? { ...item, status: 'queued', _dispatched: false, speed: '-', eta: '-' }
          : item
      )
    }));
    
    const downloadsToSave = get().downloads.filter(item => runnableIds.includes(item.id));
    downloadsToSave.forEach(d => {
      const toSave = { ...d };
      delete toSave.fraction; delete toSave.speed; delete toSave.eta;
      invoke('db_save_download', { id: toSave.id, status: toSave.status, queueId: toSave.queueId, data: JSON.stringify(toSave) }).catch(console.error);
    });

    await get().processQueue();
    return runnableIds.length;
  },
  pauseQueue: async (queueId) => {
    const activeIds = get().downloads
      .filter(item => item.queueId === queueId && item.status === 'downloading')
      .map(item => item.id);

    if (activeIds.length === 0) return 0;

    set((state) => ({
      downloads: state.downloads.map(item =>
        activeIds.includes(item.id)
          ? { ...item, status: 'paused', speed: '-', eta: '-' }
          : item
      )
    }));
    
    const downloadsToSave = get().downloads.filter(item => activeIds.includes(item.id));
    downloadsToSave.forEach(d => {
      const toSave = { ...d };
      delete toSave.fraction; delete toSave.speed; delete toSave.eta;
      invoke('db_save_download', { id: toSave.id, status: toSave.status, queueId: toSave.queueId, data: JSON.stringify(toSave) }).catch(console.error);
    });

    await Promise.all(activeIds.map(id => invoke('pause_download', { id }).catch(() => {})));
    syncSystemIntegrations();
    return activeIds.length;
  },
  addQueue: (name) => {
    const id = crypto.randomUUID();
    const q = { id, name, isMain: false };
    set((state) => ({
      queues: [...state.queues, q]
    }));
    invoke('db_save_queue', { id, data: JSON.stringify(q) }).catch(console.error);
  },
  renameQueue: (id, name) => {
    let updatedQ: Queue | null = null;
    set((state) => ({
      queues: state.queues.map(q => {
        if (q.id === id) {
          const newQ = { ...q, name };
          updatedQ = newQ;
          return newQ;
        }
        return q;
      })
    }));
    if (updatedQ) {
      invoke('db_save_queue', { id, data: JSON.stringify(updatedQ) }).catch(console.error);
    }
  },
  removeQueue: (id) => {
    if (id === MAIN_QUEUE_ID) return;
    
    const affectedDownloads = get().downloads.filter(d => d.queueId === id);
    
    set((state) => ({
      queues: state.queues.filter(q => q.id !== id),
      downloads: state.downloads.map(d => 
        d.queueId === id ? { ...d, queueId: MAIN_QUEUE_ID } : d
      )
    }));
    invoke('db_delete_queue', { id }).catch(console.error);
    
    // Also we need to save the updated downloads to DB
    affectedDownloads.forEach(d => {
      const toSave = { ...d, queueId: MAIN_QUEUE_ID };
      delete toSave.fraction; delete toSave.speed; delete toSave.eta;
      invoke('db_save_download', { id: toSave.id, status: toSave.status, queueId: MAIN_QUEUE_ID, data: JSON.stringify(toSave) }).catch(console.error);
    });
  },
  initDB: async () => {
    try {
      const queuesStr = await invoke('db_get_all_queues');
      const queues = queuesStr.map(q => JSON.parse(q));
      
      const downloadsStr = await invoke('db_get_all_downloads');
      const downloads = downloadsStr.map(d => JSON.parse(d));
      
      set(state => ({
        queues: queues.length > 0 ? queues : state.queues,
        downloads: downloads.length > 0 ? downloads : state.downloads
      }));
      
      // Auto resume downloads that were active
      const active = get().downloads.filter(d => d.status === 'downloading');
      const settings = useSettingsStore.getState();
      active.forEach(item => {
        if (item.isMedia) {
          invoke('start_media_download', {
            id: item.id,
            url: item.url,
            destination: item.destination || '~/Downloads',
            filename: item.fileName,
            formatSelector: item.mediaFormatSelector || null,
            cookieSource: null,
            speedLimit: item.speedLimit || settings.globalSpeedLimit || null,
            username: item.username || null,
            password: item.password || null,
            headers: item.headers || null,
            proxy: null,
            userAgent: null,
            maxTries: null
          }).catch(console.error);
        } else {
          invoke('start_download', {
            id: item.id,
            url: item.url,
            destination: item.destination || '~/Downloads',
            filename: item.fileName,
            connections: item.connections ?? null,
            speedLimit: item.speedLimit || settings.globalSpeedLimit || null,
            username: item.username || null,
            password: item.password || null,
            headers: item.headers || null,
            checksum: item.checksum || null,
            cookies: item.cookies || null,
            mirrors: item.mirrors || null,
            userAgent: null,
            maxTries: null,
            proxy: null
          }).catch(console.error);
        }
      });
      
      void get().processQueue();
    } catch (e) {
      console.error("Failed to init DB", e);
    }
  },
  processQueue: async () => {
    if (isProcessingQueue) return;
    isProcessingQueue = true;
    
    try {
      const { downloads, updateDownload } = get();
      const settings = useSettingsStore.getState();
      const concurrentLimit = settings.maxConcurrentDownloads || 3;
      const activeCount = downloads.filter(d => d.status === 'downloading').length;
      let availableSlots = concurrentLimit - activeCount;

      if (availableSlots <= 0) return;

      const itemsToStart = downloads.filter(d => d.status === 'queued' && !d._dispatched);
      
      for (const item of itemsToStart) {
        if (availableSlots <= 0) break;
        availableSlots--;

        // Mark as dispatched so we don't send it again on the next pass
        updateDownload(item.id, { _dispatched: true });
        try {
          const login = getSiteLogin(item.url, settings);
          let keychainPassword = null;
          if (login) {
            try {
              keychainPassword = await invoke('get_keychain_password', { id: login.id });
            } catch (e) {
              console.warn("Could not fetch keychain password for login:", e);
            }
          }
          
          const destPath = item.destination || 
                           (settings.downloadDirectories && settings.downloadDirectories[item.category]) || 
                           settings.defaultDownloadPath || 
                           '~/Downloads';

          if (item.isMedia) {
            await invoke('start_media_download', {
              id: item.id,
              url: item.url,
              destination: destPath,
              filename: item.fileName,
              formatSelector: item.mediaFormatSelector || null,
              cookieSource: settings.mediaCookieSource !== 'none' ? settings.mediaCookieSource : null,
              speedLimit: item.speedLimit || settings.globalSpeedLimit || null,
              username: item.username || (login ? login.username : null),
              password: item.password || keychainPassword,
              headers: item.headers || null,
              proxy: getProxyArgs(settings),
              userAgent: settings.customUserAgent || null,
              maxTries: settings.maxAutomaticRetries
            });
          } else {
            await invoke('start_download', {
              id: item.id,
              url: item.url,
              destination: destPath,
              filename: item.fileName,
              connections: item.connections || settings.perServerConnections || null,
              speedLimit: item.speedLimit || settings.globalSpeedLimit || null,
              username: item.username || (login ? login.username : null),
              password: item.password || keychainPassword,
              headers: item.headers || null,
              checksum: item.checksum || null,
              cookies: item.cookies || null,
              mirrors: item.mirrors || null,
              userAgent: settings.customUserAgent || null,
              maxTries: settings.maxAutomaticRetries,
              proxy: getProxyArgs(settings)
            });
          }
        } catch (e) {
          console.error("Failed to start queued download:", e);
          updateDownload(item.id, { status: 'failed' });
        }
      }
    } finally {
      isProcessingQueue = false;
    }
  }
}));