setup.sh

#!/bin/bash

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

GRAY_START='\x1b[90m'
GRAY_END='\x1b[39m'

APP_IMAGE="hub.simplito.com/deepfellow/deepfellow-workspace:latest"
APP_CONTAINER_NAME="workspace-app"
DEV_MODE=false
DEEPFELLOW_ENV_PATH="$HOME/.deepfellow/server/.env"
DEEPFELLOW_INFRA_ENV_PATH="$HOME/.deepfellow/infra/.env"

DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
ENV_PATH=$DIR/.env

# Ensure ~/.local/bin is in PATH for this session (pipx installs there)
export PATH="$HOME/.local/bin:$HOME/.local/share:$PATH"

for arg in "$@"; do
    case "$arg" in
        --dev) DEV_MODE=true ;;
    esac
done

if $DEV_MODE; then
    APP_IMAGE="df-workspace"
fi

# ---------- Default configuration ----------

CFG_DOCKER_NETWORK="deepfellow-workspace"
CFG_INFRA_NAME="App infra"
CFG_INFRA_PORT="8086"
CFG_INFRA_URL="http://infra:8086"
CFG_SERVER_PORT="8000"
CFG_APP_MONGODB_NAME="app-mongo"
CFG_APP_MONGODB_PORT="27018"
CFG_VECTORDB_TYPE="milvus"
CFG_VECTORDB_URL="http://milvus:19530"
CFG_VECTORDB_DATABASE="deepfellow"
CFG_EMBEDDING_MODEL="mxbai-embed-large"
CFG_EMBEDDING_SIZE="1024"
CFG_APP_PORT="3000"
CFG_APP_MONGODB_DATABASE="workspace_dev"
CFG_LOG_LEVEL="debug"
CFG_SESSION_TITLE_MODEL="gemma4:e4b"
CFG_SMALL_MODEL="qwen3.5:4b"

APP_ENV_PATH="$DIR/.env.app"

# ---------- Helper functions ----------

print_start() {
    echo -e "${GRAY_START}┌${GRAY_END}  $1"
}

print_end() {
    echo -e "${GRAY_START}└${GRAY_END}  $1"
}

print_info() {
    echo -e "${GRAY_START}│${GRAY_END}  $1"
}

print_error() {
    echo -e "${GRAY_START}│${GRAY_END}  ${RED}[ERROR]${NC} $1"
}

print_warning() {
    echo -e "${GRAY_START}│${GRAY_END}  ${YELLOW}[WARNING]${NC} $1"
}

# Prompt the user for input with a default value.
# Usage:  prompt_input_default "Label" VARIABLE "default_value"
# Empty input uses the default; otherwise value is trimmed.
prompt_input_default() {
    local label="$1"
    local varname="$2"
    local default_val="$3"
    local value=""

    read -rp "$(echo -e "${GRAY_START}│${GRAY_END}  ${label} [${default_val}]: ")" value
    value=$(echo "$value" | xargs)

    if [ -z "$value" ]; then
        value="$default_val"
    fi

    eval "$varname='$value'"
}

# Prompt the user for input with validation.
# Usage:  prompt_input "Label" VARIABLE
#         prompt_input --secret "Label" VARIABLE
# The value is trimmed; empty input re-prompts.
prompt_input() {
    local secret=false
    if [ "$1" = "--secret" ]; then
        secret=true
        shift
    fi

    local label="$1"
    local varname="$2"
    local value=""

    while true; do
        if $secret; then
            read -rsp "$(echo -e "${GRAY_START}│${GRAY_END}  ${label}: ")" value
            echo  # move to next line after hidden input
        else
            read -rp "$(echo -e "${GRAY_START}│${GRAY_END}  ${label}: ")" value
        fi

        # Trim whitespace
        value=$(echo "$value" | xargs)

        if [ -n "$value" ]; then
            eval "$varname='$value'"
            return 0
        fi

        print_warning "Value cannot be empty"
    done
}

# Prompt the user for a yes/no confirmation.
# Usage:  prompt_confirm "Do you want to continue?"
# Returns 0 (yes) or 1 (no).
prompt_confirm() {
    local label="$1"
    local answer=""

    while true; do
        read -rp "$(echo -e "${GRAY_START}│${GRAY_END}  ${label} [y/n]: ")" answer
        case "$answer" in
            [Yy]) return 0 ;;
            [Nn]) return 1 ;;
            *) print_warning "Please enter y or n" ;;
        esac
    done
}

# Print an error with command output and prompt the user to continue or abort.
# Usage:  prompt_error "Something failed"
# Reads $RUN_OUTPUT automatically if set.
prompt_error() {
    local message="$1"
    print_error "$message"
    if [ -n "$RUN_OUTPUT" ]; then
        print_error "$RUN_OUTPUT"
    fi
    if ! prompt_confirm "Do you want to continue anyway?"; then
        print_end "Installation aborted"
        exit 1
    fi
}

# Run a command with a spinner and live tail of the last 10 log lines.
# Usage:  run_with_status "Installing widgets..." "npm install 2>&1"
# After return the full command output is available in $RUN_OUTPUT.
run_with_status() {
    local MESSAGE="$1"
    local CMD="$2"

    local SPINNER_CHARS=("⠋" "⠙" "⠹" "⠸" "⠼" "⠴" "⠦" "⠧" "⠇" "⠏")
    local MAX_LINES=10
    local PREV_LINE_COUNT=0
    local TERM_WIDTH
    TERM_WIDTH=$(tput cols 2>/dev/null || echo 80)

    local TEMP_LOG
    TEMP_LOG=$(mktemp)

    # Run command in a subshell so pipes / redirects in CMD work
    bash -c "$CMD" > "$TEMP_LOG" 2>&1 &
    local CMD_PID=$!

    local i=0
    tput civis 2>/dev/null          # hide cursor

    while kill -0 "$CMD_PID" 2>/dev/null; do
        local SPIN="${SPINNER_CHARS[$((i % ${#SPINNER_CHARS[@]}))]}"

        # Erase the lines we drew last iteration
        if [ "$PREV_LINE_COUNT" -gt 0 ]; then
            for ((j = 0; j < PREV_LINE_COUNT; j++)); do
                printf "\033[A\033[2K"
            done
        fi

        # Spinner line (replaces the │ gutter)
        printf "${GRAY_START}%s${GRAY_END}  %s\n" "$SPIN" "$MESSAGE"

        # Last N log lines
        local LINE_COUNT=1          # 1 = the spinner line above
        if [ -s "$TEMP_LOG" ]; then
            while IFS= read -r line; do
                local TRUNCATED="${line:0:$((TERM_WIDTH - 5))}"
                printf "${GRAY_START}│${GRAY_END}  ${GRAY_START}%s${GRAY_END}\n" "$TRUNCATED"
                ((LINE_COUNT++))
            done < <(tail -n "$MAX_LINES" "$TEMP_LOG")
        fi

        PREV_LINE_COUNT=$LINE_COUNT
        sleep 0.08
        ((i++))
    done

    # Clear the spinner / log tail so the caller can print its own result
    if [ "$PREV_LINE_COUNT" -gt 0 ]; then
        for ((j = 0; j < PREV_LINE_COUNT; j++)); do
            printf "\033[A\033[2K"
        done
    fi

    tput cnorm 2>/dev/null          # restore cursor

    wait "$CMD_PID"
    local EXIT_CODE=$?

    RUN_OUTPUT=$(cat "$TEMP_LOG")
    rm -f "$TEMP_LOG"

    return $EXIT_CODE
}

# ---------- Installation state checks ----------

# Infra is considered installed only if its env file (with secrets) exists.
infra_installed() {
    [ -f "$DEEPFELLOW_INFRA_ENV_PATH" ]
}

# Server is considered installed only if its env file (with secrets) exists.
server_installed() {
    [ -f "$DEEPFELLOW_ENV_PATH" ]
}

# A meaningful existing installation requires at least one component's env file.
# The ~/.deepfellow directory alone is not enough — it can exist without infra
# or server being installed.
deepfellow_installed() {
    infra_installed || server_installed
}

# ---------- Dependency checks ----------

check_python_version() {
    local VERSION
    VERSION=$(python3 --version 2>&1 | awk '{print $2}')
    local MAJOR MINOR
    MAJOR=$(echo "$VERSION" | cut -d. -f1)
    MINOR=$(echo "$VERSION" | cut -d. -f2)

    if [ "$MAJOR" -lt 3 ] || ([ "$MAJOR" -eq 3 ] && [ "$MINOR" -lt 13 ]); then
        print_error "Python 3.13+ required, found Python $VERSION"
        return 1
    fi
    return 0
}

check_docker_version() {
    local VERSION
    VERSION=$(docker --version | awk '{print $3}' | sed 's/,$//')
    local MAJOR
    MAJOR=$(echo "$VERSION" | cut -d. -f1)

    if [ "$MAJOR" -lt 28 ]; then
        print_error "Docker 28.0.0+ required, found Docker $VERSION"
        return 1
    fi
    return 0
}

check_git() {
    if ! command -v git &> /dev/null; then
        print_error "Git is not installed"
        return 1
    fi
    return 0
}

# ---------- Installation functions ----------

install_deepfellow_cli() {
    if run_with_status "Installing DeepFellow CLI..." \
           "curl -fsSL https://deepfellow.ai/install.sh | bash"; then
        print_info "DeepFellow CLI installed successfully"
        print_info "Installed to: $(which deepfellow 2>/dev/null || echo "$HOME/.local/bin/deepfellow")"
    else
        print_error "Failed to install DeepFellow CLI"
        print_error "$RUN_OUTPUT"
        return 1
    fi
}

install_deepfellow_infra() {
    if run_with_status "Installing DeepFellow Infra..." \
        "deepfellow --non-interactive infra install --port $CFG_INFRA_PORT --infra-name \"$CFG_INFRA_NAME\" --infra-url $CFG_INFRA_URL --docker-network $CFG_DOCKER_NETWORK"; then
        print_info "DeepFellow Infra installed successfully"
    else
        print_error "Failed to install DeepFellow Infra"
        print_error "$RUN_OUTPUT"
        return 1
    fi

    # Extract Infra API key from the generated env file
    if [ -f "$DEEPFELLOW_INFRA_ENV_PATH" ]; then
        CFG_INFRA_API_KEY=$(grep -E '^DF_INFRA_API_KEY=' "$DEEPFELLOW_INFRA_ENV_PATH" | cut -d'=' -f2-)
    else
        print_warning "Infra env file not found at $DEEPFELLOW_INFRA_ENV_PATH"
    fi

    # Fallback: try to extract from install output
    if [ -z "$CFG_INFRA_API_KEY" ]; then
        CFG_INFRA_API_KEY=$(echo "$RUN_OUTPUT" | grep -oE 'API_KEY=(.+)' | head -1 | cut -d'=' -f2-)
    fi

    if [ -z "$CFG_INFRA_API_KEY" ]; then
        print_error "Could not extract Infra API key from installation"
        print_error "Check $DEEPFELLOW_INFRA_ENV_PATH for DF_INFRA_API_KEY"
        return 1
    fi
    print_info "Infra API key extracted"

    "deepfellow" infra stop >/dev/null 2>&1 || true

    print_info "Starting DeepFellow Infra..."
    if OUTPUT=$("deepfellow" infra start 2>&1); then
        print_info "DeepFellow Infra started successfully"
    else
        print_error "Failed to start DeepFellow Infra"
        print_error "$OUTPUT"
        return 1
    fi

    return 0
}

setup_docker_environment() {
    # Create the shared Docker network
    if docker network ls --format '{{.Name}}' | grep -q "^${CFG_DOCKER_NETWORK}$"; then
        print_info "Docker network '$CFG_DOCKER_NETWORK' already exists"
    else
        if docker network create "$CFG_DOCKER_NETWORK" > /dev/null 2>&1; then
            print_info "Docker network '$CFG_DOCKER_NETWORK' created"
        else
            print_error "Failed to create Docker network '$CFG_DOCKER_NETWORK'"
            return 1
        fi
    fi

    return 0
}

setup_app_mongodb() {
    if docker ps --format '{{.Names}}' | grep -q "^${CFG_APP_MONGODB_NAME}$"; then
        print_info "App MongoDB container '$CFG_APP_MONGODB_NAME' is already running"
        return 0
    fi

    if docker ps -a --format '{{.Names}}' | grep -q "^${CFG_APP_MONGODB_NAME}$"; then
        print_info "Removing stopped app MongoDB container..."
        docker rm "$CFG_APP_MONGODB_NAME" > /dev/null 2>&1
    fi

    if ! run_with_status "Starting app MongoDB container..." \
        "docker run -d \
            --name $CFG_APP_MONGODB_NAME \
            --network $CFG_DOCKER_NETWORK \
            --restart unless-stopped \
            -p ${CFG_APP_MONGODB_PORT}:27017 \
            mongo:latest --replSet rs0"; then
        prompt_error "Failed to start app MongoDB container"
        return 1
    fi
    print_info "App MongoDB container '$CFG_APP_MONGODB_NAME' started on port $CFG_APP_MONGODB_PORT"

    sleep 3
    if ! run_with_status "Initializing app MongoDB replica set..." \
        "docker exec $CFG_APP_MONGODB_NAME mongosh --quiet --eval 'rs.initiate()'"; then
        prompt_error "Failed to initialize app MongoDB replica set"
        return 1
    fi
    print_info "App MongoDB replica set initialized"

    return 0
}

install_deepfellow_server() {
    local cmd="deepfellow --non-interactive server install"
    cmd="$cmd --port $CFG_SERVER_PORT"
    cmd="$cmd --infra-url $CFG_INFRA_URL"
    cmd="$cmd --infra-api-key $CFG_INFRA_API_KEY"
    cmd="$cmd --docker-network $CFG_DOCKER_NETWORK"
    cmd="$cmd --vectordb-type $CFG_VECTORDB_TYPE"
    cmd="$cmd --vectordb-url $CFG_VECTORDB_URL"
    cmd="$cmd --vectordb-database-name $CFG_VECTORDB_DATABASE"
    cmd="$cmd --embedding-model $CFG_EMBEDDING_MODEL"
    cmd="$cmd --embedding-size $CFG_EMBEDDING_SIZE"

    if run_with_status "Installing DeepFellow Server..." "$cmd"; then
        print_info "DeepFellow Server installed successfully"
        return 0
    else
        print_error "Failed to install DeepFellow Server"
        print_error "$RUN_OUTPUT"
        return 1
    fi
}

start_deepfellow_server() {
    "deepfellow" server stop >/dev/null 2>&1 || true

    if run_with_status "Starting DeepFellow Server..." "deepfellow server start" ; then
        print_info "DeepFellow Server started successfully"
        return 0
    else
        print_error "Failed to start DeepFellow Server"
        print_error "$RUN_OUTPUT"
        return 1
    fi
}

generate_app_env() {
    local BETTER_AUTH_SECRET
    BETTER_AUTH_SECRET=$(openssl rand -hex 32)

    local MONGO_DB_URL="mongodb://${CFG_APP_MONGODB_NAME}:27017"

    local TRUSTED_ORIGINS="http://127.0.0.1:${CFG_APP_PORT},http://localhost:${CFG_APP_PORT}"
    if [ -n "$APP_URL" ]; then
        TRUSTED_ORIGINS="${TRUSTED_ORIGINS},${APP_URL}"
    fi

    cat > "$APP_ENV_PATH" <<EOF
HOST=0.0.0.0
PORT=$CFG_APP_PORT
BETTER_AUTH_SECRET=$BETTER_AUTH_SECRET
TRUSTED_ORIGINS=$TRUSTED_ORIGINS
BOOTSTRAP_ADMIN_EMAIL=$ADMIN_EMAIL
BOOTSTRAP_ADMIN_PASSWORD=$ADMIN_PASSWORD
BOOTSTRAP_ADMIN_NAME=$ADMIN_NAME
MONGO_DB_URL=$MONGO_DB_URL
MONGO_DB_NAME=$CFG_APP_MONGODB_DATABASE
DF_SERVER_URL=http://server:${CFG_SERVER_PORT}
DF_SERVER_API_KEY=$DF_SERVER_API_KEY
DF_SERVER_ORGANIZATION_ID=$DF_SERVER_ORGANIZATION_ID
DF_SERVER_PROJECT_ID=$DF_SERVER_PROJECT_ID
LOG_LEVEL=$CFG_LOG_LEVEL
SESSION_TITLE_MODEL=${LLM_MODEL:-$CFG_SESSION_TITLE_MODEL}
SMALL_MODEL=${LLM_MODEL:-$CFG_SMALL_MODEL}
EOF

    print_info "App environment file generated at $APP_ENV_PATH"
    return 0
}

setup_app_container() {
    # Stop and remove existing container if present
    if docker ps -a --format '{{.Names}}' | grep -q "^${APP_CONTAINER_NAME}$"; then
        print_info "Removing existing app container..."
        docker stop "$APP_CONTAINER_NAME" 2>/dev/null
        docker rm "$APP_CONTAINER_NAME" 2>/dev/null
    fi

    local ENV_FLAG=""
    if [ -f "$APP_ENV_PATH" ]; then
        ENV_FLAG="--env-file $APP_ENV_PATH"
        print_info "Passing app env file to container: $APP_ENV_PATH"
    else
        print_warning "App env file not found at $APP_ENV_PATH"
        print_warning "App container will start without environment variables"
    fi

    print_info "Creating and starting app container..."
    if docker run -d \
        --name "$APP_CONTAINER_NAME" \
        --restart unless-stopped \
        --network "$CFG_DOCKER_NETWORK" \
        -p "${CFG_APP_PORT}:${CFG_APP_PORT}" \
        $ENV_FLAG \
        "$APP_IMAGE" &>/dev/null; then
        print_info "App container '$APP_CONTAINER_NAME' started successfully"
        return 0
    else
        print_error "Failed to start app container"
        return 1
    fi
}

# =============================================
#  Main setup flow
# =============================================

print_start "DeepFellow + App Setup"

# ------ Step 1: Check system dependencies ------

print_info "Checking system dependencies..."

if ! command -v python3 &> /dev/null; then
    print_error "Python 3 is not installed."
    print_end "Please install Python 3.13+ from https://www.python.org"
    exit 1
fi

if ! check_python_version; then
    print_end "Python version requirement not met"
    exit 1
fi
print_info "Python: $(python3 --version)"

if ! command -v docker &> /dev/null; then
    print_error "Docker is not installed."
    print_end "Please install Docker 28+ from https://docs.docker.com/get-docker/"
    exit 1
fi

if ! check_docker_version; then
    print_end "Docker version requirement not met"
    exit 1
fi
print_info "Docker: $(docker --version)"

if ! docker info &> /dev/null; then
    print_error "Docker daemon is not running."
    print_end "Please start Docker service."
    exit 1
fi
print_info "Docker daemon is running"

if ! check_git; then
    print_end "Please install Git 2.30+ from https://git-scm.com"
    exit 1
fi
print_info "Git: $(git --version)"

print_info "All system dependencies satisfied"

# ------ Check if required ports are available ------

REQUIRED_PORTS=(
    "$CFG_INFRA_PORT:Infra"
    "$CFG_SERVER_PORT:DeepFellow Server"
    "$CFG_APP_MONGODB_PORT:App MongoDB"
    "$CFG_APP_PORT:App"
)

BUSY_PORTS=()
for entry in "${REQUIRED_PORTS[@]}"; do
    port="${entry%%:*}"
    label="${entry#*:}"
    if lsof -iTCP:"$port" -sTCP:LISTEN -t &> /dev/null; then
        BUSY_PORTS+=("  - port $port ($label)")
    fi
done

if [ ${#BUSY_PORTS[@]} -gt 0 ]; then
    print_error "The following ports are already in use:"
    for line in "${BUSY_PORTS[@]}"; do
        print_error "$line"
    done
    print_end "Free these ports and try again."
    exit 1
fi
print_info "All required ports are available"

# ------ Step 2: Install DeepFellow CLI, Infra & Server ------

EXISTING_ENV_LOADED=false

# Load existing configuration if a component is actually installed (env files
# with secrets present), not merely if the ~/.deepfellow directory exists.
if deepfellow_installed; then
    # Load Infra API key
    if [ -f "$DEEPFELLOW_INFRA_ENV_PATH" ]; then
        CFG_INFRA_API_KEY=$(grep -E '^DF_INFRA_API_KEY=' "$DEEPFELLOW_INFRA_ENV_PATH" | cut -d'=' -f2-)
        if [ -n "$CFG_INFRA_API_KEY" ]; then
            print_info "Infra API key loaded from existing installation"
        fi
    fi

    # Load existing app env if available
    if [ -f "$APP_ENV_PATH" ]; then
        DF_SERVER_API_KEY=$(grep -E '^DF_SERVER_API_KEY=' "$APP_ENV_PATH" | cut -d'=' -f2-)
        DF_SERVER_ORGANIZATION_ID=$(grep -E '^DF_SERVER_ORGANIZATION_ID=' "$APP_ENV_PATH" | cut -d'=' -f2-)
        DF_SERVER_PROJECT_ID=$(grep -E '^DF_SERVER_PROJECT_ID=' "$APP_ENV_PATH" | cut -d'=' -f2-)
        ADMIN_EMAIL=$(grep -E '^BOOTSTRAP_ADMIN_EMAIL=' "$APP_ENV_PATH" | cut -d'=' -f2-)
        ADMIN_PASSWORD=$(grep -E '^BOOTSTRAP_ADMIN_PASSWORD=' "$APP_ENV_PATH" | cut -d'=' -f2-)
        ADMIN_NAME=$(grep -E '^BOOTSTRAP_ADMIN_NAME=' "$APP_ENV_PATH" | cut -d'=' -f2-)
        LLM_MODEL=$(grep -E '^SESSION_TITLE_MODEL=' "$APP_ENV_PATH" | cut -d'=' -f2-)
        APP_URL=$(grep -E '^TRUSTED_ORIGINS=' "$APP_ENV_PATH" | cut -d'=' -f2- | tr ',' '\n' | grep -v -E '^http://(127\.0\.0\.1|localhost):' | head -1)

        if [ -n "$DF_SERVER_API_KEY" ] && [ -n "$DF_SERVER_ORGANIZATION_ID" ] && [ -n "$DF_SERVER_PROJECT_ID" ]; then
            EXISTING_ENV_LOADED=true
            print_info "Loaded existing configuration from $APP_ENV_PATH"
        fi
    fi

    print_warning "Existing DeepFellow installation detected at $HOME/.deepfellow"
    if prompt_confirm "Do you want to override the current installation?"; then
        deepfellow server stop >/dev/null 2>&1 || true
        deepfellow infra stop >/dev/null 2>&1 || true

        # Remove DeepFellow-related containers
        docker rm -f server mongo milvus etcd minio infra 2>/dev/null || true

        # Remove DeepFellow-related volumes (silently, log only on error)
        VOLUME_ERRORS=""
        if ! OUTPUT=$(docker volume rm -f server_mongo server_milvus server_etcd server_minio 2>&1); then
            VOLUME_ERRORS="$OUTPUT"
        fi
        for vol in $(docker volume ls -q 2>/dev/null); do
            case "$vol" in
                *mongo*|*milvus*|*etcd*|*minio*|*deepfellow*)
                    if ! OUTPUT=$(docker volume rm -f "$vol" 2>&1); then
                        VOLUME_ERRORS="${VOLUME_ERRORS}\n${OUTPUT}"
                    fi
                    ;;
            esac
        done

        if [ -n "$VOLUME_ERRORS" ]; then
            print_error "Some volumes could not be removed:"
            print_error "$VOLUME_ERRORS"
        fi

        # Remove local env files and reset loaded state
        rm -f "$ENV_PATH" "$APP_ENV_PATH"
        print_info "Removed existing .env and .env.app files"
        EXISTING_ENV_LOADED=false
        unset DF_SERVER_API_KEY DF_SERVER_ORGANIZATION_ID DF_SERVER_PROJECT_ID
        unset ADMIN_EMAIL ADMIN_PASSWORD ADMIN_NAME LLM_MODEL APP_URL
    else
        print_info "Using existing DeepFellow installation"
        if [ -z "$CFG_INFRA_API_KEY" ]; then
            print_error "Could not read Infra API key from existing installation"
            print_end "Try reinstalling or check $DEEPFELLOW_INFRA_ENV_PATH"
            exit 1
        fi
    fi
fi

if ! infra_installed; then
    print_info "Installing DeepFellow"

    if ! install_deepfellow_cli; then
        print_end "DeepFellow CLI installation failed"
        exit 1
    fi

    if ! install_deepfellow_infra; then
        print_end "DeepFellow Infra installation failed"
        exit 1
    fi
fi

if ! setup_docker_environment; then
    print_end "Docker environment setup failed"
    exit 1
fi

if ! install_deepfellow_server; then
    print_end "DeepFellow Server installation failed"
    exit 1
fi

if ! start_deepfellow_server; then
    print_warning "DeepFellow Server did not start automatically — you may need to run 'deepfellow server start' manually"
fi

if [ "$EXISTING_ENV_LOADED" = "false" ]; then
    print_info "Set up the admin account:"
    print_info "(Password must be 10+ chars with lowercase, uppercase, digit, and special character)"

    ADMIN_CREATED=false
    while [ "$ADMIN_CREATED" = "false" ]; do
        prompt_input "Admin name" ADMIN_NAME
        prompt_input "Admin email" ADMIN_EMAIL
        prompt_input --secret "Admin password" ADMIN_PASSWORD

        if run_with_status "Creating admin account..." "deepfellow server create-admin \"$ADMIN_NAME\" \"$ADMIN_EMAIL\" --password \"$ADMIN_PASSWORD\""; then
            ADMIN_CREATED=true
        else
            print_error "Admin creation failed"
            print_error "$RUN_OUTPUT"
            if ! prompt_confirm "Do you want to try again?"; then
                print_end "Installation aborted"
                exit 1
            fi
            print_info "Please re-enter admin credentials:"
        fi
    done

    print_info "Enter the URL where the app will be accessed (e.g., https://myapp.example.com)"
    print_info "Leave empty if you only need localhost access"
    prompt_input_default "App URL" APP_URL ""

    DF_SERVER_ADDR="http://localhost:${CFG_SERVER_PORT}"

    if ! run_with_status "Logging in to DeepFellow Server..." "deepfellow server login --server \"$DF_SERVER_ADDR\" --email \"$ADMIN_EMAIL\" --password \"$ADMIN_PASSWORD\""; then
        prompt_error "Failed to login to DeepFellow Server"
    fi
    print_info "Logged in successfully"

    if ! run_with_status "Creating organization..." "deepfellow server organization create --server \"$DF_SERVER_ADDR\" \"Workspace\" 2>&1"; then
        prompt_error "Failed to create organization"
    fi
    DF_SERVER_ORGANIZATION_ID=$(echo "$RUN_OUTPUT" | grep -oE '[0-9a-f]{24}' | head -1)
    if [ -z "$DF_SERVER_ORGANIZATION_ID" ]; then
        prompt_error "Could not extract organization ID from output"
    fi
    print_info "Organization created: $DF_SERVER_ORGANIZATION_ID"

    if ! run_with_status "Creating project..." "deepfellow server project create --server \"$DF_SERVER_ADDR\" \"$DF_SERVER_ORGANIZATION_ID\" \"Default\" 2>&1"; then
        prompt_error "Failed to create project"
    fi
    DF_SERVER_PROJECT_ID=$(echo "$RUN_OUTPUT" | grep -oE '[0-9a-f]{24}' | head -1)
    if [ -z "$DF_SERVER_PROJECT_ID" ]; then
        prompt_error "Could not extract project ID from output"
    fi
    print_info "Project created: $DF_SERVER_PROJECT_ID"

    if ! run_with_status "Creating project API key..." "deepfellow server project api-key create --server \"$DF_SERVER_ADDR\" \"$DF_SERVER_ORGANIZATION_ID\" \"$DF_SERVER_PROJECT_ID\" \"app\" 2>&1"; then
        prompt_error "Failed to create project API key"
    fi
    DF_SERVER_API_KEY=$(echo "$RUN_OUTPUT" | sed -n 's/^[[:space:]]*value: //p' | head -1)
    if [ -z "$DF_SERVER_API_KEY" ]; then
        prompt_error "Could not extract API key from output"
    fi
    print_info "Project API key created"

    print_info "Installing LLM"

    prompt_input_default "Model name" LLM_MODEL "gemma4:e4b"

    if ! run_with_status "Installing Ollama service..." "deepfellow infra service install ollama --spec '{\"hardware\":\"GPU\", \"keep_alive\":\"-1\", \"is_flash_attention\":true,\"context_length\":250000}'"; then
        prompt_error "Failed to install Ollama service"
        exit 1
    fi
    print_info "Ollama service installed successfully"

    if ! run_with_status "Installing model ${LLM_MODEL}..." "deepfellow infra model install ollama ${LLM_MODEL}"; then
        print_warning "Model installation failed, retrying in 5 seconds..."
        sleep 5
        if ! run_with_status "Retrying model ${LLM_MODEL} installation..." "deepfellow infra model install ollama ${LLM_MODEL}"; then
            prompt_error "Failed to install model ${LLM_MODEL}"
        fi
    fi
    print_info "Model ${LLM_MODEL} installed successfully"

    if ! run_with_status "Installing embedding model ${CFG_EMBEDDING_MODEL}..." "deepfellow infra model install ollama ${CFG_EMBEDDING_MODEL}"; then
        print_warning "Embedding model installation failed, retrying in 5 seconds..."
        sleep 5
        if ! run_with_status "Retrying embedding model ${CFG_EMBEDDING_MODEL} installation..." "deepfellow infra model install ollama ${CFG_EMBEDDING_MODEL}"; then
            prompt_error "Failed to install embedding model ${CFG_EMBEDDING_MODEL}"
        fi
    fi
    print_info "Embedding model ${CFG_EMBEDDING_MODEL} installed successfully"

    if ! run_with_status "Installing fast model ${CFG_SMALL_MODEL}..." "deepfellow infra model install ollama ${CFG_SMALL_MODEL}"; then
        print_warning "Embedding model installation failed, retrying in 5 seconds..."
        sleep 5
        if ! run_with_status "Retrying fast model ${CFG_SMALL_MODEL} installation..." "deepfellow infra model install ollama ${CFG_SMALL_MODEL}"; then
            prompt_error "Failed to install embedding model ${CFG_SMALL_MODEL}"
        fi
    fi
    print_info "Fast model ${CFG_SMALL_MODEL} installed successfully"

    print_info "DeepFellow installation complete"
else
    print_info "Using existing DeepFellow configuration"
fi


# ------ Step 3: Set up the App container ------

if ! setup_app_mongodb; then
    print_end "App MongoDB setup failed"
    exit 1
fi

if ! generate_app_env; then
    RUN_OUTPUT=""
    prompt_error "Failed to generate app environment file"
fi

print_info "Setting up App container ($APP_IMAGE)"

if ! $DEV_MODE; then
    if ! run_with_status "Pulling app image..." "docker pull $APP_IMAGE"; then
        RUN_OUTPUT=""
        prompt_error "Failed to pull app image $APP_IMAGE"
    fi
fi

if ! setup_app_container; then
    RUN_OUTPUT=""
    prompt_error "Deepfellow Workspace container setup failed"
fi

# Save a local copy of the env for reference
if [ -f "$DEEPFELLOW_ENV_PATH" ]; then
    cp "$DEEPFELLOW_ENV_PATH" "$ENV_PATH"
    {
        echo ""
        echo "# --- Deepfellow Workspace settings ---"
        echo "APP_IMAGE=$APP_IMAGE"
        echo "APP_CONTAINER_NAME=$APP_CONTAINER_NAME"
    } >> "$ENV_PATH"
    print_info "Full environment saved to $ENV_PATH"
fi

print_info "DeepFellow and Deepfellow Workspace are running."
print_info ""
print_info "Note: To use 'deepfellow' command in future terminal sessions, add this to your shell profile:"
print_info "  export PATH=\"\$HOME/.local/bin:\$PATH\""

if [ -z "$APP_URL" ]; then
    print_info ""
    print_warning "No app URL was specified during installation."
    print_warning "If you need to access the app from a custom URL, you must:"
    print_warning "  1. Add your URL to TRUSTED_ORIGINS in $APP_ENV_PATH"
    print_warning "  2. Recreate the docker container: docker rm -f $APP_CONTAINER_NAME && ./setup.sh"
fi

print_end "Setup complete!"