#!/usr/bin/env bash
# 把模型注册表里的模型下载到应用数据目录（QStandardPaths::AppDataLocation）。
#
# 用法:
#   ./fetch-models.sh                 # 列出所有模型和安装状态
#   ./fetch-models.sh <id> [...]      # 下载指定模型
#   ./fetch-models.sh --all           # 下载全部
#
# 目录规则（与程序里 ModelStore 保持一致）：
#   ${XDG_DATA_HOME:-~/.local/share}/live-caption/models/<id>/
set -euo pipefail

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
REG="$ROOT/data/models.json"
DEST="${XDG_DATA_HOME:-$HOME/.local/share}/live-caption/models"

[ -f "$REG" ] || { echo "找不到注册表: $REG" >&2; exit 1; }

ids=$(python3 -c "
import json,sys
d=json.load(open('$REG'))
print(' '.join(m['id'] for m in d['models']))")

show_status() {
  python3 - "$REG" "$DEST" <<'PY'
import json, os, sys
reg, dest = sys.argv[1], sys.argv[2]
d = json.load(open(reg))
for m in d['models']:
    mid = m['id']; mdir = os.path.join(dest, mid)
    files = m['files']
    have = sum(1 for f in files.values() if os.path.exists(os.path.join(mdir, f['name'])))
    total = len(files)
    mark = "已安装" if have == total else ("部分" if have else "未安装")
    langs = ",".join(m['languages'][:4]) + ("…" if len(m['languages']) > 4 else "")
    print(f"  [{have}/{total} {mark}] {mid:22s} {langs:18s} {m.get('license','')}")
    print(f"        {m['name']}")
PY
}

download_one() {
  local id="$1"
  local dir="$DEST/$id"
  mkdir -p "$dir"
  echo "==> $id → $dir"
  python3 - "$REG" "$id" <<'PY' | while IFS=$'\t' read -r name url; do
import json, sys
d = json.load(open(sys.argv[1])); mid = sys.argv[2]
m = next(x for x in d['models'] if x['id'] == mid)
for f in m['files'].values():
    print(f"{f['name']}\t{f['url']}")
PY
    if [ -f "$dir/$name" ]; then
      echo "    已存在，跳过: $name"
      continue
    fi
    echo "    下载 $name"
    if command -v aria2c >/dev/null 2>&1; then
      aria2c -x16 -s16 -k4M --summary-interval=0 -d "$dir" -o "$name" "$url" >/dev/null
    else
      curl -fL --progress-bar -o "$dir/$name" "$url"
    fi
  done
  echo "    完成"
}

case "${1:-}" in
  "" ) echo "模型存储目录: $DEST"; echo; show_status ;;
  --all) shift; for id in $ids; do download_one "$id"; done ;;
  *) for id in "$@"; do
       echo "$ids" | tr ' ' '\n' | grep -qx "$id" || { echo "未知模型: $id" >&2; exit 1; }
       download_one "$id"
     done ;;
esac