#!/bin/bash
set -eo pipefail

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

# Utility functions
info() { echo -e "${BLUE}[INFO]${NC} $1"; }
success() { echo -e "${GREEN}[SUCCESS]${NC} $1"; }
error() { echo -e "${RED}[ERROR]${NC} $1" >&2; }
warning() { echo -e "${YELLOW}[WARNING]${NC} $1"; }

# Check if running on macOS
[[ "$OSTYPE" != "darwin"* ]] && { error "This script is for macOS only."; exit 1; }

# Request sudo access
echo "Please enter your sudo password to proceed with the installation:"
sudo -v

# Keep sudo credentials fresh
while true; do sudo -n true; sleep 60; kill -0 "$$" || exit; done 2>/dev/null &

# Ollama installation and configuration
install_ollama() {
    info "Installing Ollama..."
    GITHUB_API_URL="https://api.github.com/repos/ollama/ollama/releases/latest"
    GITHUB_RESPONSE=$(curl -s $GITHUB_API_URL)
    OLLAMA_DARWIN_ZIP_URL=$(echo "$GITHUB_RESPONSE" | grep -o '"browser_download_url": "[^"]*Ollama-darwin\.zip"' | sed 's/"browser_download_url": "\(.*\)"/\1/')
    OLLAMA_DARWIN_BIN_URL=$(echo "$GITHUB_RESPONSE" | grep -o '"browser_download_url": "[^"]*ollama-darwin"' | sed 's/"browser_download_url": "\(.*\)"/\1/')

    [[ -z "$OLLAMA_DARWIN_ZIP_URL" || -z "$OLLAMA_DARWIN_BIN_URL" ]] && { error "Failed to find the download URLs. Exiting..."; exit 1; }

    curl --fail --show-error --location --progress-bar -L "$OLLAMA_DARWIN_ZIP_URL" -o Ollama-darwin.zip
    unzip -q Ollama-darwin.zip
    sudo mv Ollama.app /Applications/
    rm Ollama-darwin.zip

    curl --fail --show-error --location --progress-bar -L "$OLLAMA_DARWIN_BIN_URL" -o ollama-darwin
    
    if [ -f "ollama-darwin" ]; then
        chmod +x ollama-darwin
        if [ ! -d "/usr/local/bin" ]; then
            sudo mkdir -p /usr/local/bin
        fi
        sudo mv ollama-darwin /usr/local/bin/ollama
        success "Ollama installation complete."
    else
        error "Failed to download ollama-darwin binary. Installation incomplete."
        exit 1
    fi
    chmod +x ollama-darwin
    sudo mv ollama-darwin /usr/local/bin/ollama

    success "Ollama installation complete."
}

configure_ollama() {
    export OLLAMA_NUM_PARALLEL=4
    export OLLAMA_MAX_LOADED_MODELS=2

    config_file="$HOME/.$(basename $SHELL)rc"
    echo "export OLLAMA_NUM_PARALLEL=4" >> "$config_file"
    echo "export OLLAMA_MAX_LOADED_MODELS=2" >> "$config_file"

    success "Ollama configuration complete."
}

start_ollama() {
    info "Starting Ollama service with custom settings..."
    nohup env OLLAMA_NUM_PARALLEL=4 OLLAMA_MAX_LOADED_MODELS=2 /usr/local/bin/ollama serve > /tmp/ollama.log 2>&1 &
    for i in {1..30}; do
        lsof -i:11434 >/dev/null 2>&1 && { success "Ollama is running on port 11434."; return 0; }
        sleep 1
        printf "."
    done
    error "Failed to start Ollama. Check /tmp/ollama.log for details."
    return 1
}

pull_models() {
    for model in "llama3" "llava" "mxbai-embed-large"; do
        if ! ollama list 2>/dev/null | grep -q "$model"; then
            info "Pulling $model..."
            ollama pull "$model"
        else
            warning "$model model is already available."
        fi
    done
}

# ReflowAI installation
install_reflowai() {
    REPO_URL="https://api.github.com/repos/Reflow-AI/reflow-macos/releases/latest"

    if [ -z "$GITHUB_TOKEN" ]; then
        error "GITHUB_TOKEN environment variable is not set."
        exit 1
    fi

    release_info=$(curl -s -H "Authorization: token $GITHUB_TOKEN" $REPO_URL)

    if command -v jq &> /dev/null; then
        download_url=$(echo "$release_info" | jq -r '.assets[0].url')
    else
        download_url=$(echo "$release_info" | grep -o '"url": "[^"]*"' | grep -o 'https://.*api.github.com/repos/Reflow-AI/reflow-macos/releases/assets/[0-9]*' | head -n 1)
    fi

    if [ -z "$download_url" ]; then
        error "Could not find download URL for the latest release."
        exit 1
    fi

    info "Downloading ReflowAI.dmg..."
    curl --fail --show-error --location --progress-bar -L -H "Authorization: token $GITHUB_TOKEN" -H "Accept: application/octet-stream" $download_url -o /tmp/ReflowAI.dmg
    success "ReflowAI.dmg downloaded successfully."

    info "Installing ReflowAI..."
    hdiutil attach /tmp/ReflowAI.dmg
    cp -R /Volumes/ReflowAI/ReflowAIApp.app /Applications/
    hdiutil detach /Volumes/ReflowAI
    rm /tmp/ReflowAI.dmg
    success "ReflowAI installed successfully."
}

# Main execution
main() {
    # Ollama setup
    if [ ! -d "/Applications/Ollama.app" ] || [ ! -f "/usr/local/bin/ollama" ]; then
        install_ollama
        configure_ollama
    else
        warning "Ollama is already installed."
    fi

    start_ollama || exit 1
    pull_models

    # ReflowAI setup
    if [ ! -d "/Applications/ReflowAIApp.app" ]; then
        install_reflowai
    else
        warning "ReflowAI is already installed."
    fi

    success "Setup complete! Ollama is now running with the specified settings."
    info "To start Ollama again with these settings, run:"
    echo "OLLAMA_NUM_PARALLEL=4 OLLAMA_MAX_LOADED_MODELS=2 /usr/local/bin/ollama serve"
}

main