#!/bin/bash

# Export Ollama models for offline deployment
# This script copies Ollama models from your local machine
# so they can be bundled into the Docker image
#
# Required models:
# - qwen3:14b (main chat model)
# - qwen3-embedding:4b (embedding model)

set -e

MODELS_DIR="ollama-models"
OLLAMA_MODELS_PATH="$HOME/.ollama"

echo "=========================================="
echo "Exporting Ollama models for offline deployment"
echo "=========================================="

# Check if Ollama is installed
if ! command -v ollama &> /dev/null; then
    echo "ERROR: Ollama is not installed or not in PATH"
    exit 1
fi

# Check if required models are available
echo ""
echo "Checking for required models..."
MISSING_MODELS=0

if ! ollama list | grep -q "qwen3:14b"; then
    echo "ERROR: qwen3:14b model not found!"
    echo "Please run: ollama pull qwen3:14b"
    MISSING_MODELS=1
fi

if ! ollama list | grep -q "qwen3-embedding:4b"; then
    echo "ERROR: qwen3-embedding:4b model not found!"
    echo "Please run: ollama pull qwen3-embedding:4b"
    MISSING_MODELS=1
fi

if [ $MISSING_MODELS -eq 1 ]; then
    echo ""
    echo "Please download the required models first:"
    echo "  ollama pull qwen3:14b"
    echo "  ollama pull qwen3-embedding:4b"
    exit 1
fi

echo "✓ All required models found"

# Check if Ollama directory exists
if [ ! -d "$OLLAMA_MODELS_PATH" ]; then
    echo "ERROR: Ollama directory not found at $OLLAMA_MODELS_PATH"
    exit 1
fi

# Create export directory
echo ""
echo "Creating export directory: $MODELS_DIR"
rm -rf "$MODELS_DIR"
mkdir -p "$MODELS_DIR"

echo ""
echo "Copying Ollama data from $OLLAMA_MODELS_PATH to $MODELS_DIR..."
echo "This may take several minutes (models are large)..."

# Copy the entire .ollama directory structure
cp -r "$OLLAMA_MODELS_PATH"/* "$MODELS_DIR/"

echo ""
echo "=========================================="
echo "Models exported successfully!"
echo "=========================================="
du -sh "$MODELS_DIR"

echo ""
echo "Directory structure:"
ls -lh "$MODELS_DIR/"

echo ""
echo "Models included:"
if [ -d "$MODELS_DIR/models/manifests/registry.ollama.ai/library" ]; then
    ls -lh "$MODELS_DIR/models/manifests/registry.ollama.ai/library/"
fi

echo ""
echo "=========================================="
echo "Next steps:"
echo "=========================================="
echo "1. Build the all-in-one Docker image:"
echo "   ./build-allinone.sh"
echo ""
echo "2. Or manually:"
echo "   docker build -f Dockerfile.allinone -t system-prompt-optimizer:allinone ."
echo ""
echo "3. Export the image:"
echo "   docker save -o system-prompt-optimizer-allinone.tar system-prompt-optimizer:allinone"
echo ""
echo "4. Transfer to target server:"
echo "   scp system-prompt-optimizer-allinone.tar user@server:/path/"
echo ""
echo "Note: The final Docker image will be very large (10-20GB) due to the models."