opendance-web/web-server/src/storage/storage.service.ts

import { Injectable, Logger, NotFoundException } from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { createHash } from 'node:crypto';
import { createReadStream, existsSync } from 'node:fs';
import { copyFile, mkdir, rename, stat, unlink, writeFile } from 'node:fs/promises';
import { dirname, extname, join, resolve, sep } from 'node:path';
import { In, Repository } from 'typeorm';
import { AssetKind } from '../common/enums';
import { AppConfig } from '../config/configuration';
import { Asset } from '../entities/asset.entity';

/** 上传产物落盘后的结果 */
export interface StoredAsset {
  asset: Asset;
  /** 可直接访问的相对 URL，如 /static/song-audio/2026/09/ab/abcd….mp3 */
  url: string;
}

export function sha256File(path: string): Promise<string> {
  return new Promise((resolvePromise, reject) => {
    const hash = createHash('sha256');
    const stream = createReadStream(path);
    stream.on('error', reject);
    stream.on('data', (chunk) => hash.update(chunk));
    stream.on('end', () => resolvePromise(hash.digest('hex')));
  });
}

/**
 * 文件存储服务。
 *
 * 约定（contracts/model-formats.md §3）：
 * - 存储路径 `${STORAGE_DIR}/{kind}/{yyyy}/{mm}/{sha256前2位}/{sha256}.{ext}`
 * - 相同 sha256 复用同一 assets 记录，不重复落盘
 * - 永不使用用户提供的原始文件名作为路径
 */
@Injectable()
export class StorageService {
  private readonly logger = new Logger(StorageService.name);
  private readonly root: string;
  private readonly tmpDir: string;

  constructor(
    private readonly config: AppConfig,
    @InjectRepository(Asset) private readonly assets: Repository<Asset>,
  ) {
    this.root = resolve(this.config.storageDir);
    this.tmpDir = join(this.root, 'tmp');
  }

  async onModuleInit(): Promise<void> {
    await mkdir(this.tmpDir, { recursive: true });
    this.logger.log(`存储根目录：${this.root}`);
  }

  get temporaryDir(): string {
    return this.tmpDir;
  }

  /** 将相对存储路径解析为绝对路径；拒绝逃逸出存储根目录的路径 */
  absolutePath(asset: Pick<Asset, 'storedPath'>): string {
    const abs = resolve(this.root, asset.storedPath);
    if (abs !== this.root && !abs.startsWith(this.root + sep)) {
      throw new NotFoundException('非法的存储路径');
    }
    return abs;
  }

  /** 资产访问 URL（静态挂载点 /static） */
  url(asset: Pick<Asset, 'storedPath'> | null | undefined): string | null {
    if (!asset) return null;
    return `/static/${asset.storedPath.split(sep).join('/')}`;
  }

  async findById(id: number | null | undefined): Promise<Asset | null> {
    if (id === null || id === undefined) return null;
    return this.assets.findOne({ where: { id } });
  }

  async findMany(ids: Array<number | null | undefined>): Promise<Map<number, Asset>> {
    const wanted = [...new Set(ids.filter((v): v is number => typeof v === 'number' && v > 0))];
    if (wanted.length === 0) return new Map();
    const rows = await this.assets.find({ where: { id: In(wanted) } });
    return new Map(rows.map((row) => [row.id, row]));
  }

  /** 批量解析 URL，避免 N+1 */
  async urlMap(ids: Array<number | null | undefined>): Promise<Map<number, string>> {
    const map = await this.findMany(ids);
    return new Map([...map.entries()].map(([id, asset]) => [id, this.url(asset) as string]));
  }

  /** 保存 multer 落盘的临时文件（去重 + 归档），并清理临时文件 */
  async persistTempFile(
    kind: AssetKind,
    tempPath: string,
    originalName: string,
    mime: string,
  ): Promise<StoredAsset> {
    try {
      const info = await stat(tempPath);
      if (info.size === 0) throw new NotFoundException('上传文件为空');
      const sha256 = await sha256File(tempPath);

      const existing = await this.assets.findOne({ where: { sha256 } });
      if (existing) {
        // 去重命中时**必须确认文件真的还在**：换过 STORAGE_DIR、或只恢复了数据库
        // 而没恢复上传目录时，旧记录会指向一个已经不存在的路径。此时若直接
        // `safeUnlink(tempPath)` 把刚上传的唯一副本删掉，上传仍返回 201，
        // 但后续预处理 / 模型分析 / 下载 / 缩略图全部 ENOENT
        // （实测：动作上传成功后 `POST /api/motions/:id/preprocess` 直接 500）。
        // 这里按**记录里的原路径**补写一份 —— 与 persistBuffer 的兜底一致。
        const absolute = this.absolutePath(existing);
        if (existsSync(absolute)) {
          await this.safeUnlink(tempPath);
          return { asset: existing, url: this.url(existing) as string };
        }
        await mkdir(dirname(absolute), { recursive: true });
        try {
          await rename(tempPath, absolute);
        } catch {
          await copyFile(tempPath, absolute);
          await this.safeUnlink(tempPath);
        }
        return { asset: existing, url: this.url(existing) as string };
      }

      const ext = extname(originalName).toLowerCase();
      const now = new Date();
      const yyyy = String(now.getUTCFullYear());
      const mm = String(now.getUTCMonth() + 1).padStart(2, '0');
      const relDir = join(kind, yyyy, mm, sha256.slice(0, 2));
      const relPath = join(relDir, `${sha256}${ext}`);
      const absDir = join(this.root, relDir);
      await mkdir(absDir, { recursive: true });

      const absTarget = join(absDir, `${sha256}${ext}`);
      try {
        await rename(tempPath, absTarget);
      } catch {
        // 跨设备时 rename 会失败，退化为复制
        await copyFile(tempPath, absTarget);
        await this.safeUnlink(tempPath);
      }

      const entity = this.assets.create({
        kind,
        originalName: originalName.slice(0, 255),
        storedPath: relPath,
        mime: mime || 'application/octet-stream',
        size: info.size,
        sha256,
      });

      try {
        const saved = await this.assets.save(entity);
        return { asset: saved, url: this.url(saved) as string };
      } catch (error) {
        // 并发上传同一文件：唯一索引冲突后回退到已存在记录
        const raced = await this.assets.findOne({ where: { sha256 } });
        if (raced) {
          await this.safeUnlink(absTarget);
          return { asset: raced, url: this.url(raced) as string };
        }
        throw error;
      }
    } catch (error) {
      await this.safeUnlink(tempPath);
      throw error;
    }
  }

  /** 保存内存中的内容（预处理产物等） */
  async persistBuffer(
    kind: AssetKind,
    content: Buffer | string,
    originalName: string,
    mime: string,
  ): Promise<StoredAsset> {
    const sha256 = createHash('sha256').update(content).digest('hex');
    const buf = Buffer.isBuffer(content) ? content : Buffer.from(content, 'utf8');

    const existing = await this.assets.findOne({ where: { sha256 } });
    if (existing) {
      // 去重命中，但要确认文件真的还在：换过 STORAGE_DIR、或只恢复了数据库没恢复
      // 文件时，旧记录会指向一个不存在的路径。此时**按原路径补写一份**，
      // 否则后续的模型分析 / 缩略图 / 下载都会静默失败（曾实测到"全 0 分析结果"）。
      const absolute = this.absolutePath(existing);
      if (!existsSync(absolute)) {
        await mkdir(dirname(absolute), { recursive: true });
        await writeFile(absolute, buf);
      }
      return { asset: existing, url: this.url(existing) as string };
    }

    const ext = extname(originalName).toLowerCase();
    const now = new Date();
    const relDir = join(kind, String(now.getUTCFullYear()), String(now.getUTCMonth() + 1).padStart(2, '0'), sha256.slice(0, 2));
    const relPath = join(relDir, `${sha256}${ext}`);
    await mkdir(join(this.root, relDir), { recursive: true });
    await writeFile(join(this.root, relPath), buf);

    const entity = this.assets.create({
      kind,
      originalName: originalName.slice(0, 255),
      storedPath: relPath,
      mime: mime || 'application/octet-stream',
      size: buf.byteLength,
      sha256,
    });
    const saved = await this.assets.save(entity);
    return { asset: saved, url: this.url(saved) as string };
  }

  /** 读取资产内容（供关键点下载等场景） */
  async readAssetText(asset: Asset): Promise<string> {
    const abs = this.absolutePath(asset);
    const chunks: Buffer[] = [];
    await new Promise<void>((resolvePromise, reject) => {
      const stream = createReadStream(abs);
      stream.on('data', (c) => chunks.push(c as Buffer));
      stream.on('error', reject);
      stream.on('end', () => resolvePromise());
    });
    return Buffer.concat(chunks).toString('utf8');
  }

  async stats(): Promise<{ files: number; totalBytes: number }> {
    const row = await this.assets
      .createQueryBuilder('a')
      .select('COUNT(a.id)', 'files')
      .addSelect('COALESCE(SUM(a.size), 0)', 'totalBytes')
      .getRawOne<{ files: string; totalBytes: string }>();
    return {
      files: Number(row?.files ?? 0),
      totalBytes: Number(row?.totalBytes ?? 0),
    };
  }

  private async safeUnlink(path: string): Promise<void> {
    try {
      await unlink(path);
    } catch {
      /* 忽略：临时文件可能已被移走 */
    }
  }
}