#!/bin/bash

# DeepSeek & Open WebUI Installation Script for Debian 13
# Usage: sudo ./install.sh

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Function to print colored output
print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Check if running as root/sudo
check_root() {
    if [[ $EUID -ne 0 ]]; then
        print_error "This script must be run with sudo or as root"
        exit 1
    fi
}

# Update system
update_system() {
    print_info "Updating system packages..."
    apt update -qq
    apt upgrade -y -qq
    print_success "System updated"
}

# Check if Docker is installed
check_docker() {
    if command -v docker &> /dev/null; then
        print_info "Docker is already installed"
        return 0
    else
        print_info "Docker is not installed"
        return 1
    fi
}

# Install Docker
install_docker() {
    print_info "Installing Docker..."
    
    # Install prerequisites
    apt install -y apt-transport-https ca-certificates curl gnupg lsb-release
    
    # Add Docker's official GPG key
    curl -fsSL https://download.docker.com/linux/debian/gpg | gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg
    
    # Set up stable repository
    echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/debian $(lsb_release -cs) stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null
    
    # Install Docker Engine
    apt update
    apt install -y docker-ce docker-ce-cli containerd.io docker-compose-plugin
    
    # Start and enable Docker
    systemctl start docker
    systemctl enable docker
    
    print_success "Docker installed and started"
}

# Check if Ollama is installed
check_ollama() {
    if docker ps -a --format '{{.Names}}' | grep -q "^ollama$"; then
        print_info "Ollama container exists"
        return 0
    else
        print_info "Ollama container does not exist"
        return 1
    fi
}

# Install Ollama
install_ollama() {
    print_info "Installing Ollama..."
    
    # Create volume for Ollama data
    docker volume create ollama_data || true
    
    # Run Ollama container
    docker run -d \
        --name ollama \
        --restart unless-stopped \
        -v ollama_data:/root/.ollama \
        -p 11434:11434 \
        ollama/ollama
    
    # Wait for Ollama to be ready
    print_info "Waiting for Ollama to start..."
    sleep 10
    
    # Check if Ollama is running
    for i in {1..30}; do
        if curl -s http://localhost:11434/api/tags > /dev/null 2>&1; then
            print_success "Ollama is running"
            return 0
        fi
        sleep 2
    done
    
    print_error "Ollama failed to start properly"
    return 1
}

# Get list of installed models
get_installed_models() {
    if docker exec ollama ollama list 2>/dev/null | grep -v "NAME" | awk '{print $1}' | grep -v "^$"; then
        return 0
    else
        return 1
    fi
}

# List available DeepSeek models
list_deepseek_models() {
    echo "deepseek-r1:1.5b"
    echo "deepseek-r1:3b"
    echo "deepseek-r1:7b"
    echo "deepseek-r1:8b"
    echo "deepseek-r1:14b"
    echo "deepseek-r1:32b"
    echo "deepseek-r1:70b"
    echo "deepseek-r1"
    echo "deepseek-r1:distill-qwen-1.5b"
    echo "deepseek-r1:distill-qwen-7b"
    echo "deepseek-r1:distill-qwen-14b"
    echo "deepseek-r1:distill-qwen-32b"
    echo "deepseek-r1:distill-llama-8b"
    echo "deepseek-r1:distill-llama-70b"
}

# Install DeepSeek models
install_models() {
    print_info "Checking installed models..."
    
    # Get installed models
    installed_models=$(get_installed_models 2>/dev/null || echo "")
    
    if [[ -n "$installed_models" ]]; then
        print_info "Currently installed models:"
        echo "$installed_models" | while read -r model; do
            echo "  - $model"
        done
    else
        print_info "No models currently installed"
    fi
    
    # Get available DeepSeek models
    available_models=$(list_deepseek_models)
    
    # Filter out already installed models
    models_to_install=()
    while IFS= read -r model; do
        if ! echo "$installed_models" | grep -q "^$model$"; then
            models_to_install+=("$model")
        fi
    done <<< "$available_models"
    
    if [ ${#models_to_install[@]} -eq 0 ]; then
        print_info "All available DeepSeek models are already installed"
        return 0
    fi
    
    # Show available models for installation
    print_info "Available DeepSeek models to install:"
    echo
    for i in "${!models_to_install[@]}"; do
        printf "%2d. %s\n" $((i+1)) "${models_to_install[$i]}"
    done
    echo
    printf "%2d. Install ALL models (Warning: This will take a long time and lots of space)\n" $((${#models_to_install[@]}+1))
    printf "%2d. Skip model installation\n" $((${#models_to_install[@]}+2))
    echo
    
    # Ask user which model to install
    read -p "Enter the number of the model to install (or multiple numbers separated by spaces): " -r model_choices
    
    # Process user choice
    if [[ -z "$model_choices" ]]; then
        print_warning "No model selected, skipping..."
        return 0
    fi
    
    for choice in $model_choices; do
        if [[ "$choice" -eq $((${#models_to_install[@]}+2)) ]]; then
            print_info "Skipping model installation"
            return 0
        elif [[ "$choice" -eq $((${#models_to_install[@]}+1)) ]]; then
            print_warning "Installing ALL models. This will take a long time..."
            for model in "${models_to_install[@]}"; do
                print_info "Pulling model: $model"
                docker exec ollama ollama pull "$model"
                print_success "Model $model installed"
            done
            return 0
        elif [[ "$choice" -ge 1 ]] && [[ "$choice" -le ${#models_to_install[@]} ]]; then
            model_index=$((choice-1))
            model_name="${models_to_install[$model_index]}"
            print_info "Pulling model: $model_name"
            docker exec ollama ollama pull "$model_name"
            print_success "Model $model_name installed"
        else
            print_warning "Invalid choice: $choice, skipping..."
        fi
    done
}

# Check if Open WebUI is installed
check_webui() {
    if docker ps -a --format '{{.Names}}' | grep -q "^open-webui$"; then
        print_info "Open WebUI container exists"
        return 0
    else
        print_info "Open WebUI container does not exist"
        return 1
    fi
}

# Install Open WebUI
install_webui() {
    print_info "Installing Open WebUI..."
    
    # Create volume for Open WebUI data
    docker volume create open-webui || true
    
    # Run Open WebUI container
    docker run -d \
        --name open-webui \
        --restart unless-stopped \
        -p 3000:8080 \
        --add-host=host.docker.internal:host-gateway \
        -v open-webui:/app/backend/data \
        ghcr.io/open-webui/open-webui:main
    
    # Wait for WebUI to be ready
    print_info "Waiting for Open WebUI to start..."
    sleep 15
    
    # Check if WebUI is running
    for i in {1..30}; do
        if curl -s http://localhost:3000 > /dev/null 2>&1; then
            print_success "Open WebUI is running"
            return 0
        fi
        sleep 2
    done
    
    print_error "Open WebUI failed to start properly"
    return 1
}

# Configure Open WebUI to connect to Ollama
configure_webui_ollama() {
    print_info "Configuring Open WebUI to connect to Ollama..."
    
    # Get the host IP for Docker internal communication
    host_ip=$(ip -4 addr show docker0 | grep -oP '(?<=inet\s)\d+(\.\d+){3}')
    
    if [[ -z "$host_ip" ]]; then
        host_ip="172.17.0.1"
        print_warning "Could not detect Docker bridge IP, using default: $host_ip"
    fi
    
    # Restart Open WebUI with Ollama connection
    docker stop open-webui 2>/dev/null || true
    docker rm open-webui 2>/dev/null || true
    
    docker run -d \
        --name open-webui \
        --restart unless-stopped \
        -p 3000:8080 \
        --add-host=host.docker.internal:host-gateway \
        -v open-webui:/app/backend/data \
        -e OLLAMA_BASE_URL=http://host.docker.internal:11434 \
        ghcr.io/open-webui/open-webui:main
    
    print_success "Open WebUI configured to use Ollama"
}

# Get system IP address
get_system_ip() {
    # Try to get the main network interface IP
    ip_addr=$(ip -4 addr show | grep -oP '(?<=inet\s)\d+(\.\d+){3}' | grep -v '127.0.0.1' | head -n1)
    
    if [[ -z "$ip_addr" ]]; then
        ip_addr="localhost"
    fi
    
    echo "$ip_addr"
}

# Display final instructions
display_instructions() {
    echo
    print_success "========================================="
    print_success "Installation Complete!"
    print_success "========================================="
    echo
    
    ip_addr=$(get_system_ip)
    
    print_info "Access Open WebUI at:"
    echo -e "${GREEN}  http://$ip_addr:3000${NC}"
    echo -e "${GREEN}  http://localhost:3000${NC} (if accessing from the same machine)"
    echo
    
    if ! check_webui; then
        print_warning "Note: This is the first installation."
        print_info "You'll need to create an admin account when you first access the WebUI."
    fi
    
    print_info "Installed components:"
    
    if check_docker; then
        echo "  ✓ Docker"
    fi
    
    if check_ollama; then
        echo "  ✓ Ollama"
        
        # Show installed models
        installed_models=$(get_installed_models 2>/dev/null || echo "")
        if [[ -n "$installed_models" ]]; then
            echo "  ✓ Models:"
            echo "$installed_models" | while read -r model; do
                echo "    - $model"
            done
        fi
    fi
    
    if check_webui; then
        echo "  ✓ Open WebUI"
    fi
    
    echo
    print_info "To install additional models, run this script again with:"
    echo "  sudo $0"
    echo
}

# Main installation flow
main() {
    clear
    print_info "========================================="
    print_info "DeepSeek & Open WebUI Installation Script"
    print_info "========================================="
    echo
    
    # Check root privileges
    check_root
    
    # Update system
    update_system
    
    # Install Docker if needed
    if ! check_docker; then
        install_docker
    fi
    
    # Install Ollama if needed
    if ! check_ollama; then
        install_ollama
    else
        # Make sure Ollama is running
        if ! docker ps --format '{{.Names}}' | grep -q "^ollama$"; then
            print_info "Starting Ollama container..."
            docker start ollama
            sleep 5
        fi
    fi
    
    # Install/manage models
    install_models
    
    # Install Open WebUI if needed
    webui_first_install=false
    if ! check_webui; then
        webui_first_install=true
        install_webui
        configure_webui_ollama
    else
        # Make sure WebUI is running
        if ! docker ps --format '{{.Names}}' | grep -q "^open-webui$"; then
            print_info "Starting Open WebUI container..."
            docker start open-webui
            sleep 5
        fi
    fi
    
    # Display final instructions
    display_instructions
}

# Run main function
main "$@"
