src/routes/api/models/+server.ts

import { json, type RequestEvent } from '@sveltejs/kit';
import prisma from '$lib/prisma';

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

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

export async function GET(request : RequestEvent) {
  const token = request.cookies.get('hf_access_token')
  let IS_ADMIN = false

  if (token) {
    const user = await tokenIsAvailable(token)
    if (user) {
      IS_ADMIN = process?.env?.SECRET_HF_ADMIN ? process?.env?.SECRET_HF_ADMIN.includes(user.sub) : false
    }
  }

  const page = parseInt(request.url.searchParams.get('page') || '0')
  const filter = request.url.searchParams.get('filter') || 'hotest'
  const search = request.url.searchParams.get('search') || ''
  const limit = parseInt(request.url.searchParams.get('limit') || '20')
  const base_model = request.url.searchParams.get('base_model') || undefined

  const orderBy: Record<string, string> = {}
  if (filter === 'hotest') {
    orderBy['likes7d'] = 'desc'
  } else if (filter === 'likes') {
    orderBy['likes'] = 'desc'
  } else {
    orderBy['createdAt'] = 'desc'
  }

  let base_model_mapped: string[] | undefined = undefined;
  if (base_model) {
    switch (base_model) {
      case "sd3":
        base_model_mapped = ['stabilityai/stable-diffusion-3-medium-diffusers']
        break;
      case "sdxl":
        base_model_mapped = ['stabilityai/stable-diffusion-xl-base-1.0']
        break;
      case "sd1":
        base_model_mapped = ['CompVis/stable-diffusion-v1-4', 'runwayml/stable-diffusion-v1-5']
        break;
      case "flux":
        base_model_mapped = ['black-forest-labs/FLUX.1-dev']
        break;
    }
  }

  const only_not_public = filter === 'staff_only';

  const cards = await prisma.model.findMany({
    where: {
      ...(
        !IS_ADMIN ? { isPublic: true } : only_not_public ? { isPublic: false } : {}
      ),
      ...(base_model_mapped ? { base_model: { in: base_model_mapped } } : {}),
      OR: [
        { id: { contains: search } },
      ]
    },
    select: {
      id: true,
      likes: true,
      downloads: true,
      likes7d: true,
      image: true,
      isPublic: true,
      gallery: {
        select: {
          id: true,
          image: true,
        },
        where: {
          isPublic: true
        },
        orderBy: {
          createdAt: 'desc'
        },
        take: 1
      },
    },
    orderBy: orderBy,
    skip: page * limit,
    take: limit,
  })

  const total_reposId = await prisma.model.count({
    where: {
      ...(IS_ADMIN ? {} : { isPublic: true }),
      ...(base_model_mapped ? { base_model: { in: base_model_mapped } } : {}),
      OR: [
        { id: { contains: search } },
      ]
    },
  })

  return json({
    cards,
    total_items: total_reposId
  })
}

export async function PATCH({ request } : 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 models = await prisma.model.findMany({
    where: {
      isPublic: true
    }
  });

  let total_updates = 0;
  for (const model of models) {
    const hugging_face_request = await fetch(`https://huggingface.co/api/models?id=${model.id}&sort=likes7d`)
    const hugging_face_model = await hugging_face_request.json()?.catch(() => {})

    let hugging_face_model2 = undefined;

    if (!model.instance_prompt) {
      const hugging_face_request2 = await fetch(`https://huggingface.co/api/models/${model.id}`)
      hugging_face_model2 = await hugging_face_request2.json()?.catch(() => {})

    }
    
    if (!hugging_face_model?.[0]) {
      continue;
    }

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

    await prisma.model.update({
      where: {
        id: model.id
      },
      data: {
        likes: hugging_face_model?.[0]?.likes,
        downloads: hugging_face_model?.[0]?.downloads,
        likes7d: hugging_face_model?.[0]?.trendingScore,
        id: hugging_face_model?.[0]?.id,
        ...(hugging_face_model2?.cardData?.instance_prompt ? {
          instance_prompt: hugging_face_model2?.cardData?.instance_prompt,
        } : {}
        ),
        ...(base_model ? {
          base_model: base_model
        } : {}
        )
      }
    })
    .then(() => {
      total_updates++
    })
    .catch(() => {})
  }

  return json({
    message: `Updated ${total_updates} models`
  })
}