#!/bin/bash
# Download all model weights for llmux.
# Run as user llm: bash scripts/download_models.sh
# Requires: pip install huggingface_hub
# Requires: HuggingFace token at ~/.cache/huggingface/token for gated models

set -euo pipefail

MODELS_DIR="${LLMUX_MODELS_DIR:-$HOME/.local/share/llmux_pod/models}"
mkdir -p "$MODELS_DIR"

echo "=== Downloading models to $MODELS_DIR ==="

download_hf() {
    local repo="$1"
    local target="$MODELS_DIR/models--${repo//\//-}"
    if [ -d "$target" ]; then
        echo "SKIP: $repo (already downloaded)"
        return
    fi
    echo "Downloading: $repo"
    huggingface-cli download "$repo" --cache-dir "$MODELS_DIR"
}

download_hf_files() {
    local repo="$1"
    shift
    echo "Downloading specific files from: $repo"
    huggingface-cli download "$repo" "$@" --cache-dir "$MODELS_DIR"
}

# 1. Qwen3.5-9B-FP8
download_hf "lovedheart/Qwen3.5-9B-FP8"

# 2. Qwen3.5-9B-FP8-Uncensored (GGUF files only)
download_hf_files "HauhauCS/Qwen3.5-9B-Uncensored-HauhauCS-Aggressive" \
    "Qwen3.5-9B-Uncensored-HauhauCS-Aggressive-Q8_0.gguf" \
    "mmproj-Qwen3.5-9B-Uncensored-HauhauCS-Aggressive-BF16.gguf"

# 3. Qwen3.5-4B
download_hf "Qwen/Qwen3.5-4B"

# 4. gpt-oss-20b
download_hf "openai/gpt-oss-20b"

# 5. gpt-oss-20b-uncensored
download_hf "aoxo/gpt-oss-20b-uncensored"

# 6. cohere-transcribe (gated — requires accepted terms)
echo "Downloading: CohereLabs/cohere-transcribe-03-2026 (gated)"
download_hf "CohereLabs/cohere-transcribe-03-2026" || \
    echo "WARNING: cohere-transcribe download failed. Have you accepted the terms at https://huggingface.co/CohereLabs/cohere-transcribe-03-2026 ?"

# 7. Chatterbox TTS
echo "Downloading: Chatterbox TTS weights (auto-downloaded by library)"
python3 -c "
from chatterbox.tts import ChatterboxTTS
import os
os.environ['CUDA_VISIBLE_DEVICES'] = ''
print('Downloading Chatterbox default...')
ChatterboxTTS.from_pretrained(device='cpu')
print('Downloading Chatterbox turbo...')
ChatterboxTTS.from_pretrained(device='cpu', variant='turbo')
print('Downloading Chatterbox multilingual...')
ChatterboxTTS.from_pretrained(device='cpu', variant='multilingual')
print('Chatterbox downloads complete.')
" || echo "WARNING: Chatterbox download failed. Check chatterbox-tts installation."

echo ""
echo "=== Download complete ==="
echo "Models directory: $MODELS_DIR"
du -sh "$MODELS_DIR"