zsjz-ai/ai-electron/electron/service/download/downloadService.js

'use strict';

import { app, shell } from 'electron';
import { getDataDir } from 'ee-core/ps';
import { logger } from 'ee-core/log';
import path from 'node:path';
import fs from 'node:fs';
import { randomUUID } from 'node:crypto';
import { pipeline } from 'node:stream/promises';
import { Readable, Transform } from 'node:stream';
import { localSettingsService } from '../setting/localSettings';

/** 本机设置里存下载目录的键名 */
const DOWNLOAD_DIR_KEY = 'downloadDir';
/** 下载清单：放在数据目录，不放进下载目录本身，免得混进用户的文件列表 */
const MANIFEST_NAME = 'downloads.json';
/** 清单最多保留多少条（文件仍在磁盘上，只是不再列出） */
const MAX_ITEMS = 300;

/**
 * 下载服务：把渲染进程拿到的文件流写进**本机下载目录**，并维护一份下载清单。
 *
 * <p>为什么放在主进程：浏览器另存（<a download>）在桌面端会落到系统默认下载目录、也留不下记录；
 * 这里统一落盘到可配置的目录，清单供「下载管理器」展示，并且能直接打开文件/所在目录。
 */
class DownloadService {
  get manifestPath() {
    return path.join(getDataDir(), MANIFEST_NAME);
  }

  /** 生效的下载目录：本机设置优先，未配置用系统「下载」目录 */
  getDownloadDir() {
    const configured = localSettingsService.get(DOWNLOAD_DIR_KEY);
    if (configured) return configured;
    try {
      return app.getPath('downloads');
    } catch {
      return path.join(getDataDir(), 'downloads');
    }
  }

  /** 设置下载目录（必须是已存在的绝对目录；传空表示恢复系统默认） */
  setDownloadDir(dir) {
    const target = typeof dir === 'string' ? dir.trim() : '';
    if (target) {
      if (!path.isAbsolute(target)) throw new Error('只支持绝对路径');
      if (!fs.existsSync(target)) throw new Error(`目录不存在：${target}`);
      if (!fs.statSync(target).isDirectory()) throw new Error('请选择目录，而不是文件');
    }
    localSettingsService.set(DOWNLOAD_DIR_KEY, target);
    return this.getDownloadDir();
  }

  /**
   * 启动时清理「僵死」的下载中记录：进程重启后没有下载线程了，
   * 这类记录若继续显示"下载中"会一直转不完，统一标成失败说明原因。
   */
  reconcileOnce() {
    if (this.reconciled) return;
    this.reconciled = true;
    const items = this.readManifest();
    const stale = items.filter((item) => item.status === 'downloading');
    if (!stale.length) return;
    this.writeManifest(
      items.map((item) =>
        item.status === 'downloading' ? { ...item, status: 'failed', message: '应用重启，下载已中断' } : item,
      ),
    );
  }

  readManifest() {
    try {
      const file = this.manifestPath;
      if (!fs.existsSync(file)) return [];
      const parsed = JSON.parse(fs.readFileSync(file, 'utf8'));
      return Array.isArray(parsed?.items) ? parsed.items.filter((item) => item && typeof item.id === 'string') : [];
    } catch (e) {
      logger.warn('[download] 读取清单失败:', e instanceof Error ? e.message : String(e));
      return [];
    }
  }

  list() {
    this.reconcileOnce();
    return this.readManifest().map((item) => ({
      ...item,
      exists: !!item.savedPath && fs.existsSync(item.savedPath),
    }));
  }

  writeManifest(items) {
    const file = this.manifestPath;
    fs.mkdirSync(path.dirname(file), { recursive: true });
    fs.writeFileSync(file, `${JSON.stringify({ items: items.slice(0, MAX_ITEMS) }, null, 2)}\n`, 'utf8');
  }

  /** 文件名只保留纯名字，挡掉路径分隔符与相对路径，避免写出目录之外 */
  sanitizeFileName(fileName) {
    const base = String(fileName || '').split(/[\\/]/).pop() || '';
    const cleaned = base.replace(/[\u0000-\u001f<>:"|?*]/g, '_').trim();
    return cleaned || `download-${Date.now()}`;
  }

  /** 重名时自动加 (1)(2)…，不覆盖用户已有文件 */
  resolveUniquePath(dir, fileName) {
    const ext = path.extname(fileName);
    const stem = fileName.slice(0, fileName.length - ext.length);
    let candidate = path.join(dir, fileName);
    let index = 0;
    while (fs.existsSync(candidate)) {
      index += 1;
      candidate = path.join(dir, `${stem} (${index})${ext}`);
    }
    return candidate;
  }

  /**
   * 保存一次下载。
   * @param params { fileName: string, mime?: string, data: ArrayBuffer|Uint8Array|string(base64), source?: string }
   */
  save(params = {}) {
    const fileName = this.sanitizeFileName(params?.fileName);
    const raw = params?.data;
    let buffer;
    if (typeof raw === 'string') {
      buffer = Buffer.from(raw, 'base64');
    } else if (raw instanceof Uint8Array) {
      buffer = Buffer.from(raw);
    } else if (raw instanceof ArrayBuffer) {
      buffer = Buffer.from(new Uint8Array(raw));
    } else if (ArrayBuffer.isView(raw)) {
      buffer = Buffer.from(raw.buffer, raw.byteOffset, raw.byteLength);
    } else {
      throw new Error('下载内容为空或格式不支持');
    }
    if (!buffer.length) throw new Error('下载内容为空');

    const dir = this.getDownloadDir();
    fs.mkdirSync(dir, { recursive: true });
    const target = this.resolveUniquePath(dir, fileName);
    fs.writeFileSync(target, buffer);

    const item = {
      id: randomUUID(),
      fileName: path.basename(target),
      savedPath: target,
      dir,
      bytes: buffer.length,
      status: 'done',
      received: buffer.length,
      total: buffer.length,
      mime: typeof params?.mime === 'string' ? params.mime : '',
      source: typeof params?.source === 'string' ? params.source : '',
      at: new Date().toISOString(),
    };
    this.writeManifest([item, ...this.list()]);
    logger.info('[download] 已保存:', target);
    return item;
  }

  /** 进度回调：由 controller 接到 webContents.send 上，把 in-flight 状态推给界面 */
  setNotifier(notifier) {
    this.notifier = typeof notifier === 'function' ? notifier : null;
  }

  notify(item) {
    try {
      this.notifier?.(item);
    } catch (e) {
      logger.warn('[download] 进度推送失败:', e instanceof Error ? e.message : String(e));
    }
  }

  /** 更新清单里的一条记录（不动其它条目），并把最新状态推给界面 */
  patchItem(id, patch) {
    const items = this.list();
    const index = items.findIndex((item) => item.id === id);
    if (index < 0) return null;
    const next = { ...items[index], ...patch };
    items[index] = next;
    this.writeManifest(items);
    this.notify(next);
    return next;
  }

  /**
   * 后台下载：主进程直接拉取 URL 落盘，边下边报进度（渲染进程不承担大文件内存与阻塞）。
   *
   * <p>先写 {@code <名字>.part}，下完再改名，避免下载中就被当成完整文件用；
   * 同名文件自动加 (1)(2)，不覆盖已有文件。
   *
   * @param params { url, fileName?, headers?, source?, mime? }
   */
  start(params = {}) {
    const url = String(params?.url || '').trim();
    if (!url) throw new Error('下载地址为空');
    if (!this.notifier) logger.warn('[download] 未设置进度回调，界面拿不到进度');

    const fileName = this.sanitizeFileName(params?.fileName || '');
    const dir = this.getDownloadDir();
    fs.mkdirSync(dir, { recursive: true });

    const item = {
      id: randomUUID(),
      fileName,
      savedPath: '',
      dir,
      bytes: 0,
      status: 'downloading',
      received: 0,
      total: 0,
      mime: typeof params?.mime === 'string' ? params.mime : '',
      source: typeof params?.source === 'string' ? params.source : '',
      at: new Date().toISOString(),
    };
    this.writeManifest([item, ...this.list()]);
    this.notify(item);

    // 后台跑，不阻塞调用方；失败也只是把这条记录标成 failed
    void this.run(item, url, params);
    return item;
  }

  async run(item, url, params) {
    const partPath = path.join(item.dir, `${item.fileName}.part`);
    let lastNotify = 0;
    try {
      const response = await fetch(url, { headers: (params?.headers || {}) });
      if (!response.ok) {
        // 后端报错时响应体是 JSON（{code,message}），尽量把原因带出来
        let message = `下载失败：HTTP ${response.status}`;
        try {
          const text = await response.text();
          const parsed = JSON.parse(text);
          if (parsed?.message) message = parsed.message;
        } catch {
          // 不是 JSON 就沿用状态码
        }
        throw new Error(message);
      }

      const total = Number(response.headers.get('content-length') || 0);
      this.patchItem(item.id, { total, mime: response.headers.get('content-type') || item.mime });

      if (response.body) {
        const counter = new Transform({
          transform: (chunk, _encoding, callback) => {
            item.received += chunk.length;
            const now = Date.now();
            if (now - lastNotify > 150) {
              lastNotify = now;
              this.patchItem(item.id, { received: item.received });
            }
            callback(null, chunk);
          },
        });
        await pipeline(Readable.fromWeb(response.body), counter, fs.createWriteStream(partPath));
      } else {
        const buffer = Buffer.from(await response.arrayBuffer());
        await fs.promises.writeFile(partPath, buffer);
        item.received = buffer.length;
      }

      const target = this.resolveUniquePath(item.dir, item.fileName);
      await fs.promises.rename(partPath, target);
      const size = fs.statSync(target).size;
      const finished = this.patchItem(item.id, {
        status: 'done',
        savedPath: target,
        fileName: path.basename(target),
        bytes: size,
        received: size,
        total: size,
      });
      logger.info('[download] 已保存:', target);
      return finished;
    } catch (e) {
      const message = e instanceof Error ? e.message : String(e);
      await fs.promises.rm(partPath, { force: true }).catch(() => undefined);
      logger.warn('[download] 下载失败:', message);
      return this.patchItem(item.id, { status: 'failed', message });
    }
  }

  find(id) {
    const item = this.list().find((entry) => entry.id === id);
    if (!item) throw new Error('下载记录不存在');
    return item;
  }

  remove(id) {
    const items = this.list();
    const next = items.filter((item) => item.id !== id);
    if (next.length === items.length) throw new Error('下载记录不存在');
    this.writeManifest(next);
    return { removed: true };
  }

  clear() {
    const count = this.list().length;
    this.writeManifest([]);
    return { cleared: count };
  }

  /** 用系统默认程序打开已下载的文件 */
  async openFile(id) {
    const item = this.find(id);
    if (!fs.existsSync(item.savedPath)) {
      throw new Error(`文件已被移动或删除：${item.savedPath}`);
    }
    const message = await shell.openPath(item.savedPath);
    if (message) throw new Error(message);
    return { path: item.savedPath };
  }

  /** 在文件管理器里定位到已下载的文件 */
  revealFile(id) {
    const item = this.find(id);
    if (!fs.existsSync(item.savedPath)) {
      throw new Error(`文件已被移动或删除：${item.savedPath}`);
    }
    shell.showItemInFolder(item.savedPath);
    return { path: item.savedPath };
  }

  /** 打开下载目录本身 */
  async openDir() {
    const dir = this.getDownloadDir();
    fs.mkdirSync(dir, { recursive: true });
    const message = await shell.openPath(dir);
    if (message) throw new Error(message);
    return { path: dir };
  }
}

export const downloadService = new DownloadService();