#!/bin/bash
cd "$(dirname "$0")"
if [ -z "$API_KEY" ]; then
  echo api key was not set 1>&2
  exit 1
fi

echo "retrieving new jsons from civitai.."
mkdir -p my/civitai-by-hash/
ls -1 {embeddings/,models/{Stable-diffusion,LyCORIS,Lora,hypernetworks}}/*.{safetensors,ckpt,pt,bin} |
parallel -n 1 -j 16 '
  t="$(echo {} |sed -r "s/\.(safetensors|ckpt|pt|bin)$/\.civitai\.info/")"
  [ -e "$t" ] || {
        h=$(sha256sum {} |tee /dev/stderr | cut -c -64)
        curl https://civitai.com/api/v1/model-versions/by-hash/$h \
          -H "User-Agent: model-metadata-updater, Authorization: \"Bearer '"$API_KEY"'.\", Content-Type: application/json" \
          -X GET | json_pp > "$t"
  }
'

echo "move files to the correct folders (if misplaced).."
# correct folders also updates additional-networks links and json file names.
ls -1 {embeddings/,models/{Stable-diffusion,LyCORIS,Lora,hypernetworks}}/*.civitai.info | while read j; do
  db=$(jq '.model.type' "$j" | tr -d '"')
  case "$db" in
    Checkpoint) t=models/Stable-diffusion;;
    LORA) t=models/Lora;;
    LoCon|LyCORIS) t=models/LyCORIS;;
    Hypernetwork)  t=models/hypernetworks;;
    TextualInversion) t=embeddings;;
    AestheticGradient) t=extensions/stable-diffusion-webui-aesthetic-gradients/aesthetic_embeddings/;;
    null) continue;;
  esac
  f="$(ls -1 "${j%.civitai.info}".{safetensors,ckpt,pt,bin} 2>/dev/null)"
  [ -z "$f" ] && continue;
  echo "$j" | sed -r 's~^((models/)?[^/]+)/(.*)\.civitai\.info$~\1 \3~' | while read d b ; do
    l=
    [ -d extensions/sd-webui-additional-networks ] && l="$(pwd)/extensions/sd-webui-additional-networks/models/lora/$(basename "$f")"
    if [ "$d" != "$t" ]; then
      mv "$d/$b".* "$t/"
      if [ "$db" != AestheticGradient ]; then
        mv "$j" "${j/${d/\//_}/${t/\//_}}"
      else
        mv "$j" "${j/${d/\//_}/esthetic_embeddings}"
      fi
      #if [[ "$f" =~ \.safetensors$ ]]; then
      if [ "$t" = "models/Lora" ]; then
        [ -e "$l" ] && rm "$l"
      fi
    fi
    #if [[ "$f" =~ \.safetensors$ ]]; then
    if [ "$t" = "models/Lora" ]; then
      [ -n "$l" -a -e "$l" ] || ln -s -T "$(pwd)/$t/$(basename "$f")" "$l"
    fi
  done
done

echo "write .txt files with text .."
ls -1 {embeddings/,models/{Stable-diffusion,LyCORIS,Lora,hypernetworks}}/*.civitai.info | while read j; do
  b="${j%.civitai.info}"
  [ -s "$b.txt" ] && continue
  jq '.trainedWords[]' "$j" | sed -n 's/"//g;$!s/$/, /p' > "$b.txt"
  # remove if empty txt
  [ -s "$b.txt" ] || rm "$b.txt"
done

ls -1 {embeddings/,models/{Stable-diffusion,LyCORIS,Lora,hypernetworks}}/*.civitai.info | while read j; do
  b="${j%.civitai.info}"
  [ -s "$b.preview.png" ] && continue
  echo will download $b;
  jq '.images[0] | "\(.width) \(.url)"'  "$j" | head -n 1 | tr -d '"' |while read w u; do
    u="$(echo "$u" | sed -r "s~(/width=)[0-9]+(/.*)\.jpeg$~\1$w\2.png~")"
    curl -H "User-Agent: model-metadata-updater, Authorization: \"Bearer $API_KEY\"" "$u" > "$b.preview.png" && break
    # remove if empty png
    [ -s "$b.preview.png" ] || rm "$b.preview.png"
  done
done