#!/bin/bash

# ========== Ollama Workbench Setup Script for Apple Silicon ==========
# This script automates the installation of Ollama Workbench with
# optimized dependencies for Apple Silicon (M1/M2/M3/M4) Macs.
# ===================================================================

# Terminal colors for better user experience
CYAN='\033[0;36m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m' # No Color

# Print styled text
print_step() {
    echo -e "${CYAN}==>${NC} $1"
}

print_success() {
    echo -e "${GREEN}✓${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}!${NC} $1"
}

print_error() {
    echo -e "${RED}✗${NC} $1"
}

# Function to check if a command exists
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Function to check Apple Silicon
is_apple_silicon() {
    if [[ "$(uname -m)" == "arm64" ]]; then
        return 0
    else
        return 1
    fi
}

# Function to create a placeholder groq_utils.py file if needed
create_placeholder_groq_utils() {
    if [ ! -f "groq_utils.py" ]; then
        cat > groq_utils.py << EOL
# Placeholder for groq_utils.py
# This file was auto-generated by the installer

GROQ_MODELS = []

def get_local_embeddings(*args, **kwargs):
    """
    Placeholder function for get_local_embeddings.
    """
    return None
EOL
        print_success "Created placeholder groq_utils.py file"
    fi
}

# Main installation process
echo -e "${CYAN}=======================================${NC}"
echo -e "${CYAN}   Ollama Workbench Setup for Apple Silicon   ${NC}"
echo -e "${CYAN}=======================================${NC}"

# Check if running on Apple Silicon
if is_apple_silicon; then
    print_success "Detected Apple Silicon Mac"
else
    print_warning "This script is optimized for Apple Silicon Macs. Some optimizations may not apply."
fi

# Check for Python 3.11+
if ! command_exists python3 || [ "$(python3 -c 'import sys; print(int(sys.version_info.major == 3 and sys.version_info.minor >= 11))')" = "0" ]; then
    print_error "Python 3.11 or higher is required"
    echo "Please install Python 3.11+ from https://www.python.org/downloads/"
    exit 1
else
    print_success "Python 3.11+ is installed"
fi

# Check for Poetry and install if needed
if ! command_exists poetry; then
    print_step "Poetry not found, installing..."
    curl -sSL https://install.python-poetry.org | python3 -
    
    # After installation, check if Poetry was installed to ~/.local/bin
    POETRY_PATH="$HOME/.local/bin/poetry"
    if [ -f "$POETRY_PATH" ]; then
        print_success "Poetry installed successfully to $POETRY_PATH"
        print_step "Adding Poetry to PATH for this session..."
        export PATH="$HOME/.local/bin:$PATH"
        
        # Remind user to update their PATH
        print_warning "To use Poetry in future terminal sessions, add this line to your shell profile (~/.zshrc or ~/.bash_profile):"
        echo "    export PATH=\"\$HOME/.local/bin:\$PATH\""
    else
        # Try to find Poetry in other common locations
        for dir in "/usr/local/bin" "/opt/homebrew/bin" "$HOME/Library/Python/*/bin"; do
            for path in $dir/poetry; do
                if [ -f "$path" ]; then
                    POETRY_PATH="$path"
                    print_success "Found Poetry at $POETRY_PATH"
                    break 2
                fi
            done
        done
        
        if [ -z "$POETRY_PATH" ]; then
            print_error "Poetry was installed but couldn't be found in the PATH"
            print_warning "You may need to restart your terminal or manually add Poetry to your PATH"
            print_warning "Try running: which poetry"
            print_warning "If that doesn't work, close and reopen your terminal, then run this script again"
            exit 1
        fi
    fi
    
    # Use the full path to Poetry
    POETRY_CMD="$POETRY_PATH"
else
    print_success "Poetry is already installed"
    POETRY_CMD="poetry"
fi

# Configure Poetry
print_step "Configuring Poetry..."
$POETRY_CMD config virtualenvs.in-project true
print_success "Poetry configured"

# Remove existing conda environment if it exists
if command_exists conda; then
    print_warning "Found conda installation"
    if conda env list | grep -q "ollamaworkbench"; then
        print_step "Removing existing conda environment..."
        conda remove --name ollamaworkbench --all -y
        print_success "Removed conda environment"
    fi
fi

# Check for existing Poetry environment
if [ -d ".venv" ]; then
    print_warning "Existing Poetry virtual environment found"
    read -p "Do you want to remove it and start fresh? (y/N): " remove_env
    if [[ $remove_env =~ ^[Yy]$ ]]; then
        print_step "Removing existing Poetry environment..."
        rm -rf .venv
        print_success "Removed existing virtual environment"
    fi
fi

# Install dependencies
print_step "Installing dependencies using Poetry..."
# First, make sure pyproject.toml exists
if [ ! -f "pyproject.toml" ]; then
    print_error "pyproject.toml not found. Please ensure you're in the correct directory."
    exit 1
fi

# Install dependencies from pyproject.toml
# Use --no-root to avoid trying to install the current project as a package
echo "Debug: Running Poetry install with --no-root flag"
$POETRY_CMD install --no-root

if [ $? -ne 0 ]; then
    print_error "Poetry install failed. Attempting to fix common issues..."
    
    # Create placeholder groq_utils.py if it doesn't exist
    create_placeholder_groq_utils
    
    # Try installing dependencies in the correct order to avoid conflicts
    print_step "Installing base dependencies..."
    $POETRY_CMD run pip install --upgrade pip
    
    # Install NumPy 1.x first to ensure compatibility
    print_step "Installing NumPy 1.x..."
    $POETRY_CMD run pip install "numpy>=1.24.3,<2.0.0"
    
    # If we're on Apple Silicon, install torch with MPS support
    if is_apple_silicon; then
        print_step "Installing PyTorch with Apple Silicon (MPS) support..."
        $POETRY_CMD run pip install torch==2.2.0 torchvision==0.17.0 --index-url https://download.pytorch.org/whl/cpu
    else
        print_step "Installing PyTorch..."
        $POETRY_CMD run pip install torch torchvision
    fi
    
    # Install core dependencies in order
    print_step "Installing core dependencies..."
    # Add debug output
    echo "Debug: Installing pandas with proper quoting"
    $POETRY_CMD run pip install "pandas>=1.5.0,<2.0.0"
    $POETRY_CMD run pip install "tiktoken>=0.7.0,<1.0.0"
    $POETRY_CMD run pip install transformers==4.38.0
    $POETRY_CMD run pip install sentence-transformers==2.5.0
    $POETRY_CMD run pip install streamlit==1.32.0
    
    # Install langchain components separately to avoid conflicts
    echo "Debug: Installing langchain components separately"
    $POETRY_CMD run pip install "langchain-core>=0.2.15,<0.3.0"
    $POETRY_CMD run pip install "langchain-community==0.2.15"
    $POETRY_CMD run pip install "langchain==0.2.15"
    # Install missing dependencies from requirements.txt
    echo "Debug: Installing missing dependencies"
    $POETRY_CMD run pip install "ollama==0.2.0" "chromadb==0.5.5" "psutil==6.0.0" "pdfkit==1.0.0" "matplotlib==3.9.2" "beautifulsoup4==4.12.3" "bs4==0.0.2" "selenium==4.24.0" "webdriver-manager==4.0.2" "PyPDF2==3.0.1" "streamlit-extras==0.3.6" "autogen==0.2.35" "pyautogen==0.2.35" "fpdf==1.7.2" "radon==6.0.1" "flake8==7.1.1"
    
    # Install additional dependencies from requirements.txt that might be missing
    echo "Debug: Installing additional dependencies from requirements.txt"
    $POETRY_CMD run pip install -r requirements.txt
    
    # Try poetry install again
    print_step "Resuming Poetry installation..."
    $POETRY_CMD install --no-interaction --no-root
fi

# Create placeholder groq_utils.py if it doesn't exist (in case we didn't hit the error path)
create_placeholder_groq_utils

# Run validation tests
print_step "Running validation tests..."
cat > validation_test.py << EOL
import sys
import os
import warnings
warnings.filterwarnings("ignore", category=UserWarning)
print("Python version:", sys.version)

# Check NumPy version first
try:
    import numpy
    print("NumPy version:", numpy.__version__)
    numpy_major = int(numpy.__version__.split('.')[0])
    if numpy_major >= 2:
        print("WARNING: NumPy 2.x detected. Some packages may not be compatible.")
        print("Consider downgrading to NumPy 1.x if you encounter issues.")
except ImportError:
    print("NumPy not found")

# Test PyTorch installation
try:
    import torch
    print("PyTorch version:", torch.__version__)
    if torch.cuda.is_available():
        print("CUDA is available")
    else:
        print("CUDA is not available")
    
    # Check MPS in multiple ways to handle different PyTorch versions
    mps_available = False
    
    # Try the standard way (PyTorch 2.0+)
    if hasattr(torch, 'backends') and hasattr(torch.backends, 'mps'):
        try:
            if torch.backends.mps.is_available():
                mps_available = True
        except AttributeError:
            pass
    
    # Try the older way or direct check
    if not mps_available and hasattr(torch, 'mps'):
        try:
            # Some versions use torch.mps.is_available()
            if hasattr(torch.mps, 'is_available'):
                if torch.mps.is_available():
                    mps_available = True
            # Some versions just check if the module exists
            else:
                mps_available = True
        except (AttributeError, ImportError):
            pass
            
    # Try creating a tensor on MPS as final check
    if not mps_available:
        try:
            device = torch.device("mps")
            x = torch.ones(1, device=device)
            mps_available = True
        except (RuntimeError, ValueError):
            pass
    
    if mps_available:
        print("MPS (Metal Performance Shaders) is available - Apple Silicon optimization working!")
    else:
        print("MPS is not available")
except ImportError as e:
    print("Failed to import torch:", e)

# Test torchvision installation
try:
    import torchvision
    print("Torchvision version:", torchvision.__version__)
except ImportError as e:
    print("Failed to import torchvision:", e)

# Test sentence-transformers installation
try:
    import sentence_transformers
    print("Sentence-transformers version:", sentence_transformers.__version__)
except ImportError as e:
    print("Failed to import sentence_transformers:", e)

# Test other key dependencies
dependencies = [
    "streamlit", "langchain", "transformers", "ollama",
    "chromadb", "pandas", "psutil", "requests", "matplotlib",
    "beautifulsoup4", "bs4", "selenium", "webdriver_manager", "PyPDF2",
    "streamlit_extras", "autogen", "pyautogen", "fpdf", "radon", "flake8"
]

# Add detailed error checking for critical dependencies
print("\n=== Critical Dependencies Check ===")
critical_deps = ["ollama", "chromadb", "psutil", "pdfkit", "matplotlib", "bs4", "selenium", "webdriver_manager", "PyPDF2", "streamlit_extras", "autogen", "pyautogen", "fpdf", "radon", "flake8"]
for dep in critical_deps:
    try:
        module = __import__(dep)
        version = getattr(module, "__version__", "unknown")
        print(f"✓ {dep} version: {version} - Successfully imported")
    except ImportError as e:
        print(f"✗ {dep} MISSING: {e}")
        print(f"  Attempting to install {dep}...")
        import subprocess
        try:
            subprocess.check_call([sys.executable, "-m", "pip", "install", dep])
            print(f"  Installed {dep}, retrying import...")
            try:
                module = __import__(dep)
                version = getattr(module, "__version__", "unknown")
                print(f"✓ {dep} version: {version} - Successfully imported after installation")
            except ImportError as e2:
                print(f"✗ {dep} still MISSING after installation attempt: {e2}")
        except subprocess.CalledProcessError:
            print(f"✗ Failed to install {dep}")

for dep in dependencies:
    try:
        module = __import__(dep)
        version = getattr(module, "__version__", "unknown")
        print(f"{dep} version: {version}")
    except ImportError as e:
        print(f"Failed to import {dep}:", e)

# Test pdfkit and wkhtmltopdf
try:
    import pdfkit
    print("pdfkit version:", getattr(pdfkit, "__version__", "unknown"))
    
    # Check if wkhtmltopdf is available
    try:
        path = pdfkit.configuration().wkhtmltopdf
        if path and os.path.exists(path):
            print(f"✓ wkhtmltopdf found at: {path}")
        else:
            print(f"✗ wkhtmltopdf not found at configured path: {path}")
            print("  PDF generation features will not work properly.")
            print("  Please install wkhtmltopdf: https://wkhtmltopdf.org/downloads.html")
    except Exception as e:
        print(f"✗ Error checking wkhtmltopdf: {e}")
        print("  PDF generation features will not work properly.")
        print("  Please install wkhtmltopdf: https://wkhtmltopdf.org/downloads.html")
except ImportError as e:
    print(f"Failed to import pdfkit: {e}")
    print("  PDF generation features will not work properly.")
    print("  Please ensure pdfkit is installed: pip install pdfkit")
    print("  And install wkhtmltopdf: https://wkhtmltopdf.org/downloads.html")

# Test import of application modules
try:
    import ollama_utils
    print("Successfully imported ollama_utils")
except ImportError as e:
    print("Failed to import ollama_utils:", e)
    print("\n=== Detailed ollama_utils import error analysis ===")
    # Try to identify which specific dependency is causing the issue
    try:
        import ollama
        print("✓ ollama module is available")
    except ImportError as e1:
        print(f"✗ ollama module is missing: {e1}")
    
    try:
        import psutil
        print("✓ psutil module is available")
    except ImportError as e2:
        print(f"✗ psutil module is missing: {e2}")
    
    try:
        import streamlit
        print("✓ streamlit module is available")
    except ImportError as e3:
        print(f"✗ streamlit module is missing: {e3}")
        
    # Check if the file exists but can't be imported
    import os
    if os.path.exists("ollama_utils.py"):
        print("✓ ollama_utils.py file exists")
        print("  The issue might be with the imports inside ollama_utils.py")
    else:
        print("✗ ollama_utils.py file does not exist in the current directory")
EOL

$POETRY_CMD run python validation_test.py
rm validation_test.py

# Install and download spaCy model with NumPy compatibility
print_step "Installing spaCy and downloading model..."
$POETRY_CMD run pip install "numpy>=1.24.3,<2.0.0" # Ensure NumPy 1.x is installed
$POETRY_CMD run pip install spacy==3.7.2 # Use a version known to work with NumPy 1.x
$POETRY_CMD run python -m spacy download en_core_web_sm

# Check for wkhtmltopdf (required by pdfkit)
if ! command_exists wkhtmltopdf; then
    print_warning "wkhtmltopdf is not installed (required by pdfkit)"
    echo "The pdfkit Python package requires wkhtmltopdf to be installed on your system."
    
    if [[ "$(uname)" == "Darwin" ]]; then
        print_step "Installing wkhtmltopdf for macOS..."
        # Create a temporary directory for the download
        TEMP_DIR=$(mktemp -d)
        WKHTML_PKG="$TEMP_DIR/wkhtmltopdf.pkg"
        
        # Download wkhtmltopdf from the official GitHub repository
        echo "Downloading wkhtmltopdf from GitHub..."
        curl -L https://github.com/wkhtmltopdf/packaging/releases/download/0.12.6-2/wkhtmltox-0.12.6-2.macos-cocoa.pkg -o "$WKHTML_PKG"
        
        if [ $? -eq 0 ] && [ -f "$WKHTML_PKG" ]; then
            echo "Installing wkhtmltopdf (you may be prompted for your password)..."
            sudo installer -pkg "$WKHTML_PKG" -target /
            
            if [ $? -eq 0 ]; then
                print_success "wkhtmltopdf installed successfully"
                # Clean up
                rm -rf "$TEMP_DIR"
            else
                print_error "Failed to install wkhtmltopdf"
                print_warning "Please install wkhtmltopdf manually from: https://wkhtmltopdf.org/downloads.html"
                print_warning "Some features that require PDF generation will not work until wkhtmltopdf is installed."
            fi
        else
            print_error "Failed to download wkhtmltopdf"
            print_warning "Please install wkhtmltopdf manually from: https://wkhtmltopdf.org/downloads.html"
            print_warning "Some features that require PDF generation will not work until wkhtmltopdf is installed."
        fi
    elif [[ "$(uname)" == "Linux" ]]; then
        print_step "Installing wkhtmltopdf for Linux..."
        echo "You can install it on Linux with: sudo apt-get install wkhtmltopdf"
        echo "Would you like to install wkhtmltopdf now? (y/N)"
        read -p "" install_wkhtml
        if [[ $install_wkhtml =~ ^[Yy]$ ]]; then
            sudo apt-get update && sudo apt-get install -y wkhtmltopdf
            if [ $? -eq 0 ]; then
                print_success "wkhtmltopdf installed successfully"
            else
                print_error "Failed to install wkhtmltopdf"
                print_warning "Please install wkhtmltopdf manually from: https://wkhtmltopdf.org/downloads.html"
                print_warning "Some features that require PDF generation will not work until wkhtmltopdf is installed."
            fi
        else
            print_warning "wkhtmltopdf installation skipped"
            print_warning "Some features that require PDF generation will not work until wkhtmltopdf is installed."
        fi
    else
        print_warning "Automatic installation of wkhtmltopdf is not supported on this platform."
        print_warning "Please install wkhtmltopdf manually from: https://wkhtmltopdf.org/downloads.html"
        print_warning "Some features that require PDF generation will not work until wkhtmltopdf is installed."
    fi
else
    print_success "wkhtmltopdf is installed"
fi

# Check if Ollama is installed
if ! command_exists ollama; then
    print_warning "Ollama is not installed"
    echo "It's recommended to install Ollama to use all features of Ollama Workbench."
    echo "You can install Ollama by running: curl -fsSL https://ollama.com/install.sh | sh"
else
    print_success "Ollama is installed"
fi

print_success "Installation complete!"
echo -e "${GREEN}=======================================${NC}"
echo -e "${GREEN}   Ollama Workbench is ready to use!   ${NC}"
echo -e "${GREEN}=======================================${NC}"
echo ""
echo -e "To run Ollama Workbench, use: ${CYAN}bash run_ollama_workbench.sh${NC}"
echo ""