#!/bin/bash
# VibeVoice-ASR Setup Script for DGX Spark
# Downloads and builds the VibeVoice-ASR container
#
# Usage:
#   curl -sL https://docs.techswan.online/scripts/vibevoice-asr/setup.sh | bash
#   curl -sL https://docs.techswan.online/scripts/vibevoice-asr/setup.sh | bash -s build
#   curl -sL https://docs.techswan.online/scripts/vibevoice-asr/setup.sh | bash -s serve

set -e

BASE_URL="https://docs.techswan.online/scripts/vibevoice-asr"
INSTALL_DIR="${VIBEVOICE_DIR:-$HOME/vibevoice-asr}"
IMAGE_NAME="vibevoice-asr:dgx-spark"
CONTAINER_NAME="vibevoice-asr"

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

log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_step() {
    echo -e "${BLUE}[STEP]${NC} $1"
}

download_file() {
    local url="$1"
    local dest="$2"
    local dir=$(dirname "$dest")
    
    mkdir -p "$dir"
    
    if command -v curl &> /dev/null; then
        curl -sL "$url" -o "$dest"
    elif command -v wget &> /dev/null; then
        wget -q "$url" -O "$dest"
    else
        log_error "curl or wget is required"
        exit 1
    fi
}

download_files() {
    log_step "Downloading VibeVoice-ASR files to $INSTALL_DIR..."
    
    mkdir -p "$INSTALL_DIR"
    cd "$INSTALL_DIR"
    
    # Core files
    local files=(
        "Dockerfile"
        "requirements-realtime.txt"
        "test_vibevoice.py"
        "vibevoice_asr_gradio_demo_patched.py"
        "run_realtime.sh"
        "run_all.sh"
    )
    
    for file in "${files[@]}"; do
        log_info "Downloading $file..."
        download_file "$BASE_URL/$file" "$INSTALL_DIR/$file"
    done
    
    # Realtime module
    local realtime_files=(
        "__init__.py"
        "models.py"
        "server.py"
        "asr_worker.py"
        "session_manager.py"
        "audio_buffer.py"
        "vad_processor.py"
    )
    
    mkdir -p "$INSTALL_DIR/realtime"
    for file in "${realtime_files[@]}"; do
        log_info "Downloading realtime/$file..."
        download_file "$BASE_URL/realtime/$file" "$INSTALL_DIR/realtime/$file"
    done
    
    # Static files
    mkdir -p "$INSTALL_DIR/static"
    log_info "Downloading static/realtime_client.html..."
    download_file "$BASE_URL/static/realtime_client.html" "$INSTALL_DIR/static/realtime_client.html"
    
    # Make scripts executable
    chmod +x "$INSTALL_DIR/run_realtime.sh" "$INSTALL_DIR/run_all.sh"
    
    log_info "All files downloaded to $INSTALL_DIR"
}

check_prerequisites() {
    log_step "Checking prerequisites..."

    # Check Docker
    if ! command -v docker &> /dev/null; then
        log_error "Docker is not installed"
        exit 1
    fi

    # Check NVIDIA Docker runtime
    if ! docker info 2>/dev/null | grep -q "Runtimes.*nvidia"; then
        log_warn "NVIDIA Docker runtime may not be configured"
    fi

    # Check GPU availability
    if command -v nvidia-smi &> /dev/null; then
        log_info "GPU detected:"
        nvidia-smi --query-gpu=name,memory.total --format=csv,noheader
    else
        log_warn "nvidia-smi not found on host"
    fi

    log_info "Prerequisites check complete"
}

build_image() {
    log_step "Building Docker image: ${IMAGE_NAME}"
    log_info "This may take several minutes..."

    cd "$INSTALL_DIR"

    docker build \
        --network=host \
        -t "$IMAGE_NAME" \
        -f Dockerfile \
        .

    log_info "Docker image built successfully: ${IMAGE_NAME}"
}

run_container() {
    local mode="${1:-interactive}"

    log_step "Running container in ${mode} mode..."

    # Stop existing container if running
    if docker ps -q -f name="$CONTAINER_NAME" | grep -q .; then
        log_warn "Stopping existing container..."
        docker stop "$CONTAINER_NAME" 2>/dev/null || true
    fi

    # Remove existing container
    docker rm "$CONTAINER_NAME" 2>/dev/null || true

    # Common Docker options for DGX Spark
    local docker_opts=(
        --gpus all
        --ipc=host
        --network=host
        --ulimit memlock=-1:-1
        --ulimit stack=-1:-1
        -e "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True"
        -v "$HOME/.cache/huggingface:/root/.cache/huggingface"
        --name "$CONTAINER_NAME"
    )

    if [ "$mode" = "interactive" ]; then
        docker run --rm -it "${docker_opts[@]}" "$IMAGE_NAME" bash
    elif [ "$mode" = "test" ]; then
        docker run --rm "${docker_opts[@]}" "$IMAGE_NAME" python /workspace/test_vibevoice.py
    elif [ "$mode" = "demo" ]; then
        log_info "Starting Gradio demo on port 7860..."
        log_info "Access the demo at: http://localhost:7860"
        docker run --rm -it "${docker_opts[@]}" "$IMAGE_NAME"
    elif [ "$mode" = "realtime" ]; then
        log_info "Starting Realtime ASR server on port 8000..."
        log_info "WebSocket API: ws://localhost:8000/ws/asr/{session_id}"
        log_info "Test client:   http://localhost:8000/static/realtime_client.html"
        docker run --rm -it "${docker_opts[@]}" "$IMAGE_NAME" \
            python -m realtime.server --host 0.0.0.0 --port 8000
    elif [ "$mode" = "serve" ]; then
        log_info "Starting all services..."
        log_info "  Gradio demo:    http://localhost:7860"
        log_info "  Realtime ASR:   http://localhost:8000"
        log_info "  Test client:    http://localhost:8000/static/realtime_client.html"
        docker run --rm -it "${docker_opts[@]}" "$IMAGE_NAME" ./run_all.sh
    else
        log_error "Unknown mode: $mode"
        exit 1
    fi
}

show_usage() {
    echo "VibeVoice-ASR Setup for DGX Spark"
    echo ""
    echo "Usage:"
    echo "  curl -sL $BASE_URL/setup.sh | bash              # Download only"
    echo "  curl -sL $BASE_URL/setup.sh | bash -s build     # Download and build"
    echo "  curl -sL $BASE_URL/setup.sh | bash -s demo      # Download, build, run demo"
    echo "  curl -sL $BASE_URL/setup.sh | bash -s serve     # Download, build, run all"
    echo ""
    echo "Commands:"
    echo "  (default)   Download files only"
    echo "  build       Download and build Docker image"
    echo "  demo        Download, build, and start Gradio demo (port 7860)"
    echo "  realtime    Download, build, and start Realtime ASR (port 8000)"
    echo "  serve       Download, build, and start both services"
    echo "  run         Run container interactively (after build)"
    echo ""
    echo "Environment variables:"
    echo "  VIBEVOICE_DIR   Installation directory (default: ~/vibevoice-asr)"
    echo ""
    echo "After installation, you can also run:"
    echo "  cd ~/vibevoice-asr"
    echo "  docker run --gpus all -p 7860:7860 vibevoice-asr:dgx-spark"
}

main() {
    local command="${1:-download}"

    echo ""
    echo "=========================================="
    echo "  VibeVoice-ASR Setup for DGX Spark"
    echo "=========================================="
    echo ""

    case "$command" in
        download)
            download_files
            echo ""
            log_info "Done! Next steps:"
            echo "  cd $INSTALL_DIR"
            echo "  docker build -t vibevoice-asr:dgx-spark ."
            echo "  docker run --gpus all -p 7860:7860 vibevoice-asr:dgx-spark"
            ;;
        build)
            download_files
            check_prerequisites
            build_image
            echo ""
            log_info "Done! To run:"
            echo "  cd $INSTALL_DIR"
            echo "  docker run --gpus all -p 7860:7860 vibevoice-asr:dgx-spark"
            ;;
        run)
            cd "$INSTALL_DIR" 2>/dev/null || { log_error "Run 'build' first"; exit 1; }
            run_container interactive
            ;;
        test)
            cd "$INSTALL_DIR" 2>/dev/null || { log_error "Run 'build' first"; exit 1; }
            run_container test
            ;;
        demo)
            download_files
            check_prerequisites
            build_image
            run_container demo
            ;;
        realtime)
            download_files
            check_prerequisites
            build_image
            run_container realtime
            ;;
        serve)
            download_files
            check_prerequisites
            build_image
            run_container serve
            ;;
        -h|--help|help)
            show_usage
            ;;
        *)
            log_error "Unknown command: $command"
            show_usage
            exit 1
            ;;
    esac
}

main "$@"