#!/usr/bin/env bash
# ACE-Step Gradio Web UI Launcher - Linux (CUDA)
# This script launches the Gradio web interface for ACE-Step

set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# ==================== Load .env Configuration ====================
# Load settings from .env file if it exists
_load_env_file() {
    local env_file="${SCRIPT_DIR}/.env"
    if [[ ! -f "$env_file" ]]; then
        return 0
    fi
    
    echo "[Config] Loading configuration from .env file..."
    
    # Read .env file and export variables
    while IFS='=' read -r key value || [[ -n "$key" ]]; do
        # Skip empty lines and comments
        [[ -z "$key" || "$key" =~ ^[[:space:]]*# ]] && continue
        
        # Trim whitespace from key and value
        key="${key#"${key%%[![:space:]]*}"}"
        key="${key%"${key##*[![:space:]]}"}"
        value="${value#"${value%%[![:space:]]*}"}"
        value="${value%"${value##*[![:space:]]}"}"
        
        # Map .env variable names to script variables
        case "$key" in
            ACESTEP_CONFIG_PATH)
                [[ -n "$value" ]] && CONFIG_PATH="--config_path $value"
                ;;
            ACESTEP_LM_MODEL_PATH)
                [[ -n "$value" ]] && LM_MODEL_PATH="--lm_model_path $value"
                ;;
            ACESTEP_INIT_LLM)
                if [[ -n "$value" && "$value" != "auto" ]]; then
                    INIT_LLM="--init_llm $value"
                fi
                ;;
            ACESTEP_DOWNLOAD_SOURCE)
                if [[ -n "$value" && "$value" != "auto" ]]; then
                    DOWNLOAD_SOURCE="--download-source $value"
                fi
                ;;
            ACESTEP_API_KEY)
                [[ -n "$value" ]] && API_KEY="--api-key $value"
                ;;
            PORT)
                [[ -n "$value" ]] && PORT="$value"
                ;;
            SERVER_NAME)
                [[ -n "$value" ]] && SERVER_NAME="$value"
                ;;
            LANGUAGE)
                [[ -n "$value" ]] && LANGUAGE="$value"
                ;;
            ACESTEP_BATCH_SIZE)
                [[ -n "$value" ]] && BATCH_SIZE="--batch_size $value"
                ;;
        esac
    done < "$env_file"
    
    echo "[Config] Configuration loaded from .env"
}

_load_env_file

# ==================== Manual Startup Configuration ====================
# Ask for manual settings
_load_manual() {
    echo -e "\033[1;36m"
    echo "======================================================"
    echo "            ACE-Step Manual Launch Mode"
    echo "======================================================"
    echo

    while true; do
        read -rp "Continue with manual configuration? (Y/N): " MANUAL_CHOICE
        case "$MANUAL_CHOICE" in
            [Yy])
                break ;;
            [Nn])
                echo -e "\033[0m"
                echo "Proceeding with automatic configuration..."
                echo
                return 0 ;;
            *)
                echo "Invalid input. Please enter Y or N." ;;
        esac
    done

    echo
    echo "-------------------- Update Settings --------------------"
    while true; do
        read -rp "Check for updates before launch? (Y/N): " UPDATE_CHOICE
        case "$UPDATE_CHOICE" in
            [Yy])
                CHECK_UPDATE="true"
                break ;;
            [Nn])
                CHECK_UPDATE="false"
                break ;;
            *)
                echo "Invalid input. Please enter Y or N." ;;
        esac
    done

    echo
    echo "-------------------- Select DiT Model --------------------"
    echo "1) acestep-v15-base"
    echo "2) acestep-v15-sft"
    echo "3) acestep-v15-turbo (Recommended)"
    echo "4) acestep-v15-turbo-rl"
    echo
    while true; do
        read -rp "Enter selection (1-4): " DIT_CHOICE
        case "$DIT_CHOICE" in
            1)
                CONFIG_PATH="--config_path acestep-v15-base"
                break ;;
            2)
                CONFIG_PATH="--config_path acestep-v15-sft"
                break ;;
            3)
                CONFIG_PATH="--config_path acestep-v15-turbo"
                break ;;
            4)
                CONFIG_PATH="--config_path acestep-v15-turbo-rl"
                break ;;
            *)
                echo "Invalid input. Please enter a number between 1 and 4." ;;
        esac
    done

    echo
    echo "-------------------- Select LM Model --------------------"
    echo "1) acestep-5Hz-lm-0.6B (Recommended)"
    echo "2) acestep-5Hz-lm-1.7B"
    echo "3) acestep-5Hz-lm-4B"
    echo "4) Launch without LM Model"
    echo
    while true; do
        read -rp "Enter selection (1-4): " LM_CHOICE
        case "$LM_CHOICE" in
            1)
                LM_MODEL_PATH="--lm_model_path acestep-5Hz-lm-0.6B"
                INIT_LLM="--init_llm true"
                break ;;
            2)
                LM_MODEL_PATH="--lm_model_path acestep-5Hz-lm-1.7B"
                INIT_LLM="--init_llm true"
                break ;;
            3)
                LM_MODEL_PATH="--lm_model_path acestep-5Hz-lm-4B"
                INIT_LLM="--init_llm true"
                break ;;
            4)
                LM_MODEL_PATH=""
                INIT_LLM="--init_llm false"
                break ;;
            *)
                echo "Invalid input. Please enter a number between 1 and 4." ;;
        esac
    done

    echo
    echo "-------------------- CPU Offload Option --------------------"
    while true; do
        read -rp "Enable CPU Offload? (Y/N): " OFFLOAD_CHOICE
        case "$OFFLOAD_CHOICE" in
            [Yy])
                OFFLOAD_TO_CPU="--offload_to_cpu true"
                break ;;
            [Nn])
                OFFLOAD_TO_CPU="--offload_to_cpu false"
                break ;;
            *)
                echo "Invalid input. Please enter Y or N." ;;
        esac
    done

    echo -e "\033[0m"
    echo "Manual configuration applied successfully."
    echo
}

_load_manual

# ==================== Configuration ====================
# Default values (used if not set in .env file)
# You can override these by uncommenting and modifying the lines below
# or by creating a .env file (recommended to survive updates)

# Server settings
: "${PORT:=7860}"
: "${SERVER_NAME:=127.0.0.1}"
# SERVER_NAME="0.0.0.0"
SHARE="${SHARE:-}"
# SHARE="--share"

# Reset LANGUAGE if it contains an invalid value (e.g. system locale like en_CA:en)
case "$LANGUAGE" in
    en|zh|he|ja) ;;
    *) unset LANGUAGE ;;
esac
# UI language: en, zh, he, ja
: "${LANGUAGE:=en}"

# Batch size: default batch size for generation (1 to GPU-dependent max)
# When not specified, defaults to min(2, GPU_max)
BATCH_SIZE="${BATCH_SIZE:-}"
# BATCH_SIZE="--batch_size 4"

# Model settings
: "${CONFIG_PATH:=--config_path acestep-v15-turbo}"
: "${LM_MODEL_PATH:=--lm_model_path acestep-5Hz-lm-0.6B}"
# OFFLOAD_TO_CPU="--offload_to_cpu true"
OFFLOAD_TO_CPU="${OFFLOAD_TO_CPU:-}"

# LLM (Language Model) initialization settings
# By default, LLM is auto-enabled/disabled based on GPU VRAM:
#   - <=6GB VRAM: LLM disabled (DiT-only mode)
#   - >6GB VRAM: LLM enabled
# Values: auto (default), true (force enable), false (force disable)
INIT_LLM="${INIT_LLM:-}"
# INIT_LLM="--init_llm auto"
# INIT_LLM="--init_llm true"
# INIT_LLM="--init_llm false"

# Download source settings
# Preferred download source: auto (default), huggingface, or modelscope
DOWNLOAD_SOURCE="${DOWNLOAD_SOURCE:-}"
# DOWNLOAD_SOURCE="--download-source modelscope"
# DOWNLOAD_SOURCE="--download-source huggingface"

# Update check on startup (set to "false" to disable)
: "${CHECK_UPDATE:=true}"
# CHECK_UPDATE="false"

# Auto-initialize models on startup
: "${INIT_SERVICE:=--init_service true}"

# API settings (enable REST API alongside Gradio)
ENABLE_API="${ENABLE_API:-}"
# ENABLE_API="--enable-api"
API_KEY="${API_KEY:-}"
# API_KEY="--api-key sk-your-secret-key"

# Authentication settings
AUTH_USERNAME="${AUTH_USERNAME:-}"
# AUTH_USERNAME="--auth-username admin"
AUTH_PASSWORD="${AUTH_PASSWORD:-}"
# AUTH_PASSWORD="--auth-password password"

# ==================== Launch ====================

# ==================== Startup Update Check ====================
_startup_update_check() {
    [[ "$CHECK_UPDATE" != "true" ]] && return 0
    command -v git &>/dev/null || return 0
    cd "$SCRIPT_DIR" || return 0
    git rev-parse --git-dir &>/dev/null 2>&1 || return 0

    local branch commit remote_commit
    branch="$(git rev-parse --abbrev-ref HEAD 2>/dev/null || echo "main")"
    commit="$(git rev-parse --short HEAD 2>/dev/null || echo "")"
    [[ -z "$commit" ]] && return 0

    echo "[Update] Checking for updates..."

    # Fetch with timeout (10s)
    local fetch_ok=0
    if command -v timeout &>/dev/null; then
        timeout 10 git fetch origin --quiet 2>/dev/null && fetch_ok=1
    elif command -v gtimeout &>/dev/null; then
        gtimeout 10 git fetch origin --quiet 2>/dev/null && fetch_ok=1
    else
        git fetch origin --quiet 2>/dev/null && fetch_ok=1
    fi

    if [[ $fetch_ok -eq 0 ]]; then
        echo "[Update] Network unreachable, skipping."
        echo
        return 0
    fi

    remote_commit="$(git rev-parse --short "origin/$branch" 2>/dev/null || echo "")"

    if [[ -z "$remote_commit" || "$commit" == "$remote_commit" ]]; then
        echo "[Update] Already up to date ($commit)."
        echo
        return 0
    fi

    echo
    echo "========================================"
    echo "  Update available!"
    echo "========================================"
    echo "  Current: $commit  ->  Latest: $remote_commit"
    echo
    echo "  Recent changes:"
    git --no-pager log --oneline "HEAD..origin/$branch" 2>/dev/null | head -10
    echo

    read -rp "Update now before starting? (Y/N): " update_choice
    if [[ "${update_choice^^}" == "Y" ]]; then
        if [[ -f "$SCRIPT_DIR/check_update.sh" ]]; then
            bash "$SCRIPT_DIR/check_update.sh"
        else
            echo "Pulling latest changes..."
            git pull --ff-only origin "$branch" 2>/dev/null || {
                echo "[Update] Update failed. Please run: git pull"
            }
        fi
    else
        echo "[Update] Skipped. Run ./check_update.sh to update later."
    fi
    echo
}
_startup_update_check

echo "Starting ACE-Step Gradio Web UI..."
echo "Server will be available at: http://${SERVER_NAME}:${PORT}"
echo

# ==================== Standard uv Workflow ====================
# Works on all platforms: x86_64 Linux (cu128), aarch64 Linux/DGX Spark (cu130),
# macOS (MPS), Windows (cu128). uv resolves the correct PyTorch wheels via
# platform-specific index mappings in pyproject.toml.

# Check if uv is installed
if ! command -v uv &>/dev/null; then
    # Try common install locations
    if [[ -x "$HOME/.local/bin/uv" ]]; then
        export PATH="$HOME/.local/bin:$PATH"
    elif [[ -x "$HOME/.cargo/bin/uv" ]]; then
        export PATH="$HOME/.cargo/bin:$PATH"
    fi
fi

if ! command -v uv &>/dev/null; then
    echo
    echo "========================================"
    echo "uv package manager not found!"
    echo "========================================"
    echo
    echo "ACE-Step requires the uv package manager."
    echo
    read -rp "Install uv now? (Y/N): " INSTALL_UV

    if [[ "${INSTALL_UV^^}" == "Y" ]]; then
        echo
        bash "$SCRIPT_DIR/install_uv.sh" --silent
        INSTALL_RESULT=$?

        if [[ $INSTALL_RESULT -eq 0 ]]; then
            echo
            echo "========================================"
            echo "uv installed successfully!"
            echo "========================================"
            echo

            export PATH="$HOME/.local/bin:$HOME/.cargo/bin:$PATH"

            if command -v uv &>/dev/null; then
                echo "uv is now available!"
                uv --version
                echo
            else
                echo
                echo "uv installed but not in PATH yet."
                echo "Please restart your terminal or run:"
                echo "  export PATH=\"\$HOME/.local/bin:\$PATH\""
                echo
                exit 1
            fi
        else
            echo
            echo "========================================"
            echo "Installation failed!"
            echo "========================================"
            echo
            echo "Please install uv manually:"
            echo "  curl -LsSf https://astral.sh/uv/install.sh | sh"
            echo
            exit 1
        fi
    else
        echo
        echo "Installation cancelled."
        echo
        echo "To use ACE-Step, please install uv:"
        echo "  curl -LsSf https://astral.sh/uv/install.sh | sh"
        echo
        exit 1
    fi
fi

echo "[Environment] Using uv package manager..."
echo

# Check if virtual environment exists
if [[ ! -d "$SCRIPT_DIR/.venv" ]]; then
    echo "[Setup] Virtual environment not found. Setting up environment..."
    echo "This will take a few minutes on first run."
    echo
    echo "Running: uv sync"
    echo

    if ! (cd "$SCRIPT_DIR" && uv sync); then
        echo
        echo "[Retry] Online sync failed, retrying in offline mode..."
        echo
        if ! (cd "$SCRIPT_DIR" && uv sync --offline); then
            echo
            echo "========================================"
            echo "[Error] Failed to setup environment"
            echo "========================================"
            echo
            echo "Both online and offline modes failed."
            echo "Please check:"
            echo "  1. Your internet connection (required for first-time setup)"
            echo "  2. Ensure you have enough disk space"
            echo "  3. Try running: uv sync manually"
            exit 1
        fi
    fi

    echo
    echo "========================================"
    echo "Environment setup completed!"
    echo "========================================"
    echo
fi

echo "Starting ACE-Step Gradio UI..."
echo

# Build command with optional parameters
ACESTEP_ARGS="acestep --port $PORT --server-name $SERVER_NAME --language $LANGUAGE"
[[ -n "$SHARE" ]] && ACESTEP_ARGS="$ACESTEP_ARGS $SHARE"
[[ -n "$CONFIG_PATH" ]] && ACESTEP_ARGS="$ACESTEP_ARGS $CONFIG_PATH"
[[ -n "$LM_MODEL_PATH" ]] && ACESTEP_ARGS="$ACESTEP_ARGS $LM_MODEL_PATH"
[[ -n "$OFFLOAD_TO_CPU" ]] && ACESTEP_ARGS="$ACESTEP_ARGS $OFFLOAD_TO_CPU"
[[ -n "$INIT_LLM" ]] && ACESTEP_ARGS="$ACESTEP_ARGS $INIT_LLM"
[[ -n "$DOWNLOAD_SOURCE" ]] && ACESTEP_ARGS="$ACESTEP_ARGS $DOWNLOAD_SOURCE"
[[ -n "$INIT_SERVICE" ]] && ACESTEP_ARGS="$ACESTEP_ARGS $INIT_SERVICE"
[[ -n "$BATCH_SIZE" ]] && ACESTEP_ARGS="$ACESTEP_ARGS $BATCH_SIZE"
[[ -n "$ENABLE_API" ]] && ACESTEP_ARGS="$ACESTEP_ARGS $ENABLE_API"
[[ -n "$API_KEY" ]] && ACESTEP_ARGS="$ACESTEP_ARGS $API_KEY"
[[ -n "$AUTH_USERNAME" ]] && ACESTEP_ARGS="$ACESTEP_ARGS $AUTH_USERNAME"
[[ -n "$AUTH_PASSWORD" ]] && ACESTEP_ARGS="$ACESTEP_ARGS $AUTH_PASSWORD"

cd "$SCRIPT_DIR" && uv run $ACESTEP_ARGS || {
    echo
    echo "[Retry] Online dependency resolution failed, retrying in offline mode..."
    echo
    uv run --offline $ACESTEP_ARGS || {
        echo
        echo "========================================"
        echo "[Error] Failed to start ACE-Step"
        echo "========================================"
        echo
        echo "Both online and offline modes failed."
        echo "Please check:"
        echo "  1. Your internet connection (for first-time setup)"
        echo "  2. If dependencies were previously installed (offline mode requires a prior successful install)"
        echo "  3. Try running: uv sync --offline"
        exit 1
    }
}