import Koa from 'koa'
import http from 'axios'
import sharp from 'sharp'
import tf from '@tensorflow/tfjs-node'
import nsfw from 'nsfwjs'

import { extractColors } from './extract-colors.js'
import FastAverageColor from './fast-average-color.js'
import { phash } from './phash.js'

const app = new Koa()

app.use(async (ctx) => {
  const { data: image_buffer } = await http({
    method: 'GET',
    url: String(ctx.request.query.url).replace(
      'i.pximg.net',
      'i.pixiv.re',
    ),
    responseType: 'arraybuffer',
  })
  const { data: image_raw, info } = await sharp(image_buffer, {
    animated: true,
  })
    .ensureAlpha()
    .raw()
    .toBuffer({ resolveWithObject: true })
  const image_phash = await phash(image_buffer)
  const colors = await extractColors(
    { data: image_raw, width: info.width, height: info.height },
    {
      pixels: info.width * info.height,
      distance: 0.05,
      splitPower: 6,
      colorValidator: (red, green, blue, alpha = 255) => alpha >= 255,
      hueDistance: 0,
      saturationDistance: 1,
      lightnessDistance: 1,
    },
  )
  const fac = new FastAverageColor()
  const {
    value: [r, g, b, a],
    hex,
    isLight: blight,
  } = fac.prepareResult(fac.getColorFromArray4(image_raw))
  const image_resize = await sharp(image_buffer, { animated: true })
    .resize({
      width: 128,
      height: 128,
      fit: 'inside',
    })
    .gif({ dither: 0 })
    .toBuffer()
  // const model = await nsfw.load(`https://raw.githubusercontent.com/infinitered/nsfwjs/refs/heads/master/models/mobilenet_v2_mid/model.json`, { type: 'graph' })
  //   // const model = await nsfw.load(`https://raw.githubusercontent.com/infinitered/nsfwjs/refs/heads/master/models/mobilenet_v2/model.json`)
  //   // const model = await nsfw.load(`https://s3.amazonaws.com/ir_public/nsfwjscdn/TFJS_nsfw_mobilenet/tfjs_quant_nsfw_mobilenet/model.json`)
  const model = await nsfw.load(`https://raw.githubusercontent.com/infinitered/nsfwjs/refs/heads/master/models/inception_v3/model.json`, { size: 299 })
  const predictions = await model.classifyGif(image_resize, { fps: 1 })
  const detail = new Map()
  predictions
    .flat(Infinity)
    .sort((a, b) => a.probability - b.probability)
    .forEach((i) => {
      detail.set(
        i.className.toLocaleLowerCase(),
        Math.floor(i.probability * 100),
      )
    })
  const fail = ['hentai', 'porn', 'sexy']
  const result = [...detail]
    .filter((i) => i.pop() > 7)
    .some((i) => fail.includes(i.pop()))
  tf.dispose()
  model.model.dispose()
  ctx.body = {
    nsfw: Object.assign({ result }, Object.fromEntries(detail)),
    colors: colors
      .sort((a, b) => b.area - a.area)
      .slice(0, 3)
      .map((i) => {
        // hex: `#${Number(十进制).toString(16).padStart(6, 0)}`
        return {
          color: parseInt(String(i.hex).replace('#', ''), 16),
          ratio: Number(i.area.toFixed(2)),
        }
      }),
    color: hex,
    // blight,
    phash: image_phash,
  }
})

app.listen(process.env.PORT || 7860)