#!/bin/zsh

# Define the file to store the update timestamps
TIMESTAMP_FILE="$HOME/.ollama_model_timestamps"

# Function to fetch and sort Ollama models by size, excluding recently updated models
function fetch_sorted_models() {
    ollama list | tail -n +2 | awk -F"\t" '
    !/seconds ago|minutes ago|hours ago/ && $1 ~ /:/ {print $1 "\t" $3}' | sort -k2 -h
}

# Function to update a model if it hasn't been updated in the past hour
function update_model() {
    local model_name=$1
    local current_time=$(date +%s)
    local last_update_time=$(grep "^$model_name " "$TIMESTAMP_FILE" 2>/dev/null | awk '{print $2}')

    # Check if the model was updated in the past hour
    if [[ -n "$last_update_time" && $((current_time - last_update_time)) -lt 3600 ]]; then
        echo "Skipping update for model: $model_name (updated less than an hour ago)"
        return
    fi

    echo "Attempting to update model: $model_name"
    if ollama pull "$model_name"; then
        # Remove any previous entry for the model
        grep -v "^$model_name " "$TIMESTAMP_FILE" > "$TIMESTAMP_FILE.tmp" && mv "$TIMESTAMP_FILE.tmp" "$TIMESTAMP_FILE"
        # Add new timestamp entry
        echo "$model_name $current_time" >> "$TIMESTAMP_FILE"
        echo "Update for model: $model_name completed."
    else
        echo "Update for model: $model_name failed."
    fi
}

# Function to update Ollama models
function update_ollama_models() {
    local sorted_models
    echo "Fetching and sorting Ollama models by size..."
    sorted_models=$(fetch_sorted_models)
    echo "$sorted_models" | while read -r line; do
        local model_name=$(echo $line | awk '{print $1}')
        if [[ "$model_name" =~ : ]]; then
            update_model "$model_name"
        else
            # echo "Skipping invalid model name: $model_name"
        fi
    done

    echo "Updated list of Ollama models:"
    ollama list | tail -n +2 | awk -F"\t" '{printf $1 "\t" $3 "\t" $4 "\t" $5 RS}'
    echo "Ollama models update complete."
}

update_ollama_models