#!/bin/bash

# Qdrant Vector Database Installation Script
# This script installs and configures Qdrant on Ubuntu 22.04

set -e

# Update system packages
apt-get update
apt-get upgrade -y

# Install required packages
apt-get install -y curl wget gnupg2 software-properties-common apt-transport-https ca-certificates

# Install Docker
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg
echo "deb [arch=amd64 signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null
apt-get update
apt-get install -y docker-ce docker-ce-cli containerd.io

# Start and enable Docker
systemctl start docker
systemctl enable docker

# Create qdrant user and directories
useradd -r -s /bin/false qdrant || true
mkdir -p /opt/qdrant/storage
mkdir -p /opt/qdrant/config
chown -R qdrant:qdrant /opt/qdrant

# Create Qdrant configuration file
cat > /opt/qdrant/config/config.yaml << EOF
service:
  host: 0.0.0.0
  http_port: 6333
  grpc_port: 6334
  enable_cors: true

storage:
  storage_path: /qdrant/storage
  snapshots_path: /qdrant/snapshots
  on_disk_payload: true

cluster:
  enabled: false

telemetry:
  disabled: true

log_level: INFO
EOF

# Create API key configuration if provided
if [ -n "${qdrant_api_key}" ] && [ "${qdrant_api_key}" != "" ]; then
cat >> /opt/qdrant/config/config.yaml << EOF

service:
  api_key: "${qdrant_api_key}"
EOF
fi

# Create systemd service for Qdrant
cat > /etc/systemd/system/qdrant.service << EOF
[Unit]
Description=Qdrant Vector Database
After=docker.service
Requires=docker.service

[Service]
Type=simple
User=root
ExecStartPre=-/usr/bin/docker stop qdrant
ExecStartPre=-/usr/bin/docker rm qdrant
ExecStart=/usr/bin/docker run --name qdrant \
  -p 6333:6333 \
  -p 6334:6334 \
  -v /opt/qdrant/storage:/qdrant/storage:z \
  -v /opt/qdrant/config/config.yaml:/qdrant/config/production.yaml:z \
  qdrant/qdrant:latest
ExecStop=/usr/bin/docker stop qdrant
Restart=always
RestartSec=10

[Install]
WantedBy=multi-user.target
EOF

# Pull Qdrant Docker image
docker pull qdrant/qdrant:latest

# Enable and start Qdrant service
systemctl daemon-reload
systemctl enable qdrant
systemctl start qdrant

# Install monitoring tools
apt-get install -y htop iotop nethogs

# Create a simple health check script
cat > /opt/qdrant/health_check.sh << 'EOF'
#!/bin/bash
response=$(curl -s -o /dev/null -w "%%{http_code}" http://localhost:6333/health)
if [ "$response" = "200" ]; then
    echo "Qdrant is healthy"
    exit 0
else
    echo "Qdrant is not responding properly (HTTP $response)"
    exit 1
fi
EOF

chmod +x /opt/qdrant/health_check.sh

# Set up log rotation for Docker logs
cat > /etc/logrotate.d/docker << EOF
/var/lib/docker/containers/*/*.log {
    rotate 7
    daily
    compress
    size=1M
    missingok
    delaycompress
    copytruncate
}
EOF

# Configure firewall (ufw)
ufw --force enable
ufw allow ssh
ufw allow 6333/tcp  # Qdrant HTTP API
ufw allow 6334/tcp  # Qdrant gRPC API

# Create a simple backup script
cat > /opt/qdrant/backup.sh << 'EOF'
#!/bin/bash
BACKUP_DIR="/opt/qdrant/backups"
DATE=$(date +%Y%m%d_%H%M%S)
mkdir -p $BACKUP_DIR

# Create snapshot via API
curl -X POST "http://localhost:6333/snapshots" \
  -H "Content-Type: application/json" \
  -d '{"snapshot_name": "backup_'$DATE'"}'

# Copy storage directory
tar -czf $BACKUP_DIR/qdrant_storage_$DATE.tar.gz -C /opt/qdrant storage/

# Keep only last 7 backups
find $BACKUP_DIR -name "*.tar.gz" -mtime +7 -delete

echo "Backup completed: $DATE"
EOF

chmod +x /opt/qdrant/backup.sh

# Set up daily backup cron job
echo "0 2 * * * root /opt/qdrant/backup.sh >> /var/log/qdrant_backup.log 2>&1" >> /etc/crontab

# Wait for Qdrant to be ready
echo "Waiting for Qdrant to start..."
for i in {1..30}; do
    if curl -s http://localhost:6333/health > /dev/null; then
        echo "Qdrant is ready!"
        break
    fi
    echo "Waiting... ($i/30)"
    sleep 10
done

# Create a default collection for image vectors
curl -X PUT "http://localhost:6333/collections/image_vectors" \
  -H "Content-Type: application/json" \
  -d '{
    "vectors": {
      "size": 512,
      "distance": "Cosine"
    },
    "optimizers_config": {
      "default_segment_number": 2
    },
    "replication_factor": 1
  }'

echo "Qdrant installation and configuration completed!"
echo "Qdrant is accessible at:"
echo "  HTTP API: http://$(curl -s ifconfig.me):6333"
echo "  gRPC API: http://$(curl -s ifconfig.me):6334"
echo "Health check: /opt/qdrant/health_check.sh"
echo "Backup script: /opt/qdrant/backup.sh"