src/routes/api/models/[id]/+server.ts

import { json, type RequestEvent } from '@sveltejs/kit';
import puppeteer from 'puppeteer';
import { promises } from 'fs';
import { randomUUID } from 'crypto';

import prisma from '$lib/prisma';
import { tokenIsAvailable } from '$lib/utils';

/** @type {import('./$types').RequestHandler} */

export async function GET({ url, params } : RequestEvent) {
  const id = params.id?.replace("@", "/")

  const full = Boolean(url.searchParams.get('full')) ?? false
  const metadata = Boolean(url.searchParams.get('metadata')) ?? false

  const model = await prisma.model.findFirst({
    where: {
      id,
    },
    select: full ? {
      id: true,
      likes: true,
      downloads: true,
      image: true,
      instance_prompt: true,
      metadata: true,
      base_model: true,
      isPublic: true,
      gallery: {
        select: {
          id: true,
          prompt: true,
          image: true,
          createdAt: true,
        },
        where: {
          isPublic: true
        },
        orderBy: {
          createdAt: 'desc'
        },
        take: 10
      },
      comments: {
        select: {
          id: true,
          createdAt: true,
          text: true,
          user: {
            select: {
              id: true,
              name: true,
              sub: true,
              picture: true,
              preferred_username: true,
            }
          }
        }
      }
    } : {
      instance_prompt: true,
      image: true,
      id: true,
      metadata: true,
      isPublic: true,
      gallery: {
        select: {
          image: true,
        },
        where: {
          isPublic: true
        },
        orderBy: {
          createdAt: 'desc'
        },
        take: 1
      },
    }
  })

  if (!model) {
    return json({
      error: {
        token: "Model params is required"
      }
    }, { status: 401 })
  }

  let infos: Record<string, string | string[]> = {};

  if (full) {
    const hf_model_request = await fetch(`https://huggingface.co/api/models/${id}`, {
      headers: {
        "Authorization": `Bearer ${process.env.SECRET_HF_TOKEN}`
      }
    })
    const hf_model_response = await hf_model_request.json();
    infos = {
      base_model: hf_model_response?.cardData?.base_model,
      license: hf_model_response?.cardData?.license,
      tags: hf_model_response?.cardData?.tags,
    }
  }

  // get puppeteer screenshot
  if (!model.metadata && metadata && model?.isPublic) {
    const browser = await puppeteer.launch({
      headless: true,
      defaultViewport: {
        width: 1920,
        height: 1080,
      },
      args: [
        '--no-sandbox',
        '--disable-setuid-sandbox',
        '--disable-dev-shm-usage',
      ],
      ignoreDefaultArgs: ['--disable-extensions']
    }).catch((err) => console.error(err))

    if (browser) {
      const page = await browser.newPage();
      await page.goto("https://" + process.env.SPACE_HOST?.split(",")[0] + '/metadata/models/' + model.id, {
        waitUntil: 'networkidle0',
      });
    
      const buffer = await page.screenshot({
        type: 'png',
      });
    
      const dir = await promises.opendir(process?.env?.PUBLIC_FILE_UPLOAD_DIR as string).catch(() => null)
      if (!dir) await promises.mkdir(process?.env?.PUBLIC_FILE_UPLOAD_DIR as string)
      const file_name_formatted = randomUUID() + "_" + "metadata" + model.id.replace("/", "-") + ".png"
      await promises.writeFile(`${process.env.PUBLIC_FILE_UPLOAD_DIR}/${file_name_formatted}`, buffer)
    
      await prisma.model.update({
        where: {
          id,
        },
        data: {
          metadata: file_name_formatted
        }
      })
    
      await browser.close();
    }

  }

  return json({
    model: {
      ...model,
      infos
    }
  })
}

export async function POST({ params, cookies, fetch } : RequestEvent) {
  const token = cookies.get('hf_access_token')
  if (!token) {
    return json({
      error: {
        token: "Token is required"
      }
    }, { status: 401 })
  }

  const user = await tokenIsAvailable(token)
  if (!user || !process.env.SECRET_HF_ADMIN?.includes(user.sub)) {
    return json({
      error: {
        token: "Wrong castle fam :^)"
      }
    }, { status: 401 })
  }

  const id = params.id?.replace("@", "/")

  const model = await prisma.model.findFirst({
    where: {
      id,
    }
  })

  if (!model) {
    return json({
      error: {
        token: "Model not found"
      }
    }, { status: 404 })
  }

  const data = await fetch(`https://huggingface.co/api/models/${model.id}`)
  const hf_model = await data.json()

  const base_model = hf_model.tags.find((tag: string) => tag.startsWith("base_model:"))?.split(":")[1] ?? null

  await prisma.model.update({
    where: {
      id,
    },
    data: {
      isPublic: true,
      instance_prompt: hf_model?.cardData?.instance_prompt,
      ...(base_model ? {
        base_model: base_model
      } : {}
      )
    }
  })

  return json({
    success: true
  })
}

export async function DELETE({ params, cookies } : RequestEvent) {
  const token = cookies.get('hf_access_token')
  if (!token) {
    return json({
      error: {
        token: "Token is required"
      }
    }, { status: 401 })
  }

  const user = await tokenIsAvailable(token)
  if (!user || !process.env.SECRET_HF_ADMIN?.includes(user.sub)) {
    return json({
      error: {
        token: "Wrong castle fam :^)"
      }
    }, { status: 401 })
  }

  const id = params.id?.replace("@", "/")

  const model = await prisma.model.findFirst({
    where: {
      id,
    }
  })

  if (!model) {
    return json({
      error: {
        token: "Model not found"
      }
    }, { status: 404 })
  }

  await prisma.model.update({
    where: {
      id,
    },
    data: {
      isPublic: false,
    }
  })

  return json({
    success: true
  })
}

export async function PATCH({ request, params } : RequestEvent) {
  const headers = Object.fromEntries(request.headers.entries());

  if (headers["x-hf-token"] !== process.env.SECRET_HF_TOKEN) {
    return Response.json({
      message: "Wrong castle fam :^)"
    }, { status: 401 });
  }

  const id = params.id?.replace("@", "/")

  const model = await prisma.model.findFirst({
    where: {
      id,
    }
  })

  if (!model) {
    return json({
      error: {
        token: "Model not found"
      }
    }, { status: 404 })
  }

  const body = await request.json()

  await prisma.model.update({
    where: {
      id,
    },
    data: {
      ...body
    }
  })

  return json({
    success: true
  })
}