import { mkdtempSync, readFileSync, rmSync, writeFileSync } from 'node:fs'
import { tmpdir } from 'node:os'
import { extname, join } from 'node:path'
import { runFfmpeg } from '~/server/utils/ffmpeg'
import { imageDimensions } from '~/server/utils/resolution'
import { padExtents, parseImageAspect, type ImageAspectId } from '~/utils/imageAspect'

export type PaddedImage = { filename: string; data: Buffer; type?: string }

function sourceExt(filename: string, type?: string) {
  const fromName = extname(filename || '').toLowerCase()
  if (fromName === '.jpg' || fromName === '.jpeg' || fromName === '.png' || fromName === '.webp') return fromName
  if (type?.includes('jpeg')) return '.jpg'
  if (type?.includes('webp')) return '.webp'
  return '.png'
}

export async function applyImageFill(opts: {
  image: PaddedImage
  mask?: PaddedImage | null
  aspect: string
}): Promise<{ image: PaddedImage; mask?: PaddedImage; padded: boolean; aspect: ImageAspectId }> {
  const aspect = parseImageAspect(opts.aspect)
  if (aspect === 'auto') {
    return { image: opts.image, mask: opts.mask || undefined, padded: false, aspect }
  }
  const size = imageDimensions(opts.image.data)
  if (!size?.width || !size.height) {
    return { image: opts.image, mask: opts.mask || undefined, padded: false, aspect }
  }
  const box = padExtents(size.width, size.height, aspect)
  if (!box) {
    return { image: opts.image, mask: opts.mask || undefined, padded: false, aspect }
  }

  const dir = mkdtempSync(join(tmpdir(), 'aigen-pad-'))
  try {
    const srcPath = join(dir, `src${sourceExt(opts.image.filename, opts.image.type)}`)
    const outPath = join(dir, 'out.png')
    writeFileSync(srcPath, opts.image.data)
    await runFfmpeg([
      '-y',
      '-i', srcPath,
      '-vf', `pad=${box.canvasW}:${box.canvasH}:${box.x}:${box.y}:black`,
      '-frames:v', '1',
      outPath
    ])
    const padded: PaddedImage = {
      filename: 'padded.png',
      data: readFileSync(outPath),
      type: 'image/png'
    }

    if (!opts.mask?.data?.length) {
      return { image: padded, padded: true, aspect }
    }

    const userPath = join(dir, `mask${sourceExt(opts.mask.filename, opts.mask.type)}`)
    const maskPath = join(dir, 'pad-mask.png')
    writeFileSync(userPath, opts.mask.data)
    await runFfmpeg([
      '-y',
      '-f', 'lavfi',
      '-i', `color=c=white:s=${box.canvasW}x${box.canvasH}:d=1:r=1`,
      '-i', userPath,
      '-filter_complex', `[1]scale=${box.srcW}:${box.srcH}[m];[0][m]overlay=${box.x}:${box.y}`,
      '-frames:v', '1',
      maskPath
    ])
    return {
      image: padded,
      mask: { filename: 'pad-mask.png', data: readFileSync(maskPath), type: 'image/png' },
      padded: true,
      aspect
    }
  } finally {
    rmSync(dir, { recursive: true, force: true })
  }
}