install_skills.sh

#!/bin/bash
#
# Databricks Skills Installer
#
# Installs Databricks skills for Claude Code into your project.
# These skills teach Claude how to work with Databricks using MCP tools.
#
# Usage:
#   # Install all skills (Databricks + MLflow + APX)
#   curl -sSL https://raw.githubusercontent.com/databricks-solutions/ai-dev-kit/main/databricks-skills/install_skills.sh | bash
#
#   # Install specific skills (can mix Databricks and MLflow skills)
#   curl -sSL https://raw.githubusercontent.com/databricks-solutions/ai-dev-kit/main/databricks-skills/install_skills.sh | bash -s -- databricks-bundles agent-evaluation
#
#   # Or run locally
#   ./install_skills.sh                              # Install all skills
#   ./install_skills.sh databricks-bundles agent-evaluation  # Install specific skills
#   ./install_skills.sh --mlflow-version v1.0.0      # Pin MLflow skills version
#   ./install_skills.sh --local                      # Install Databricks skills from local directory
#   ./install_skills.sh --install-to-genie           # Install then upload ./.claude/skills to workspace (Genie Code / Assistant)
#   ./install_skills.sh --install-to-genie --profile prod --local
#   ./install_skills.sh --list                       # List available skills
#   ./install_skills.sh --help                       # Show help
#

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

# Configuration
REPO_URL="https://github.com/databricks-solutions/ai-dev-kit"
REPO_RAW_URL="https://raw.githubusercontent.com/databricks-solutions/ai-dev-kit/main"
SKILLS_DIR=".claude/skills"
INSTALL_FROM_LOCAL=false
INSTALL_TO_GENIE=false
# Databricks CLI profile for workspace upload (only used with --install-to-genie)
DB_PROFILE="${DATABRICKS_CONFIG_PROFILE:-DEFAULT}"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# MLflow skills configuration
MLFLOW_REPO_RAW_URL="https://raw.githubusercontent.com/mlflow/skills"
MLFLOW_REPO_REF="main"

# Databricks skills (hosted in this repo)
DATABRICKS_SKILLS="databricks-agent-bricks databricks-ai-functions databricks-aibi-dashboards databricks-bundles databricks-app-python databricks-config databricks-dbsql databricks-docs databricks-genie databricks-iceberg databricks-jobs databricks-lakebase-autoscale databricks-lakebase-provisioned databricks-metric-views databricks-mlflow-evaluation databricks-model-serving databricks-python-sdk databricks-execution-compute databricks-spark-declarative-pipelines databricks-spark-structured-streaming databricks-synthetic-data-gen databricks-unity-catalog databricks-unstructured-pdf-generation databricks-vector-search databricks-zerobus-ingest spark-python-data-source"

# MLflow skills (fetched from mlflow/skills repo)
MLFLOW_SKILLS="agent-evaluation analyze-mlflow-chat-session analyze-mlflow-trace instrumenting-with-mlflow-tracing mlflow-onboarding querying-mlflow-metrics retrieving-mlflow-traces searching-mlflow-docs"

# APX skills configuration (fetched from databricks-solutions/apx repo)
APX_REPO_RAW_URL="https://raw.githubusercontent.com/databricks-solutions/apx"
APX_REPO_REF="main"
APX_REPO_SKILL_PATH="skills/apx"

# APX skills
APX_SKILLS="databricks-app-apx"

# All available skills
ALL_SKILLS="$DATABRICKS_SKILLS $MLFLOW_SKILLS $APX_SKILLS"

# Get skill description
get_skill_description() {
    case "$1" in
        # Databricks skills
        "databricks-agent-bricks") echo "Knowledge Assistants, Genie Spaces, Supervisor Agents" ;;
        "databricks-ai-functions") echo "Built-in AI Functions (classify, extract, query, forecast, parse, etc.), doc processing & custom RAG" ;;
        "databricks-aibi-dashboards") echo "Databricks AI/BI Dashboards - create and manage dashboards" ;;
        "databricks-bundles") echo "Databricks Asset Bundles - deployment and configuration" ;;                                                                                                                          
        "databricks-app-apx") echo "Databricks Apps with React/Next.js (APX framework)" ;;                                                                                                                                     
        "databricks-app-python") echo "Databricks Apps with Python (Dash, Streamlit) and foundation model integration" ;;     
        "databricks-config") echo "Profile authentication setup for Databricks" ;;
        "databricks-dbsql") echo "Databricks SQL - SQL scripting, MVs, geospatial, AI functions, federation" ;;
        "databricks-docs") echo "Documentation reference via llms.txt" ;;
        "databricks-genie") echo "Genie Spaces - create, curate, and query via Conversation API" ;;
        "databricks-iceberg") echo "Apache Iceberg - managed tables, UniForm, IRC, Snowflake interop, migration" ;;
        "databricks-jobs") echo "Databricks Lakeflow Jobs - workflow orchestration" ;;
        "databricks-python-sdk") echo "Databricks Python SDK, Connect, and REST API" ;;
        "databricks-execution-compute") echo "Execute code and manage compute on Databricks - serverless, clusters, and SQL warehouses" ;;
        "databricks-unity-catalog") echo "System tables for lineage, audit, billing" ;;
        "databricks-lakebase-autoscale") echo "Lakebase Autoscale - managed PostgreSQL with autoscaling" ;;
        "databricks-lakebase-provisioned") echo "Lakebase Provisioned - data connections and reverse ETL" ;;
        "databricks-metric-views") echo "Unity Catalog Metric Views - governed business metrics in YAML" ;;
        "databricks-model-serving") echo "Model Serving - deploy MLflow models and AI agents" ;;
        "databricks-mlflow-evaluation") echo "MLflow evaluation and trace analysis" ;;
        "databricks-spark-declarative-pipelines") echo "Spark Declarative Pipelines (SDP/LDP/DLT)" ;;
        "spark-python-data-source") echo "Spark custom Python data sources" ;;
        "databricks-spark-structured-streaming") echo "Spark Structured Streaming patterns and best practices" ;;
        "databricks-synthetic-data-gen") echo "Synthetic test data generation" ;;
        "databricks-unstructured-pdf-generation") echo "Generate synthetic PDFs for RAG" ;;
        "databricks-vector-search") echo "Vector Search - endpoints, indexes, and queries for RAG" ;;
        "databricks-zerobus-ingest") echo "Zerobus Ingest - gRPC data ingestion into Delta tables" ;;
        # MLflow skills (from mlflow/skills repo)
        "agent-evaluation") echo "End-to-end agent evaluation workflow" ;;
        "analyze-mlflow-chat-session") echo "Debug multi-turn conversations" ;;
        "analyze-mlflow-trace") echo "Debug traces, spans, and assessments" ;;
        "instrumenting-with-mlflow-tracing") echo "Add MLflow tracing to Python/TypeScript" ;;
        "mlflow-onboarding") echo "MLflow setup guide for new users" ;;
        "querying-mlflow-metrics") echo "Aggregated metrics and time-series analysis" ;;
        "retrieving-mlflow-traces") echo "Trace search and filtering" ;;
        "searching-mlflow-docs") echo "Search MLflow documentation" ;;
        # APX skills (from databricks-solutions/apx repo)
        "databricks-app-apx") echo "Databricks Apps with React/Next.js (APX framework)" ;;
        *) echo "Unknown skill" ;;
    esac
}

# Get extra files for a Databricks skill (besides SKILL.md)
get_skill_extra_files() {
    case "$1" in
        "databricks-agent-bricks") echo "1-knowledge-assistants.md 2-supervisor-agents.md" ;;
        "databricks-ai-functions") echo "1-task-functions.md 2-ai-query.md 3-ai-forecast.md 4-document-processing-pipeline.md" ;;
        "databricks-aibi-dashboards") echo "widget-reference.md sql-patterns.md" ;;
        "databricks-genie") echo "spaces.md conversation.md" ;;
        "databricks-bundles") echo "alerts_guidance.md SDP_guidance.md" ;;
        "databricks-iceberg") echo "1-managed-iceberg-tables.md 2-uniform-and-compatibility.md 3-iceberg-rest-catalog.md 4-snowflake-interop.md 5-external-engine-interop.md" ;;
        "databricks-app-apx") echo "backend-patterns.md best-practices.md frontend-patterns.md" ;;
        "databricks-app-python") echo "1-authorization.md 2-app-resources.md 3-frameworks.md 4-deployment.md 5-lakebase.md 6-mcp-approach.md examples/llm_config.py examples/fm-minimal-chat.py examples/fm-parallel-calls.py examples/fm-structured-outputs.py" ;;
        "databricks-jobs") echo "task-types.md triggers-schedules.md notifications-monitoring.md examples.md" ;;
        "databricks-python-sdk") echo "doc-index.md examples/1-authentication.py examples/2-clusters-and-jobs.py examples/3-sql-and-warehouses.py examples/4-unity-catalog.py examples/5-serving-and-vector-search.py" ;;
        "databricks-unity-catalog") echo "5-system-tables.md" ;;
        "databricks-lakebase-autoscale") echo "projects.md branches.md computes.md connection-patterns.md reverse-etl.md" ;;
        "databricks-lakebase-provisioned") echo "connection-patterns.md reverse-etl.md" ;;
        "databricks-metric-views") echo "yaml-reference.md patterns.md" ;;
        "databricks-model-serving") echo "1-classical-ml.md 2-custom-pyfunc.md 3-genai-agents.md 4-tools-integration.md 5-development-testing.md 6-logging-registration.md 7-deployment.md 8-querying-endpoints.md 9-package-requirements.md" ;;
        "databricks-mlflow-evaluation") echo "references/CRITICAL-interfaces.md references/GOTCHAS.md references/patterns-context-optimization.md references/patterns-datasets.md references/patterns-evaluation.md references/patterns-scorers.md references/patterns-trace-analysis.md references/user-journeys.md" ;;
        "databricks-spark-declarative-pipelines") echo "1-ingestion-patterns.md 2-streaming-patterns.md 3-scd-patterns.md 4-performance-tuning.md 5-python-api.md 6-dlt-migration.md 7-advanced-configuration.md 8-project-initialization.md" ;;
        "databricks-spark-structured-streaming") echo "checkpoint-best-practices.md kafka-streaming.md merge-operations.md multi-sink-writes.md stateful-operations.md stream-static-joins.md stream-stream-joins.md streaming-best-practices.md trigger-and-cost-optimization.md" ;;
        "databricks-vector-search") echo "index-types.md end-to-end-rag.md" ;;
        "databricks-zerobus-ingest") echo "1-setup-and-authentication.md 2-python-client.md 3-multilanguage-clients.md 4-protobuf-schema.md 5-operations-and-limits.md" ;;
        *) echo "" ;;
    esac
}

# Check if a skill is from MLflow repo
is_mlflow_skill() {
    local skill=$1
    for mlflow_skill in $MLFLOW_SKILLS; do
        if [ "$skill" = "$mlflow_skill" ]; then
            return 0
        fi
    done
    return 1
}

# Check if a skill is from APX repo
is_apx_skill() {
    local skill=$1
    for apx_skill in $APX_SKILLS; do
        if [ "$skill" = "$apx_skill" ]; then
            return 0
        fi
    done
    return 1
}

# Get extra files for an MLflow skill (besides SKILL.md)
get_mlflow_skill_extra_files() {
    case "$1" in
        "agent-evaluation") echo "references/dataset-preparation.md references/scorers-constraints.md references/scorers.md references/setup-guide.md references/tracing-integration.md references/troubleshooting.md scripts/analyze_results.py scripts/create_dataset_template.py scripts/list_datasets.py scripts/run_evaluation_template.py scripts/setup_mlflow.py scripts/validate_agent_tracing.py scripts/validate_auth.py scripts/validate_environment.py scripts/validate_tracing_runtime.py" ;;
        "analyze-mlflow-chat-session") echo "scripts/discover_schema.sh scripts/inspect_turn.sh" ;;
        "analyze-mlflow-trace") echo "references/trace-structure.md" ;;
        "instrumenting-with-mlflow-tracing") echo "references/advanced-patterns.md references/distributed-tracing.md references/feedback-collection.md references/production.md references/python.md references/typescript.md" ;;
        "mlflow-onboarding") echo "" ;;
        "querying-mlflow-metrics") echo "references/api_reference.md scripts/fetch_metrics.py" ;;
        "retrieving-mlflow-traces") echo "" ;;
        "searching-mlflow-docs") echo "" ;;
        *) echo "" ;;
    esac
}

# Show usage
show_help() {
    echo -e "${BLUE}Databricks Skills Installer for Claude Code${NC}"
    echo ""
    echo "Usage:"
    echo "  ./install_skills.sh [options] [skill1 skill2 ...]"
    echo ""
    echo "Options:"
    echo "  --help, -h              Show this help message"
    echo "  --list, -l              List all available skills"
    echo "  --all, -a               Install all skills (default if no skills specified)"
    echo "  --local                 Install from local files instead of downloading"
    echo "  --install-to-genie      After install, upload ./.claude/skills to workspace"
    echo "                          /Users/<you>/.assistant/skills for Genie Code (uses cwd; requires databricks CLI)"
    echo "  --profile <name>        Databricks config profile for workspace upload (default: DEFAULT or \$DATABRICKS_CONFIG_PROFILE)"
    echo "  --mlflow-version <ref>  Pin MLflow skills to specific version/branch/tag (default: main)"
    echo "  --apx-version <ref>    Pin APX skills to specific version/branch/tag (default: main)"
    echo ""
    echo "Examples:"
    echo "  ./install_skills.sh                          # Install all skills"
    echo "  ./install_skills.sh databricks-spark-declarative-pipelines  # Install specific Databricks skill"
    echo "  ./install_skills.sh agent-evaluation         # Install specific MLflow skill"
    echo "  ./install_skills.sh databricks-bundles agent-evaluation  # Mix of both sources"
    echo "  ./install_skills.sh --mlflow-version v1.0.0  # Pin MLflow skills version"
    echo "  ./install_skills.sh --apx-version v1.0.0    # Pin APX skills version"
    echo "  ./install_skills.sh --local                  # Install all from local directory"
    echo "  ./install_skills.sh --install-to-genie       # Install all, then upload to workspace for Genie Code"
    echo "  ./install_skills.sh --install-to-genie --profile prod  # Same with explicit Databricks CLI profile"
    echo "  ./install_skills.sh --list                   # List available skills"
    echo ""
    echo -e "${GREEN}Databricks Skills:${NC}"
    for skill in $DATABRICKS_SKILLS; do
        echo "  - $skill: $(get_skill_description "$skill")"
    done
    echo ""
    echo -e "${GREEN}MLflow Skills (from github.com/mlflow/skills):${NC}"
    for skill in $MLFLOW_SKILLS; do
        echo "  - $skill: $(get_skill_description "$skill")"
    done
    echo ""
    echo -e "${GREEN}APX Skills (from github.com/databricks-solutions/apx):${NC}"
    for skill in $APX_SKILLS; do
        echo "  - $skill: $(get_skill_description "$skill")"
    done
    echo ""
}

# List available skills
list_skills() {
    echo -e "${BLUE}Available Skills:${NC}"
    echo ""
    echo -e "${GREEN}Databricks Skills:${NC}"
    for skill in $DATABRICKS_SKILLS; do
        echo -e "  ${GREEN}$skill${NC}"
        echo -e "    $(get_skill_description "$skill")"
    done
    echo ""
    echo -e "${GREEN}MLflow Skills (from github.com/mlflow/skills):${NC}"
    for skill in $MLFLOW_SKILLS; do
        echo -e "  ${GREEN}$skill${NC}"
        echo -e "    $(get_skill_description "$skill")"
    done
    echo ""
    echo -e "${GREEN}APX Skills (from github.com/databricks-solutions/apx):${NC}"
    for skill in $APX_SKILLS; do
        echo -e "  ${GREEN}$skill${NC}"
        echo -e "    $(get_skill_description "$skill")"
    done
    echo ""
}

# Validate skill name
is_valid_skill() {
    local skill=$1
    for valid_skill in $ALL_SKILLS; do
        if [ "$skill" = "$valid_skill" ]; then
            return 0
        fi
    done
    return 1
}

# Upload one skill folder to workspace /Users/<user>/.assistant/skills/<skill>/
upload_skill_to_genie_workspace() {
    local skill_dir="$1"
    local skills_path="$2"
    local db_profile="$3"

    skill_dir="${skill_dir%/}"
    local skill_name
    skill_name=$(basename "$skill_dir")

    if [[ "$skill_name" == "."* ]] || [[ "$skill_name" == "TEMPLATE" ]] || [[ ! -d "$skill_dir" ]]; then
        return 0
    fi
    if [[ ! -f "$skill_dir/SKILL.md" ]]; then
        return 0
    fi

    echo -e "  ${GREEN}Uploading${NC} $skill_name"
    databricks workspace mkdirs "$skills_path/$skill_name" --profile "$db_profile" 2>/dev/null || true

    while IFS= read -r -d '' file; do
        rel_path="${file#$skill_dir/}"
        dest_path="$skills_path/$skill_name/$rel_path"
        parent_dir=$(dirname "$dest_path")
        if [[ "$parent_dir" != "$skills_path/$skill_name" ]]; then
            databricks workspace mkdirs "$parent_dir" --profile "$db_profile" 2>/dev/null || true
        fi
        databricks workspace import "$dest_path" --file "$file" --profile "$db_profile" --format AUTO --overwrite 2>/dev/null || true
    done < <(find "$skill_dir" -type f \( -name "*.md" -o -name "*.py" -o -name "*.yaml" -o -name "*.yml" -o -name "*.sh" \) -print0)
}

# Upload all skills under SKILLS_DIR to workspace for Genie Code / Assistant agent mode
install_skills_to_genie_workspace() {
    if ! command -v databricks >/dev/null 2>&1; then
        echo -e "${RED}Error: databricks CLI not found. Install it to use --install-to-genie.${NC}"
        return 1
    fi

    local abs_skills_dir
    if [ ! -d "$SKILLS_DIR" ]; then
        echo -e "${RED}Error: Skills directory not found: ${SKILLS_DIR} (run from the directory where skills were installed).${NC}"
        return 1
    fi
    abs_skills_dir="$(cd "$SKILLS_DIR" && pwd)"

    echo ""
    echo -e "${BLUE}════════════════════════════════════════════════════════════${NC}"
    echo -e "${BLUE}  Uploading skills to workspace (Genie Code)${NC}"
    echo -e "${BLUE}════════════════════════════════════════════════════════════${NC}"
    echo -e "Databricks profile: ${DB_PROFILE}"
    echo -e "Local skills (cwd-relative): ${SKILLS_DIR}/ → ${abs_skills_dir}"
    echo ""

    local user_name
    user_name=$(databricks current-user me --profile "$DB_PROFILE" --output json 2>/dev/null | python3 -c "import sys, json; d=json.load(sys.stdin); print(d.get('userName', ''))" 2>/dev/null || echo "")
    if [ -z "$user_name" ]; then
        echo -e "${RED}Error: Could not determine workspace user. Check authentication and --profile.${NC}"
        return 1
    fi

    local skills_path="/Users/$user_name/.assistant/skills"
    echo -e "Workspace user: ${user_name}"
    echo -e "Workspace path: ${skills_path}"
    echo ""

    echo -e "${GREEN}Creating workspace skills directory...${NC}"
    databricks workspace mkdirs "$skills_path" --profile "$DB_PROFILE" 2>/dev/null || true

    echo -e "${GREEN}Uploading skills...${NC}"
    local skill_dir
    for skill_dir in "$abs_skills_dir"/*/; do
        [ -d "$skill_dir" ] || continue
        upload_skill_to_genie_workspace "$skill_dir" "$skills_path" "$DB_PROFILE"
    done

    echo ""
    echo -e "${GREEN}Workspace listing:${NC}"
    databricks workspace list "$skills_path" --profile "$DB_PROFILE" 2>/dev/null || echo -e "  ${YELLOW}(Could not list workspace path)${NC}"

    local skills_count
    skills_count=$(find "$abs_skills_dir" -maxdepth 1 -type d -exec test -f {}/SKILL.md \; -print 2>/dev/null | wc -l | tr -d ' ')
    echo ""
    echo -e "${GREEN}Genie Code upload complete.${NC} ${skills_count} skills under ${skills_path}"
    echo ""
    return 0
}

# Function to download a Databricks skill
download_databricks_skill() {
    local skill_name=$1
    local skill_dir="$SKILLS_DIR/$skill_name"

    if [ "$INSTALL_FROM_LOCAL" = true ]; then
        # Copy from local files
        echo -e "  Copying from local..."
        local source_dir="$SCRIPT_DIR/${skill_name}"

        # Check if source directory exists
        if [ ! -d "$source_dir" ]; then
            echo -e "  ${RED}✗${NC} Source directory not found: $source_dir"
            rm -rf "$skill_dir"
            return 1
        fi

        # Copy SKILL.md (required)
        if [ -f "$source_dir/SKILL.md" ]; then
            cp "$source_dir/SKILL.md" "$skill_dir/SKILL.md"
            echo -e "  ${GREEN}✓${NC} Copied SKILL.md"
        else
            echo -e "  ${RED}✗${NC} SKILL.md not found in $source_dir"
            rm -rf "$skill_dir"
            return 1
        fi

        # Copy skill-specific extra files
        local extra_files=$(get_skill_extra_files "$skill_name")
        if [ -n "$extra_files" ]; then
            for extra_file in $extra_files; do
                if [ -f "$source_dir/${extra_file}" ]; then
                    # Create subdirectory if needed
                    local extra_file_dir=$(dirname "$skill_dir/${extra_file}")
                    mkdir -p "$extra_file_dir"
                    cp "$source_dir/${extra_file}" "$skill_dir/${extra_file}"
                    echo -e "  ${GREEN}✓${NC} Copied ${extra_file}"
                else
                    echo -e "  ${YELLOW}○${NC} Optional file ${extra_file} not found"
                fi
            done
        fi
    else
        # Download from URL
        echo -e "  Downloading from Databricks repo..."

        # Download SKILL.md (required)
        if curl -sSL -f "${REPO_RAW_URL}/databricks-skills/${skill_name}/SKILL.md" -o "$skill_dir/SKILL.md" 2>/dev/null; then
            echo -e "  ${GREEN}✓${NC} Downloaded SKILL.md"
        else
            echo -e "  ${RED}✗${NC} Failed to download SKILL.md"
            rm -rf "$skill_dir"
            return 1
        fi

        # Download skill-specific extra files
        local extra_files=$(get_skill_extra_files "$skill_name")
        if [ -n "$extra_files" ]; then
            for extra_file in $extra_files; do
                # Create subdirectory if needed
                local extra_file_dir=$(dirname "$skill_dir/${extra_file}")
                mkdir -p "$extra_file_dir"
                if curl -sSL -f "${REPO_RAW_URL}/databricks-skills/${skill_name}/${extra_file}" -o "$skill_dir/${extra_file}" 2>/dev/null; then
                    echo -e "  ${GREEN}✓${NC} Downloaded ${extra_file}"
                else
                    echo -e "  ${YELLOW}○${NC} Optional file ${extra_file} not found"
                fi
            done
        fi
    fi

    return 0
}

# Function to download an MLflow skill
download_mlflow_skill() {
    local skill_name=$1
    local skill_dir="$SKILLS_DIR/$skill_name"

    echo -e "  Downloading from MLflow repo (${MLFLOW_REPO_REF})..."

    # Download SKILL.md (required)
    if curl -sSL -f "${MLFLOW_REPO_RAW_URL}/${MLFLOW_REPO_REF}/${skill_name}/SKILL.md" -o "$skill_dir/SKILL.md" 2>/dev/null; then
        echo -e "  ${GREEN}✓${NC} Downloaded SKILL.md"
    else
        echo -e "  ${RED}✗${NC} Failed to download SKILL.md from MLflow repo"
        rm -rf "$skill_dir"
        return 1
    fi

    # Download skill-specific extra files
    local extra_files=$(get_mlflow_skill_extra_files "$skill_name")
    if [ -n "$extra_files" ]; then
        for extra_file in $extra_files; do
            # Create subdirectory if needed
            local extra_file_dir=$(dirname "$skill_dir/${extra_file}")
            mkdir -p "$extra_file_dir"
            if curl -sSL -f "${MLFLOW_REPO_RAW_URL}/${MLFLOW_REPO_REF}/${skill_name}/${extra_file}" -o "$skill_dir/${extra_file}" 2>/dev/null; then
                echo -e "  ${GREEN}✓${NC} Downloaded ${extra_file}"
            else
                echo -e "  ${YELLOW}○${NC} Optional file ${extra_file} not found"
            fi
        done
    fi

    return 0
}

# Get extra files for an APX skill (besides SKILL.md)
get_apx_skill_extra_files() {
    case "$1" in
        "databricks-app-apx") echo "backend-patterns.md frontend-patterns.md" ;;
        *) echo "" ;;
    esac
}

# Function to download an APX skill
download_apx_skill() {
    local skill_name=$1
    local skill_dir="$SKILLS_DIR/$skill_name"
    local base_url="${APX_REPO_RAW_URL}/${APX_REPO_REF}/${APX_REPO_SKILL_PATH}"

    echo -e "  Downloading from APX repo (${APX_REPO_REF})..."

    # Download SKILL.md (required)
    if curl -sSL -f "${base_url}/SKILL.md" -o "$skill_dir/SKILL.md" 2>/dev/null; then
        echo -e "  ${GREEN}✓${NC} Downloaded SKILL.md"
    else
        echo -e "  ${RED}✗${NC} Failed to download SKILL.md from APX repo"
        rm -rf "$skill_dir"
        return 1
    fi

    # Download skill-specific extra files
    local extra_files=$(get_apx_skill_extra_files "$skill_name")
    if [ -n "$extra_files" ]; then
        for extra_file in $extra_files; do
            if curl -sSL -f "${base_url}/${extra_file}" -o "$skill_dir/${extra_file}" 2>/dev/null; then
                echo -e "  ${GREEN}✓${NC} Downloaded ${extra_file}"
            else
                echo -e "  ${YELLOW}○${NC} Optional file ${extra_file} not found"
            fi
        done
    fi

    return 0
}

# Function to download a skill (routes to appropriate download function)
download_skill() {
    local skill_name=$1
    local skill_dir="$SKILLS_DIR/$skill_name"

    echo -e "\n${BLUE}Processing skill: ${skill_name}${NC}"

    # Remove existing skill directory to ensure clean install
    if [ -d "$skill_dir" ]; then
        rm -rf "$skill_dir"
    fi

    # Create skill directory
    mkdir -p "$skill_dir"

    # Route to appropriate download function
    if is_mlflow_skill "$skill_name"; then
        if [ "$INSTALL_FROM_LOCAL" = true ]; then
            echo -e "  ${RED}✗${NC} MLflow skills cannot be installed from local (they are fetched from github.com/mlflow/skills)"
            rm -rf "$skill_dir"
            return 1
        fi
        download_mlflow_skill "$skill_name"
    elif is_apx_skill "$skill_name"; then
        if [ "$INSTALL_FROM_LOCAL" = true ]; then
            echo -e "  ${RED}✗${NC} APX skills cannot be installed from local (they are fetched from github.com/databricks-solutions/apx)"
            rm -rf "$skill_dir"
            return 1
        fi
        download_apx_skill "$skill_name"
    else
        download_databricks_skill "$skill_name"
    fi

    local result=$?
    if [ $result -eq 0 ]; then
        echo -e "  ${GREEN}✓ Installed successfully${NC}"
    fi
    return $result
}

# Parse arguments
SKILLS_TO_INSTALL=""

while [ $# -gt 0 ]; do
    case $1 in
        --help|-h)
            show_help
            exit 0
            ;;
        --list|-l)
            list_skills
            exit 0
            ;;
        --all|-a)
            SKILLS_TO_INSTALL="$ALL_SKILLS"
            shift
            ;;
        --local)
            INSTALL_FROM_LOCAL=true
            shift
            ;;
        --install-to-genie|--deploy-to-assistant)
            INSTALL_TO_GENIE=true
            shift
            ;;
        --profile)
            if [ -z "$2" ] || [ "${2:0:1}" = "-" ]; then
                echo -e "${RED}Error: --profile requires a profile name${NC}"
                exit 1
            fi
            DB_PROFILE="$2"
            shift 2
            ;;
        --mlflow-version)
            if [ -z "$2" ] || [ "${2:0:1}" = "-" ]; then
                echo -e "${RED}Error: --mlflow-version requires a version/ref argument${NC}"
                exit 1
            fi
            MLFLOW_REPO_REF="$2"
            shift 2
            ;;
        --apx-version)
            if [ -z "$2" ] || [ "${2:0:1}" = "-" ]; then
                echo -e "${RED}Error: --apx-version requires a version/ref argument${NC}"
                exit 1
            fi
            APX_REPO_REF="$2"
            shift 2
            ;;
        -*)
            echo -e "${RED}Unknown option: $1${NC}"
            echo "Use --help for usage information."
            exit 1
            ;;
        *)
            # Validate skill name
            if is_valid_skill "$1"; then
                SKILLS_TO_INSTALL="$SKILLS_TO_INSTALL $1"
            else
                echo -e "${RED}Unknown skill: $1${NC}"
                echo ""
                echo "Available skills:"
                for skill in $ALL_SKILLS; do
                    echo "  - $skill"
                done
                echo ""
                echo "Use --list for more details."
                exit 1
            fi
            shift
            ;;
    esac
done

# If no skills specified, install all
if [ -z "$SKILLS_TO_INSTALL" ]; then
    SKILLS_TO_INSTALL="$ALL_SKILLS"
fi

# Header
echo -e "${BLUE}╔════════════════════════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║        Databricks Skills Installer for Claude Code         ║${NC}"
echo -e "${BLUE}╚════════════════════════════════════════════════════════════╝${NC}"
echo ""

# Check if we're in a git repo or project directory
if [ ! -d ".git" ] && [ ! -f "pyproject.toml" ] && [ ! -f "package.json" ] && [ ! -f "databricks.yml" ]; then
    echo -e "${YELLOW}Warning: This doesn't look like a project root directory.${NC}"
    echo -e "Current directory: $(pwd)"
    read -p "Continue anyway? (y/N): " confirm
    if [ "$confirm" != "y" ] && [ "$confirm" != "Y" ]; then
        echo "Aborted."
        exit 1
    fi
fi

# Create .claude/skills directory if it doesn't exist
if [ ! -d "$SKILLS_DIR" ]; then
    echo -e "${GREEN}Creating $SKILLS_DIR directory...${NC}"
    mkdir -p "$SKILLS_DIR"
fi

# Show what will be installed
echo -e "${GREEN}Skills to install:${NC}"
for skill in $SKILLS_TO_INSTALL; do
    echo -e "  - $skill"
done

if [ "$INSTALL_FROM_LOCAL" = true ]; then
    echo -e "\n${BLUE}Installing from local directory: ${SCRIPT_DIR}${NC}"
else
    echo -e "\n${BLUE}Installing from: ${REPO_URL}${NC}"
fi

# Download each skill
echo -e "\n${GREEN}Installing Databricks skills...${NC}"
installed=0
failed=0

for skill in $SKILLS_TO_INSTALL; do
    if download_skill "$skill"; then
        installed=$((installed + 1))
    else
        failed=$((failed + 1))
    fi
done

# Summary
echo ""
echo -e "${BLUE}════════════════════════════════════════════════════════════${NC}"
echo -e "${GREEN}Installation complete!${NC}"
echo -e "  Installed: ${installed} skills"
if [ $failed -gt 0 ]; then
    echo -e "  ${RED}Failed: ${failed} skills${NC}"
fi
echo ""
echo -e "${BLUE}Skills installed to: ${SKILLS_DIR}/${NC}"
echo ""
echo -e "Installed skills:"
for skill in $SKILLS_TO_INSTALL; do
    if [ -d "$SKILLS_DIR/$skill" ]; then
        echo -e "  ${GREEN}✓${NC} $skill"
    fi
done

if [ "$INSTALL_TO_GENIE" = true ]; then
    install_skills_to_genie_workspace || exit 1
fi