gongxiangp2/lib/performance-survey-image-processing.ts

import sharp from 'sharp'
import { sliceImage } from '@/lib/image-processor'

export const PERFORMANCE_SURVEY_IMAGE_LIMITS = {
  maxSourceCount: 12,
  maxSourceBytes: 20 * 1024 * 1024,
  maxPixelsPerImage: 60_000_000,
  maxWidth: 16_000,
  sliceHeight: 1500,
  maxSlicesPerImage: 24,
  maxSlicesPerBatch: 48,
  targetSliceSizeKB: 500,
} as const

export type PerformanceSurveyImageInput = {
  buffer: Buffer
  mimeType: string
  sourceName: string
}

export type ProcessedPerformanceSurveyImage = {
  buffer: Buffer
  sourceIndex: number
  sliceIndex: number
  sliceCount: number
  sourceName: string
}

type PerformanceSurveyImagePlan = {
  input: PerformanceSurveyImageInput
  sourceIndex: number
  width: number
  height: number
  sliceCount: number
}

export class PerformanceSurveyImageProcessingError extends Error {}

export function getPerformanceSurveyImageSliceCount(height: number) {
  return Math.ceil(height / PERFORMANCE_SURVEY_IMAGE_LIMITS.sliceHeight)
}

export function assertPerformanceSurveyImageDimensions(width: number, height: number, sourceName = '图片') {
  if (!Number.isInteger(width) || !Number.isInteger(height) || width <= 0 || height <= 0) {
    throw new PerformanceSurveyImageProcessingError(`图片「${sourceName}」尺寸无效`)
  }
  if (width > PERFORMANCE_SURVEY_IMAGE_LIMITS.maxWidth) {
    throw new PerformanceSurveyImageProcessingError(
      `图片「${sourceName}」宽度超过 ${PERFORMANCE_SURVEY_IMAGE_LIMITS.maxWidth}px`,
    )
  }
  if (width * height > PERFORMANCE_SURVEY_IMAGE_LIMITS.maxPixelsPerImage) {
    throw new PerformanceSurveyImageProcessingError(`图片「${sourceName}」像素总量过大`)
  }

  const sliceCount = getPerformanceSurveyImageSliceCount(height)
  if (sliceCount > PERFORMANCE_SURVEY_IMAGE_LIMITS.maxSlicesPerImage) {
    throw new PerformanceSurveyImageProcessingError(
      `图片「${sourceName}」过长，最多允许 ${PERFORMANCE_SURVEY_IMAGE_LIMITS.maxSlicesPerImage} 个切片`,
    )
  }
  return sliceCount
}

async function buildPerformanceSurveyImagePlan(
  input: PerformanceSurveyImageInput,
  sourceIndex: number,
): Promise<PerformanceSurveyImagePlan> {
  if (!input.mimeType.startsWith('image/')) {
    throw new PerformanceSurveyImageProcessingError(`文件「${input.sourceName}」不是图片`)
  }
  if (input.buffer.length > PERFORMANCE_SURVEY_IMAGE_LIMITS.maxSourceBytes) {
    throw new PerformanceSurveyImageProcessingError(`图片「${input.sourceName}」超过 20MB`)
  }

  try {
    const metadata = await sharp(input.buffer, {
      limitInputPixels: PERFORMANCE_SURVEY_IMAGE_LIMITS.maxPixelsPerImage,
      sequentialRead: true,
    }).metadata()
    if (!metadata.width || !metadata.height) {
      throw new PerformanceSurveyImageProcessingError(`图片「${input.sourceName}」尺寸无效`)
    }

    return {
      input,
      sourceIndex,
      width: metadata.width,
      height: metadata.height,
      sliceCount: assertPerformanceSurveyImageDimensions(metadata.width, metadata.height, input.sourceName),
    }
  } catch (error) {
    if (error instanceof PerformanceSurveyImageProcessingError) throw error
    throw new PerformanceSurveyImageProcessingError(`图片「${input.sourceName}」无法读取或像素总量过大`)
  }
}

export async function processPerformanceSurveyImages(
  inputs: PerformanceSurveyImageInput[],
): Promise<ProcessedPerformanceSurveyImage[]> {
  if (inputs.length === 0) {
    throw new PerformanceSurveyImageProcessingError('每个案例至少上传 1 张图片')
  }
  if (inputs.length > PERFORMANCE_SURVEY_IMAGE_LIMITS.maxSourceCount) {
    throw new PerformanceSurveyImageProcessingError(
      `单个案例最多上传 ${PERFORMANCE_SURVEY_IMAGE_LIMITS.maxSourceCount} 张图片`,
    )
  }

  const plans: PerformanceSurveyImagePlan[] = []
  let plannedSliceCount = 0
  for (let sourceIndex = 0; sourceIndex < inputs.length; sourceIndex += 1) {
    const plan = await buildPerformanceSurveyImagePlan(inputs[sourceIndex], sourceIndex)
    plannedSliceCount += plan.sliceCount
    if (plannedSliceCount > PERFORMANCE_SURVEY_IMAGE_LIMITS.maxSlicesPerBatch) {
      throw new PerformanceSurveyImageProcessingError(
        `单个案例处理后最多生成 ${PERFORMANCE_SURVEY_IMAGE_LIMITS.maxSlicesPerBatch} 个切片`,
      )
    }
    plans.push(plan)
  }

  const output: ProcessedPerformanceSurveyImage[] = []
  for (const plan of plans) {
    let slices: Buffer[]
    try {
      slices = await sliceImage(
        plan.input.buffer,
        PERFORMANCE_SURVEY_IMAGE_LIMITS.sliceHeight,
        PERFORMANCE_SURVEY_IMAGE_LIMITS.targetSliceSizeKB,
      )
    } catch {
      throw new PerformanceSurveyImageProcessingError(`图片「${plan.input.sourceName}」处理失败`)
    }

    if (slices.length === 0 || slices.length > plan.sliceCount) {
      throw new PerformanceSurveyImageProcessingError(`图片「${plan.input.sourceName}」切片数量异常`)
    }

    slices.forEach((buffer, sliceIndex) => {
      output.push({
        buffer,
        sourceIndex: plan.sourceIndex,
        sliceIndex,
        sliceCount: slices.length,
        sourceName: plan.input.sourceName,
      })
    })
  }

  return output
}