#!/bin/bash

# RusTorch Universal Jupyter Installer
# 環境自動検出ワンライナーインストーラー

set -e

# Default installation path for launcher script
DEFAULT_INSTALL_PATH="$HOME/bin"
INSTALL_PATH="${RUSTORCH_INSTALL_PATH:-$DEFAULT_INSTALL_PATH}"

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

echo "🚀 RusTorch Universal Jupyter Installer"
echo "🚀 RusTorch万能Jupyterインストーラー"
echo ""

# Function to detect OS
detect_os() {
    case "$OSTYPE" in
        darwin*)  echo "macos" ;;
        linux*)   echo "linux" ;;
        msys*|mingw*|cygwin*) echo "windows" ;;
        *)        echo "unknown" ;;
    esac
}

# Function to detect CPU architecture
detect_cpu() {
    local arch=$(uname -m)
    case "$arch" in
        x86_64|amd64) echo "x64" ;;
        arm64|aarch64) echo "arm64" ;;
        armv7l) echo "arm32" ;;
        *) echo "unknown" ;;
    esac
}

# Function to detect GPU capabilities
detect_gpu() {
    local gpu_support=""
    
    # Check for NVIDIA GPU (CUDA)
    if command -v nvidia-smi >/dev/null 2>&1; then
        gpu_support="cuda"
    # Check for Apple Silicon (Metal)
    elif [[ "$(detect_os)" == "macos" ]] && [[ "$(detect_cpu)" == "arm64" ]]; then
        gpu_support="metal"
    # Check for AMD/Intel GPU (OpenCL) - basic check
    elif command -v clinfo >/dev/null 2>&1; then
        gpu_support="opencl"
    else
        gpu_support="cpu"
    fi
    
    echo "$gpu_support"
}

# Function to check browser compatibility for WebGPU
check_webgpu_support() {
    local os=$(detect_os)
    # WebGPU works best on Chrome/Chromium on all platforms
    if command -v google-chrome >/dev/null 2>&1 || 
       command -v chromium-browser >/dev/null 2>&1 || 
       command -v chromium >/dev/null 2>&1; then
        echo "true"
    else
        echo "false"
    fi
}

# Function to create launcher script
create_launcher() {
    local install_type="$1"
    local project_path="$2"
    local launcher_name="rustorch-jupyter"
    local launcher_path="$INSTALL_PATH/$launcher_name"
    
    # Create install directory if it doesn't exist
    mkdir -p "$INSTALL_PATH"
    
    cat > "$launcher_path" << EOF
#!/bin/bash
# RusTorch Jupyter Launcher - Generated by Universal Installer
# Installation type: $install_type
# Default project path: $project_path

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

# Default configuration
DEFAULT_RUSTORCH_PROJECT_DIR="$project_path"
LAUNCHER_SCRIPT="\$0"

# Function to show help
show_help() {
    echo "🚀 RusTorch Jupyter Launcher"
    echo "🚀 RusTorch Jupyterランチャー"
    echo ""
    echo "Usage / 使用方法:"
    echo "  \$0                          # Launch with default directory"
    echo "  \$0 --dir <path>             # Launch with custom directory"
    echo "  \$0 --uninstall              # Uninstall RusTorch Jupyter"
    echo "  \$0 --help                   # Show this help"
    echo ""
    echo "Examples / 例:"
    echo "  \$0"
    echo "  \$0 --dir /path/to/rustorch"
    echo "  \$0 --uninstall"
    echo ""
    echo "Default directory / デフォルトディレクトリ:"
    echo "  \$DEFAULT_RUSTORCH_PROJECT_DIR"
    echo ""
}

# Function to uninstall
uninstall_rustorch() {
    echo -e "\${RED}🗑️  RusTorch Jupyter Uninstaller\${NC}"
    echo -e "\${RED}🗑️  RusTorch Jupyterアンインストーラー\${NC}"
    echo ""
    
    echo "The following will be removed / 以下が削除されます:"
    echo "  • Launcher script: \$LAUNCHER_SCRIPT"
    echo "  • ランチャースクリプト: \$LAUNCHER_SCRIPT"
    if [[ -d "\$DEFAULT_RUSTORCH_PROJECT_DIR" ]]; then
        echo "  • Project directory: \$DEFAULT_RUSTORCH_PROJECT_DIR"
        echo "  • プロジェクトディレクトリ: \$DEFAULT_RUSTORCH_PROJECT_DIR"
    fi
    echo ""
    
    read -p "Are you sure? (y/N): " -n 1 -r
    echo
    if [[ \$REPLY =~ ^[Yy]\$ ]]; then
        echo "🗑️  Removing RusTorch installation..."
        echo "🗑️  RusTorchインストールを削除中..."
        
        # Remove project directory
        if [[ -d "\$DEFAULT_RUSTORCH_PROJECT_DIR" ]]; then
            rm -rf "\$DEFAULT_RUSTORCH_PROJECT_DIR"
            echo "✅ Removed project directory / プロジェクトディレクトリを削除しました"
        fi
        
        # Remove launcher script
        rm -f "\$LAUNCHER_SCRIPT"
        echo "✅ Removed launcher script / ランチャースクリプトを削除しました"
        
        echo ""
        echo -e "\${GREEN}🎉 RusTorch Jupyter successfully uninstalled!\${NC}"
        echo -e "\${GREEN}🎉 RusTorch Jupyterが正常にアンインストールされました！\${NC}"
        echo ""
        echo "💡 To reinstall, run:"
        echo "💡 再インストールするには:"
        echo "   curl -sSL https://raw.githubusercontent.com/JunSuzukiJapan/rustorch/main/install_jupyter.sh | bash"
        
        exit 0
    else
        echo "❌ Uninstall cancelled / アンインストールがキャンセルされました"
        exit 0
    fi
}

# Function to verify and launch RusTorch
launch_rustorch() {
    local project_dir="\$1"
    
    echo "🚀 RusTorch Jupyter Launcher"
    echo "🚀 RusTorch Jupyterランチャー"
    echo ""
    echo -e "📍 Project directory: \${BLUE}\$project_dir\${NC}"
    echo -e "📍 プロジェクトディレクトリ: \${BLUE}\$project_dir\${NC}"
    
    # Check if project directory exists
    if [[ ! -d "\$project_dir" ]]; then
        echo -e "\${RED}❌ RusTorch project directory not found: \$project_dir\${NC}"
        echo -e "\${RED}❌ RusTorchプロジェクトディレクトリが見つかりません: \$project_dir\${NC}"
        echo ""
        echo "🔄 Please check the path or reinstall using:"
        echo "🔄 パスを確認するか、再インストールしてください:"
        echo "   curl -sSL https://raw.githubusercontent.com/JunSuzukiJapan/rustorch/main/install_jupyter.sh | bash"
        exit 1
    fi
    
    # Store current directory to return to it later
    ORIGINAL_DIR="\$(pwd)"
    
    # Navigate to project directory
    cd "\$project_dir" || {
        echo -e "\${RED}❌ Failed to navigate to \$project_dir\${NC}"
        echo -e "\${RED}❌ \$project_dir への移動に失敗しました\${NC}"
        exit 1
    }
    
    # Verify this is a RusTorch project
    if [[ ! -f "Cargo.toml" ]] || ! grep -q "rustorch" "Cargo.toml" 2>/dev/null; then
        echo -e "\${RED}❌ Invalid RusTorch project directory: \$project_dir\${NC}"
        echo -e "\${RED}❌ 無効なRusTorchプロジェクトディレクトリ: \$project_dir\${NC}"
        echo ""
        echo "🔄 Please check the path or reinstall using:"
        echo "🔄 パスを確認するか、再インストールしてください:"
        echo "   curl -sSL https://raw.githubusercontent.com/JunSuzukiJapan/rustorch/main/install_jupyter.sh | bash"
        exit 1
    fi
    
    echo -e "\${GREEN}✅ RusTorch project verified\${NC}"
    echo -e "\${GREEN}✅ RusTorchプロジェクトを確認しました\${NC}"
    echo ""
    
    # Use the existing quick launcher
    if [[ -f "./scripts/start_jupyter_quick.sh" ]]; then
        echo "🚀 Starting Jupyter environment..."
        echo "🚀 Jupyter環境を開始中..."
        ./scripts/start_jupyter_quick.sh
    elif [[ -f "./start_jupyter_quick.sh" ]]; then
        echo "🚀 Starting Jupyter environment..."
        echo "🚀 Jupyter環境を開始中..."
        ./start_jupyter_quick.sh
    else
        echo -e "\${RED}❌ start_jupyter_quick.sh not found in project directory\${NC}"
        echo -e "\${RED}❌ プロジェクトディレクトリにstart_jupyter_quick.shが見つかりません\${NC}"
        echo ""
        echo "🔄 Please update your RusTorch installation:"
        echo "🔄 RusTorchインストールを更新してください:"
        echo "   cd \$project_dir && git pull"
        exit 1
    fi
    
    # Return to original directory if Jupyter exits
    cd "\$ORIGINAL_DIR"
}

# Parse command line arguments
while [[ \$# -gt 0 ]]; do
    case \$1 in
        --dir|-d)
            CUSTOM_DIR="\$2"
            shift 2
            ;;
        --uninstall)
            uninstall_rustorch
            ;;
        --help|-h)
            show_help
            exit 0
            ;;
        *)
            echo "Unknown option: \$1"
            echo "Use --help for usage information"
            exit 1
            ;;
    esac
done

# Determine which directory to use
if [[ -n "\$CUSTOM_DIR" ]]; then
    # Use custom directory if provided
    RUSTORCH_PROJECT_DIR="\$CUSTOM_DIR"
    echo "Using custom directory: \$RUSTORCH_PROJECT_DIR"
    echo "カスタムディレクトリを使用: \$RUSTORCH_PROJECT_DIR"
    echo ""
else
    # Use default directory
    RUSTORCH_PROJECT_DIR="\$DEFAULT_RUSTORCH_PROJECT_DIR"
fi

# Launch RusTorch
launch_rustorch "\$RUSTORCH_PROJECT_DIR"
EOF
    
    chmod +x "$launcher_path"
    echo "$launcher_path"
}

# Function to add to PATH
add_to_path() {
    local install_dir="$1"
    local shell_rc=""
    
    # Detect shell and corresponding RC file
    if [[ "$SHELL" == *"zsh"* ]]; then
        shell_rc="$HOME/.zshrc"
    elif [[ "$SHELL" == *"bash"* ]]; then
        if [[ -f "$HOME/.bashrc" ]]; then
            shell_rc="$HOME/.bashrc"
        else
            shell_rc="$HOME/.bash_profile"
        fi
    elif [[ "$SHELL" == *"fish"* ]]; then
        shell_rc="$HOME/.config/fish/config.fish"
    fi
    
    # Check if PATH already contains the directory
    if [[ ":$PATH:" != *":$install_dir:"* ]]; then
        if [[ -n "$shell_rc" ]]; then
            echo "" >> "$shell_rc"
            echo "# Added by RusTorch installer" >> "$shell_rc"
            if [[ "$shell_rc" == *"fish"* ]]; then
                echo "set -gx PATH $install_dir \$PATH" >> "$shell_rc"
            else
                echo "export PATH=\"$install_dir:\$PATH\"" >> "$shell_rc"
            fi
            echo "✅ Added $install_dir to PATH in $shell_rc"
        fi
    fi
}

# Main detection and installation logic
main() {
    echo "🔍 Environment Detection / 環境検出"
    echo "=================================="
    
    local os=$(detect_os)
    local cpu=$(detect_cpu)
    local gpu=$(detect_gpu)
    local webgpu_support=$(check_webgpu_support)
    
    echo "OS: $os"
    echo "CPU: $cpu" 
    echo "GPU: $gpu"
    echo "WebGPU Support: $webgpu_support"
    echo ""
    
    # Determine best installation strategy
    local install_strategy=""
    local install_command=""
    
    # Store GPU-specific info for fallback
    local gpu_specific_strategy=""
    local gpu_specific_command=""
    
    if [[ "$gpu" == "cuda" ]]; then
        gpu_specific_strategy="CUDA GPU acceleration"
        gpu_specific_command="./scripts/start_jupyter.sh"
    elif [[ "$gpu" == "metal" ]]; then
        gpu_specific_strategy="Metal GPU acceleration"
        gpu_specific_command="./scripts/start_jupyter.sh"
    elif [[ "$webgpu_support" == "true" ]]; then
        gpu_specific_strategy="WebGPU browser acceleration"
        gpu_specific_command="./scripts/start_jupyter_webgpu.sh"
    else
        gpu_specific_strategy="CPU-optimized with Rust kernel"
        gpu_specific_command="./scripts/quick_start_rust_kernel.sh"
    fi

    # Default to hybrid
    install_strategy="Hybrid Python + Rust dual-kernel"
    install_command="./scripts/start_jupyter_hybrid.sh"
    
    echo -e "${GREEN}🎯 Default (Recommended): Hybrid Environment${NC}"
    echo -e "${GREEN}🎯 デフォルト（推奨）: ハイブリッド環境${NC}"
    echo "     🦀🐍 Dual-kernel Jupyter with both Python and Rust support"
    echo "     🦀🐍 Python環境とRust環境の両方をサポートするデュアルカーネル"
    echo ""
    echo -e "${BLUE}Alternative: $gpu_specific_strategy${NC}"
    echo "     GPU-optimized single environment"
    echo "     GPU最適化された単一環境"
    echo ""
    
    # Ask user for setup type (auto-select default for non-interactive mode)
    echo "Choose installation type:"
    echo "[1] Hybrid (Default - Python + Rust dual-kernel) ⭐"
    echo "[2] GPU-Optimized ($gpu_specific_strategy)"
    echo "[q] Cancel"
    echo ""
    
    # Check if running interactively
    if [[ -t 0 ]]; then
        read -p "Select [1/2/q] (default: 1): " -n 1 -r
        echo
    else
        echo "Non-interactive mode detected. Auto-selecting default option [1]..."
        REPLY="1"
    fi
    
    case $REPLY in
        1|"")
            echo "Using hybrid setup (default)..."
            ;;
        2)
            install_strategy="$gpu_specific_strategy"
            install_command="$gpu_specific_command"
            echo "Using GPU-optimized setup..."
            ;;
        q|Q)
            echo "Installation cancelled."
            exit 0
            ;;
        *)
            echo "Invalid selection. Using hybrid setup (default)..."
            ;;
    esac
    
    echo ""
    echo "🚀 Starting Installation / インストール開始"
    echo "==========================================="
    
    # Set up working directory and track project path
    local project_dir=""
    local rustorch_home="$HOME/.rustorch"
    mkdir -p "$rustorch_home"
    
    # Primary approach: Clone the complete repository
    if command -v git >/dev/null 2>&1; then
        echo "📦 Setting up complete RusTorch project..."
        
        if [[ -d "$rustorch_home/rustorch" ]]; then
            # Check if it's a valid git repository
            if [[ -d "$rustorch_home/rustorch/.git" ]]; then
                echo "🔄 Updating existing RusTorch installation..."
                cd "$rustorch_home/rustorch"
                git pull
            else
                echo "🔄 Existing directory is not a git repository. Recreating..."
                rm -rf "$rustorch_home/rustorch"
                cd "$rustorch_home"
                git clone https://github.com/JunSuzukiJapan/rustorch.git
                cd rustorch
            fi
        else
            echo "📥 Cloning RusTorch repository..."
            cd "$rustorch_home"
            git clone https://github.com/JunSuzukiJapan/rustorch.git
            cd rustorch
        fi
        
        chmod +x *.sh
        eval "$install_command"
        project_dir="$rustorch_home/rustorch"
        cd - > /dev/null
    else
        echo "⚠️  Git not found. Using fallback approach..."
        echo "🔄 Attempting script download method..."
        
        # Fallback: Download essential files individually
        local script_name=$(basename "$install_command")
        echo "📥 Downloading $script_name..."
        
        if curl -sSL "https://raw.githubusercontent.com/JunSuzukiJapan/rustorch/main/$script_name" -o "$script_name"; then
            chmod +x "$script_name"
            echo "✅ Downloaded $script_name successfully"
            
            # Create installation directory and download essential files
            if [[ -d "$rustorch_home/rustorch" ]]; then
                echo "🔄 Updating existing installation..."
                cd "$rustorch_home/rustorch"
            else
                echo "📦 Setting up RusTorch installation in $rustorch_home..."
                mkdir -p "$rustorch_home/rustorch"
                cd "$rustorch_home/rustorch"
                
                # Download essential project files
                echo "📥 Downloading essential project files..."
                curl -sSL "https://raw.githubusercontent.com/JunSuzukiJapan/rustorch/main/Cargo.toml" -o "Cargo.toml" || echo "Warning: Could not download Cargo.toml"
                
                # Download all necessary scripts
                mkdir -p scripts
                for script in "start_jupyter_quick.sh" "start_jupyter_hybrid.sh" "download_notebooks.sh"; do
                    curl -sSL "https://raw.githubusercontent.com/JunSuzukiJapan/rustorch/main/scripts/$script" -o "scripts/$script" || echo "Warning: Could not download $script"
                    chmod +x "scripts/$script" 2>/dev/null || true
                done
            fi
            
            # Copy and run installer script
            if [[ ! -f "./$script_name" ]]; then
                cp "$OLDPWD/$script_name" . 2>/dev/null || echo "Note: Script already available"
            elif ! cmp -s "$OLDPWD/$script_name" "./$script_name" 2>/dev/null; then
                cp "$OLDPWD/$script_name" . 2>/dev/null || echo "Note: Using existing script"
            else
                echo "✅ Script already exists and is up to date"
            fi
            
            eval "./$script_name"
            project_dir="$rustorch_home/rustorch"
            cd - > /dev/null
            
            # Clean up
            if [[ -f "$script_name" ]] && [[ "$PWD/$script_name" != "$rustorch_home/rustorch/$script_name" ]]; then
                rm -f "$script_name"
            fi
        else
            echo "❌ Failed to download $script_name"
            echo "❌ Both git and curl approaches failed. Please install git or check your internet connection."
            exit 1
        fi
    fi
    
    echo ""
    echo "📦 Creating launcher script / ランチャースクリプト作成"
    echo "=================================================="
    
    # Create launcher script with project path
    local launcher_path=$(create_launcher "$install_strategy" "$project_dir")
    echo -e "${GREEN}✅ Launcher created: $launcher_path${NC}"
    
    # Add to PATH
    add_to_path "$INSTALL_PATH"
    
    echo ""
    echo -e "${GREEN}🎉 Installation Complete! / インストール完了！${NC}"
    echo "=============================================="
    echo ""
    echo "📋 Usage / 使用方法:"
    echo "  1. Restart your terminal / ターミナルを再起動"
    echo "  2. Run anywhere: ${BLUE}rustorch-jupyter${NC}"
    echo "  3. Or directly: ${BLUE}$launcher_path${NC}"
    echo ""
    echo "🛠️  Advanced Options / 高度なオプション:"
    echo "  ${BLUE}rustorch-jupyter --help${NC}              # Show help / ヘルプを表示"
    echo "  ${BLUE}rustorch-jupyter --dir <path>${NC}        # Use custom directory / カスタムディレクトリを使用"
    echo "  ${BLUE}rustorch-jupyter --uninstall${NC}         # Uninstall / アンインストール"
    echo ""
    echo "📍 Installation details / インストール詳細:"
    echo "  • Project location: ${BLUE}$project_dir${NC}"
    echo "  • プロジェクトの場所: ${BLUE}$project_dir${NC}"
    echo "  • Launcher script: ${BLUE}$launcher_path${NC}"
    echo "  • ランチャースクリプト: ${BLUE}$launcher_path${NC}"
    echo ""
    echo "🔧 Customization / カスタマイズ:"
    echo "  Install path: ${BLUE}RUSTORCH_INSTALL_PATH=/custom/path $0${NC}"
    echo ""
    echo "💡 Note: You can now run ${BLUE}rustorch-jupyter${NC} from any directory!"
    echo "💡 注意: ${BLUE}rustorch-jupyter${NC}はどのディレクトリからでも実行できます！"
    echo ""
    echo "🗑️  To uninstall later / 後でアンインストールするには:"
    echo "   ${BLUE}rustorch-jupyter --uninstall${NC}"
    echo ""
}

# Handle command line arguments
if [[ "$1" == "--help" ]] || [[ "$1" == "-h" ]]; then
    echo "RusTorch Universal Jupyter Installer"
    echo ""
    echo "Usage:"
    echo "  $0                    # Install with auto-detection"
    echo "  $0 --help            # Show this help"
    echo ""
    echo "Environment Variables:"
    echo "  RUSTORCH_INSTALL_PATH # Custom install path (default: ~/bin)"
    echo ""
    exit 0
fi

# Run main installation
main