const express = require('express');
const fs = require('fs');
const path = require('path');
const cv = require('opencv4nodejs');
const Jimp = require('jimp');
const axios = require('axios');
const sharp = require('sharp');

const app = express();
const port = 7860;

// 从文件夹加载所有分类器
async function loadClassifiersFromFolder(folderPath) {
  const classifiers = [];
  const files = fs.readdirSync(folderPath);
  for (const file of files) {
    const filePath = path.join(folderPath, file);
    const classifier = new cv.CascadeClassifier(filePath);
    classifiers.push(classifier);
  }
  return classifiers;
}

// 下载并转换图像到内存中
async function downloadAndConvertImage(url) {
  try {
    const response = await axios({
      method: 'get',
      url: url,
      responseType: 'arraybuffer'
    });

    const imageBuffer = await sharp(response.data)
      .toFormat('jpeg')
      .toBuffer();

    console.log('Image downloaded and converted successfully.');
    return imageBuffer;
  } catch (error) {
    console.error('Error downloading or converting image:', error);
    throw error;
  }
}

// 图像处理函数
async function processImage(imageBuffer, classifiersFolder, replacementImageBuffer) {
  try {
    const classifiers = await loadClassifiersFromFolder(classifiersFolder);

    // 使用 Jimp 验证图片文件是否有效
    const jimpImage = await Jimp.read(imageBuffer);
    if (!jimpImage) {
      throw new Error('Failed to read the image with Jimp. The image file might be corrupted or not supported.');
    }

    const img = cv.imdecode(imageBuffer);
    if (img.empty) {
      throw new Error('Failed to read the image with OpenCV. The image file might be corrupted or not supported.');
    }

    const grayImg = img.bgrToGray();
    const allFaces = [];
    classifiers.forEach(classifier => {
      const faces = classifier.detectMultiScale(grayImg).objects;
      allFaces.push(...faces);
    });

    const replacementFace = await Jimp.read(replacementImageBuffer);
    allFaces.forEach((faceRect, i) => {
      const resizedReplacementFace = replacementFace.resize(faceRect.width, faceRect.height)
      const faceRegion = img.getRegion(faceRect)
  
      const replacementBuffer = resizedReplacementFace.bitmap.data
      const centerX = faceRect.width / 2
      const centerY = faceRect.height / 2
      const maxRadius = Math.min(centerX, centerY)
      for (let y = 0; y < faceRect.height; y++) {
        for (let x = 0; x < faceRect.width; x++) {
          const distance = Math.sqrt((x - centerX) ** 2 + (y - centerY) ** 2)
          if (distance <= maxRadius) {
            const idx = (y * faceRect.width + x) << 2
            const [r, g, b, a] = replacementBuffer.slice(idx, idx + 4)
  
            if (a > 0) {
              faceRegion.set(y, x, new cv.Vec3(b, g, r))
            }
          }
        }
      }
    })
    const outputBuffer = cv.imencode('.jpg', img);
    return outputBuffer;
  } catch (error) {
    console.error('Error during image processing:', error);
    throw error;
  }
}

// Express接口
app.get('/process-image', async (req, res) => {
  const { imageUrl, replacementImageUrl } = req.query;

  if (!imageUrl || !replacementImageUrl) {
    return res.status(400).send('Both image URL and replacement image URL are required');
  }

  const classifiersFolder = './classifiers';
  let replacementImageBuffer;
  try {
    // 下载并转换输入图片和替换图片到内存中
    const imageBuffer = await downloadAndConvertImage(imageUrl);
    if (replacementImageUrl === 'replace') {
      // 使用默认的本地替换图像路径
      const defaultReplacementImagePath = './replacement_face.png';
      replacementImageBuffer = await Jimp.read(defaultReplacementImagePath);
    } else {
      // 下载并转换替换图片到内存中
      replacementImageBuffer = await downloadAndConvertImage(replacementImageUrl);
    }

    // 处理图片
    const outputBuffer = await processImage(imageBuffer, classifiersFolder, replacementImageBuffer);

    // 返回处理后的图片
    res.set('Content-Type', 'image/jpeg');
    res.send(outputBuffer);
  } catch (error) {
    console.error('Error processing image:', error);
    res.status(500).send('Error processing image');
  }
});

app.listen(port, () => {
  console.log(`Server is running on http://localhost:${port}`);
});