app/api/route.ts

import { PrismaClient } from '@prisma/client'

import list_styles from "@/assets/list_styles.json"
import { UploaderDataset } from './uploader'
import { isTextNSFW } from '@/utils/checker/prompt'
import { isImageNSFW } from '@/utils/checker/image'

const prisma = new PrismaClient()

export async function POST(
  request: Request,
) {
  const global_headers = {
    Authorization: `Bearer ${process.env.NEXT_APP_HF_TOKEN}`,
    'Content-Type': 'application/json',
    ['x-use-cache']: "0"
  }

  const { inputs, style } = await request.json()
  const findStyle = list_styles.find((item) => item.name === style)

  const textIsNSFW = await isTextNSFW(inputs, global_headers)
  if (textIsNSFW) return Response.json({ status: 401, ok: false, message: "Prompt is not safe for work." });

  const response = await fetch(`${process.env.NEXT_PUBLIC_API_URL}/models/stabilityai/stable-diffusion-xl-base-1.0`, {
    method: 'POST',
    body: JSON.stringify({
      inputs: findStyle?.prompt.replace("{prompt}", inputs) ?? inputs,
      negative_prompt: findStyle?.negative_prompt ?? "",
    }),
    headers: global_headers,
  })

  const res = await response.clone().json().catch(() => ({}));
  if (res?.error) return Response.json({ status: response.status, ok: false, message: res.error });

  const blob = await response.blob()
  
  const imageIsNSFW = await isImageNSFW(blob, global_headers)
  if (imageIsNSFW) return Response.json({ status: 401, ok: false, message: "Image is not safe for work." });

  const name = Date.now() + `-${inputs.replace(/[^a-zA-Z0-9]/g, '-').slice(0, 10).toLowerCase()}`
  const { ok, message } = await UploaderDataset(blob, name)

  if (!ok) return Response.json({ status: 500, ok: false, message });

  const new_image = await prisma.collection.create({
    data: {
      prompt: inputs,
      file_name: name,
    },
  })
  
  return Response.json({ image: new_image, status: 200, ok: true });

}