#!/bin/bash
# TuFT Installation Script
# Usage: /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/agentscope-ai/tuft/main/scripts/install.sh)"

set -euo pipefail

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

# Configuration
TUFT_HOME="${TUFT_HOME:-$HOME/.tuft}"
TUFT_BIN="$TUFT_HOME/bin"
TUFT_VENV="$TUFT_HOME/venv"
PYTHON_VERSION="3.12"
TUFT_PYPI_PACKAGE="tuft"
TUFT_GIT_REPO="https://github.com/agentscope-ai/tuft.git"
INSTALL_FROM_SOURCE=false
LOCAL_SOURCE_PATH=""
CLEAN_INSTALL=false

# Print functions
print_step() {
    echo -e "${BLUE}==>${NC} $1"
}

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

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

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

# Parse command line arguments
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --from-source)
                INSTALL_FROM_SOURCE=true
                shift
                ;;
            --local-source)
                LOCAL_SOURCE_PATH="$2"
                shift 2
                ;;
            --clean)
                CLEAN_INSTALL=true
                shift
                ;;
            --help|-h)
                echo "TuFT Installation Script"
                echo ""
                echo "Usage: install.sh [options]"
                echo ""
                echo "Options:"
                echo "  --from-source         Install from GitHub instead of PyPI"
                echo "  --local-source PATH   Install from local source directory (for development/CI)"
                echo "  --clean               Remove existing installation before installing"
                echo "  --help, -h            Show this help message"
                echo ""
                echo "The script installs TuFT with full backend support (GPU, persistence, flash-attn)."
                echo ""
                echo "Environment Variables:"
                echo "  TUFT_HOME             Installation directory (default: ~/.tuft)"
                exit 0
                ;;
            *)
                print_error "Unknown option: $1"
                exit 1
                ;;
        esac
    done
}

# Detect OS and architecture
detect_platform() {
    OS="$(uname -s)"
    ARCH="$(uname -m)"

    case "$OS" in
        Linux)
            PLATFORM="linux"
            ;;
        Darwin)
            PLATFORM="macos"
            ;;
        *)
            print_error "Unsupported operating system: $OS"
            exit 1
            ;;
    esac

    case "$ARCH" in
        x86_64|amd64)
            ARCH="x86_64"
            ;;
        arm64|aarch64)
            ARCH="aarch64"
            ;;
        *)
            print_error "Unsupported architecture: $ARCH"
            exit 1
            ;;
    esac

    print_step "Detected platform: $PLATFORM ($ARCH)"
}

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

# Install uv if not present
install_uv() {
    if command_exists uv; then
        print_step "uv is already installed"
        return
    fi

    print_step "Installing uv (Python package manager)..."
    curl -LsSf https://astral.sh/uv/install.sh | sh

    # Source the env file to get uv in PATH
    if [ -f "$HOME/.local/bin/env" ]; then
        source "$HOME/.local/bin/env"
    elif [ -f "$HOME/.cargo/env" ]; then
        source "$HOME/.cargo/env"
    fi

    # Add to current PATH if still not found
    if ! command_exists uv; then
        export PATH="$HOME/.local/bin:$HOME/.cargo/bin:$PATH"
    fi

    if ! command_exists uv; then
        print_error "Failed to install uv. Please install it manually: https://docs.astral.sh/uv/getting-started/installation/"
        exit 1
    fi

    print_success "uv installed successfully"
}

# Create tuft directory structure
create_directories() {
    print_step "Creating TuFT directory structure..."
    mkdir -p "$TUFT_HOME"
    mkdir -p "$TUFT_BIN"
    mkdir -p "$TUFT_HOME/checkpoints"
    mkdir -p "$TUFT_HOME/configs"
    mkdir -p "$TUFT_HOME/scripts"
}

# Create Python virtual environment and install tuft
install_tuft() {
    print_step "Creating Python $PYTHON_VERSION virtual environment..."

    # Remove existing venv if present
    if [ -d "$TUFT_VENV" ]; then
        rm -rf "$TUFT_VENV"
    fi

    uv venv --python "$PYTHON_VERSION" "$TUFT_VENV"

    print_step "Installing TuFT package..."

    # Determine package source
    if [ -n "$LOCAL_SOURCE_PATH" ]; then
        print_step "Installing from local source: $LOCAL_SOURCE_PATH"
        PACKAGE_SPEC="$LOCAL_SOURCE_PATH"
    elif [ "$INSTALL_FROM_SOURCE" = true ]; then
        print_step "Installing from GitHub: $TUFT_GIT_REPO"
        PACKAGE_SPEC="git+${TUFT_GIT_REPO}"
    else
        print_step "Installing from PyPI: $TUFT_PYPI_PACKAGE"
        PACKAGE_SPEC="$TUFT_PYPI_PACKAGE"
    fi

    # Install tuft with all extras (backend, persistence)
    uv pip install --python "$TUFT_VENV/bin/python" "${PACKAGE_SPEC}[backend,persistence]"

    print_success "TuFT installed successfully"
}

# URL for the flash-attn installation script
FLASH_ATTN_SCRIPT_URL="https://raw.githubusercontent.com/agentscope-ai/tuft/main/scripts/install_flash_attn.py"

# Install flash-attn from precompiled wheels (avoids lengthy compilation)
# Also stores the script locally for later use by install-backend command
install_flash_attn() {
    print_step "Installing flash-attn from precompiled wheels..."

    local script_path="$TUFT_HOME/scripts/install_flash_attn.py"

    # Copy or download the flash-attn install script to local storage
    if [ -n "$LOCAL_SOURCE_PATH" ] && [ -f "$LOCAL_SOURCE_PATH/scripts/install_flash_attn.py" ]; then
        print_step "Using local flash-attn install script"
        cp "$LOCAL_SOURCE_PATH/scripts/install_flash_attn.py" "$script_path"
    else
        # Download the script from GitHub and store locally
        if ! curl -fsSL "$FLASH_ATTN_SCRIPT_URL" -o "$script_path"; then
            print_warning "Could not download flash-attn install script, skipping"
            return
        fi
    fi

    # Run the script and check exit code
    if "$TUFT_VENV/bin/python" "$script_path"; then
        print_success "flash-attn installation complete"
    else
        print_warning "flash-attn installation failed. This is optional, so installation will continue."
    fi
}

# Create the tuft wrapper script
# Note: The wrapper is intentionally embedded in this install script (heredoc) rather than
# being a separate file. This ensures the wrapper is always in sync with the install script
# version and simplifies distribution. When updating the wrapper, edit the heredoc below.
# The wrapper provides CLI commands (launch, version, upgrade, etc.) that delegate to the
# Python module while handling configuration defaults and environment setup.
create_wrapper() {
    print_step "Creating tuft command wrapper..."

    cat > "$TUFT_BIN/tuft" << 'WRAPPER_EOF'
#!/bin/bash
# TuFT CLI Wrapper
# This script provides a convenient interface to the TuFT server
# Generated by install.sh - edit the heredoc in install.sh to modify

set -e

TUFT_HOME="${TUFT_HOME:-$HOME/.tuft}"
TUFT_VENV="$TUFT_HOME/venv"
TUFT_PYTHON="$TUFT_VENV/bin/python"

# Verify installation
if [ ! -f "$TUFT_PYTHON" ]; then
    echo "Error: TuFT installation not found at $TUFT_HOME"
    echo "Please reinstall TuFT using:"
    echo '  /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/agentscope-ai/tuft/main/scripts/install.sh)"'
    exit 1
fi

# Handle commands
case "${1:-}" in
    launch)
        shift
        # Pass all arguments directly to the CLI (single source of truth)
        exec "$TUFT_PYTHON" -m tuft launch "$@"
        ;;

    version|--version|-v)
        "$TUFT_PYTHON" -c "import tuft; print(f'TuFT version: {tuft.__version__}')" 2>/dev/null || \
        "$TUFT_PYTHON" -c "from importlib.metadata import version; print(f'TuFT version: {version(\"tuft\")}')"
        ;;

    upgrade)
        shift
        # Parse upgrade options
        UPGRADE_FROM_SOURCE=false
        UPGRADE_LOCAL_SOURCE=""
        while [[ $# -gt 0 ]]; do
            case "$1" in
                --from-source)
                    UPGRADE_FROM_SOURCE=true
                    shift
                    ;;
                --local-source)
                    UPGRADE_LOCAL_SOURCE="$2"
                    shift 2
                    ;;
                *)
                    echo "Unknown option: $1"
                    echo "Usage: tuft upgrade [--from-source | --local-source PATH]"
                    exit 1
                    ;;
            esac
        done

        echo "Upgrading TuFT..."
        if [ -n "$UPGRADE_LOCAL_SOURCE" ]; then
            echo "Upgrading from local source: $UPGRADE_LOCAL_SOURCE"
            uv pip install --python "$TUFT_PYTHON" --upgrade "${UPGRADE_LOCAL_SOURCE}[backend,persistence]"
        elif [ "$UPGRADE_FROM_SOURCE" = true ]; then
            echo "Upgrading from GitHub..."
            uv pip install --python "$TUFT_PYTHON" --upgrade "git+https://github.com/agentscope-ai/tuft.git#egg=tuft[backend,persistence]"
        else
            uv pip install --python "$TUFT_PYTHON" --upgrade "tuft[backend,persistence]"
        fi

        # Also update flash-attn
        echo ""
        echo "Updating flash-attn..."
        FLASH_SCRIPT_PATH="$TUFT_HOME/scripts/install_flash_attn.py"
        if [ -n "$UPGRADE_LOCAL_SOURCE" ] && [ -f "$UPGRADE_LOCAL_SOURCE/scripts/install_flash_attn.py" ]; then
            cp "$UPGRADE_LOCAL_SOURCE/scripts/install_flash_attn.py" "$FLASH_SCRIPT_PATH"
        elif [ ! -f "$FLASH_SCRIPT_PATH" ]; then
            FLASH_SCRIPT_URL="https://raw.githubusercontent.com/agentscope-ai/tuft/main/scripts/install_flash_attn.py"
            mkdir -p "$TUFT_HOME/scripts"
            curl -fsSL "$FLASH_SCRIPT_URL" -o "$FLASH_SCRIPT_PATH" 2>/dev/null || true
        fi
        if [ -f "$FLASH_SCRIPT_PATH" ]; then
            "$TUFT_PYTHON" "$FLASH_SCRIPT_PATH" || echo "Warning: flash-attn update failed (optional)"
        fi

        echo ""
        echo "TuFT upgraded successfully!"
        ;;

    uninstall)
        echo "Uninstalling TuFT..."
        read -p "This will remove $TUFT_HOME. Are you sure? [y/N] " -n 1 -r
        echo
        if [[ "$REPLY" =~ ^[Yy]$ ]]; then
            rm -rf "$TUFT_HOME"
            echo "TuFT uninstalled. Please remove $TUFT_HOME/bin from your PATH."
        else
            echo "Uninstall cancelled."
        fi
        ;;

    help|--help|-h)
        echo "TuFT - Tenant-unified Fine-Tuning Server"
        echo ""
        echo "Usage: tuft <command> [options]"
        echo ""
        echo "Commands:"
        echo "  launch            Start the TuFT server"
        echo "  version           Show TuFT version"
        echo "  upgrade           Upgrade TuFT to the latest version"
        echo "                    Options: --from-source, --local-source PATH"
        echo "  uninstall         Remove TuFT installation"
        echo "  help              Show this help message"
        echo ""
        echo "Launch options: Run 'tuft launch --help' for all available options."
        echo ""
        echo "Environment Variables:"
        echo "  TUFT_HOME            Installation directory (default: ~/.tuft)"
        echo "  TUFT_CONFIG          Default config file path"
        echo "  TUFT_HOST            Default host for launch command"
        echo "  TUFT_PORT            Default port for launch command"
        echo "  TUFT_CHECKPOINT_DIR  Default checkpoint directory"
        echo "  TUFT_LOG_LEVEL       Default log level"
        echo ""
        echo "Examples:"
        echo "  tuft launch --config tuft_config.yaml"
        echo "  tuft launch --port 10610 --config /path/to/tuft_config.yaml"
        echo "  tuft launch  # uses default config at ~/.tuft/configs/tuft_config.yaml"
        echo "  tuft upgrade"
        echo ""
        echo "Documentation: https://github.com/agentscope-ai/tuft"
        ;;

    "")
        # No command provided, show help
        "$0" help
        ;;

    *)
        # Pass through to the tuft module for any other commands
        exec "$TUFT_PYTHON" -m tuft "$@"
        ;;
esac
WRAPPER_EOF

    chmod +x "$TUFT_BIN/tuft"
    print_success "Wrapper script created at $TUFT_BIN/tuft"
}

# Create example configuration
create_example_config() {
    if [ ! -f "$TUFT_HOME/configs/tuft_config.yaml.example" ]; then
        print_step "Creating example configuration..."
        cat > "$TUFT_HOME/configs/tuft_config.yaml.example" << 'CONFIG_EOF'
# TuFT Server Configuration Example
# Copy this file to tuft_config.yaml and customize for your setup

model_owner: local

supported_models:
  - model_name: Qwen/Qwen3-8B
    model_path: Qwen/Qwen3-8B  # HuggingFace model ID or local path
    max_model_len: 32768
    tensor_parallel_size: 1
    temperature: 0.7
    top_p: 1.0
    top_k: -1

  # Add more models as needed:
  # - model_name: meta-llama/Llama-2-7b-hf
  #   model_path: /path/to/local/model
  #   max_model_len: 4096
  #   tensor_parallel_size: 1

# API Key authentication
# Format: api_key: user_identifier
authorized_users:
  my-api-key: default
  # Add more API keys as needed:
  # another-key: another-user

# Optional: Persistence configuration
# persistence:
#   mode: DISABLE  # Options: DISABLE, REDIS, FILE
#   redis_url: "redis://localhost:6379/0"
#   namespace: "persistence-tuft-server"
CONFIG_EOF
    fi
}

# Update shell configuration to add tuft to PATH
update_shell_config() {
    print_step "Configuring shell PATH..."

    SHELL_NAME="$(basename "$SHELL")"
    SHELL_CONFIG=""

    case "$SHELL_NAME" in
        bash)
            if [ -f "$HOME/.bash_profile" ]; then
                SHELL_CONFIG="$HOME/.bash_profile"
            else
                SHELL_CONFIG="$HOME/.bashrc"
            fi
            ;;
        zsh)
            SHELL_CONFIG="$HOME/.zshrc"
            ;;
        fish)
            SHELL_CONFIG="$HOME/.config/fish/config.fish"
            ;;
        *)
            print_warning "Unknown shell: $SHELL_NAME. Please add $TUFT_BIN to your PATH manually."
            return
            ;;
    esac

    # Check if PATH is already configured
    if [ -n "$SHELL_CONFIG" ] && [ -f "$SHELL_CONFIG" ]; then
        if grep -q "TUFT_HOME" "$SHELL_CONFIG" 2>/dev/null; then
            print_step "PATH already configured in $SHELL_CONFIG"
            return
        fi
    fi

    # Add to shell config
    # Use $HOME literal so the config remains portable
    if [ -n "$SHELL_CONFIG" ]; then
        if [ "$SHELL_NAME" = "fish" ]; then
            mkdir -p "$(dirname "$SHELL_CONFIG")"
            echo "" >> "$SHELL_CONFIG"
            echo "# TuFT" >> "$SHELL_CONFIG"
            echo 'set -gx TUFT_HOME $HOME/.tuft' >> "$SHELL_CONFIG"
            echo 'fish_add_path $TUFT_HOME/bin' >> "$SHELL_CONFIG"
        else
            echo "" >> "$SHELL_CONFIG"
            echo "# TuFT" >> "$SHELL_CONFIG"
            echo 'export TUFT_HOME="$HOME/.tuft"' >> "$SHELL_CONFIG"
            echo 'export PATH="$TUFT_HOME/bin:$PATH"' >> "$SHELL_CONFIG"
        fi
        print_success "Added TuFT to PATH in $SHELL_CONFIG"
    fi
}

# Print completion message
print_completion() {
    echo ""
    echo -e "${GREEN}============================================${NC}"
    echo -e "${GREEN}  TuFT installation complete!${NC}"
    echo -e "${GREEN}============================================${NC}"
    echo ""
    echo "Installation directory: $TUFT_HOME"
    echo ""
    echo "To get started:"
    echo ""
    echo "  1. Restart your terminal or run:"
    echo "     source ~/.$(basename "$SHELL")rc"
    echo ""
    echo "  2. Create a server configuration file:"
    echo "     cp $TUFT_HOME/configs/tuft_config.yaml.example $TUFT_HOME/configs/tuft_config.yaml"
    echo "     # Edit the file to configure your models and API keys"
    echo ""
    echo "  3. Launch the TuFT server:"
    echo "     tuft launch"
    echo ""
    echo "For more information:"
    echo "  tuft help"
    echo "  https://github.com/agentscope-ai/tuft"
    echo ""
}

# Main installation flow
main() {
    parse_args "$@"

    echo ""
    echo -e "${BLUE}============================================${NC}"
    echo -e "${BLUE}  TuFT Installer${NC}"
    echo -e "${BLUE}  Tenant-unified Fine-Tuning Server${NC}"
    echo -e "${BLUE}============================================${NC}"
    echo ""

    print_step "Installing with full backend support (GPU, persistence, flash-attn)"

    if [ -n "$LOCAL_SOURCE_PATH" ]; then
        print_step "Installing from local source: $LOCAL_SOURCE_PATH"
    elif [ "$INSTALL_FROM_SOURCE" = true ]; then
        print_step "Installing from GitHub (source)"
    else
        print_step "Installing from PyPI"
    fi

    # Clean existing installation if requested
    if [ "$CLEAN_INSTALL" = true ] && [ -d "$TUFT_HOME" ]; then
        print_step "Cleaning existing installation at $TUFT_HOME..."
        rm -rf "$TUFT_HOME"
        print_success "Existing installation removed"
    fi

    detect_platform
    install_uv
    create_directories
    install_tuft
    install_flash_attn
    create_wrapper
    create_example_config
    update_shell_config
    print_completion
}

# Run main
main "$@"