import { spawn } from 'node:child_process'

const REMBG_COMMAND = process.env.REMBG_COMMAND ?? 'rembg'
const REMBG_TIMEOUT_MS = Number(process.env.REMBG_TIMEOUT_MS ?? 120000)

export class BackgroundRemovalError extends Error {
  constructor(message: string, readonly statusCode = 500) {
    super(message)
    this.name = 'BackgroundRemovalError'
  }
}

export function removeImageBackground(input: Buffer) {
  return new Promise<Buffer>((resolve, reject) => {
    const chunks: Buffer[] = []
    const errorChunks: Buffer[] = []
    const child = spawn(REMBG_COMMAND, ['i'], {
      stdio: ['pipe', 'pipe', 'pipe'],
      windowsHide: true
    })

    const timeout = setTimeout(() => {
      child.kill('SIGTERM')
      reject(new BackgroundRemovalError('背景移除超时，请稍后重试', 504))
    }, REMBG_TIMEOUT_MS)

    child.stdout.on('data', (chunk: Buffer) => {
      chunks.push(chunk)
    })

    child.stderr.on('data', (chunk: Buffer) => {
      errorChunks.push(chunk)
    })

    child.on('error', (error: NodeJS.ErrnoException) => {
      clearTimeout(timeout)

      if (error.code === 'ENOENT') {
        reject(new BackgroundRemovalError('服务器未安装 rembg，请先安装后再使用背景移除功能', 503))
        return
      }

      reject(error)
    })

    child.on('close', (code) => {
      clearTimeout(timeout)

      if (code !== 0) {
        const stderr = Buffer.concat(errorChunks).toString('utf8').trim()
        reject(new BackgroundRemovalError(stderr || `rembg 处理失败，退出码 ${code}`))
        return
      }

      const output = Buffer.concat(chunks)

      if (output.length === 0) {
        reject(new BackgroundRemovalError('rembg 未返回图片结果'))
        return
      }

      resolve(output)
    })

    child.stdin.end(input)
  })
}
