#!/bin/bash

# Docker Management Script for VectorChat Playwright Container
# Usage: ./docker-manage.sh [command] [options]

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
IMAGE_NAME="vectorchat"
IMAGE_TAG="playwright-latest"
CONTAINER_NAME="vectorchat-daemon"
COMPOSE_FILE="docker-compose.playwright.yml"
DOCKERFILE="Dockerfile.playwright"

# Helper functions
print_header() {
    echo -e "${BLUE}================================${NC}"
    echo -e "${BLUE}$1${NC}"
    echo -e "${BLUE}================================${NC}"
}

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

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

print_warning() {
    echo -e "${YELLOW}⚠️  $1${NC}"
}

# GPU Detection and Backend Selection
detect_gpu_backend() {
    print_header "Detecting GPU Backend"

    # Check for NVIDIA GPU
    if command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi >/dev/null 2>&1; then
        print_success "NVIDIA GPU detected"
        echo "cuda"
        return
    fi

    # Check for AMD GPU (ROCm)
    if command -v rocm-smi >/dev/null 2>&1 && rocm-smi >/dev/null 2>&1; then
        print_success "AMD GPU (ROCm) detected"
        echo "rocm"
        return
    fi

    # Check for Intel GPU (Level Zero)
    if command -v intel-gpu-top >/dev/null 2>&1 || [ -d /sys/class/drm/card0/device/driver/module/drivers/i915 ]; then
        print_success "Intel GPU detected"
        echo "intel"
        return
    fi

    print_warning "No GPU detected, using CPU backend"
    echo "cpu"
}

check_gpu_runtime() {
    BACKEND="$1"

    case "$BACKEND" in
        cuda)
            if ! docker info | grep -q "Runtimes.*nvidia"; then
                print_warning "NVIDIA runtime not available. Install nvidia-container-toolkit."
                print_warning "Run: curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg"
                print_warning "Run: curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list"
                print_warning "Run: sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit"
                print_warning "Run: sudo nvidia-ctk runtime configure --runtime=docker"
                print_warning "Run: sudo systemctl restart docker"
                return 1
            fi
            ;;
        rocm)
            print_warning "ROCm runtime support requires AMD's ROCm Docker base image or manual device mounting."
            ;;
        intel)
            print_warning "Intel GPU runtime support requires oneAPI and device mounting."
            ;;
    esac
    return 0
}

get_gpu_docker_args() {
    BACKEND="$1"

    case "$BACKEND" in
        cuda)
            echo "--gpus all"
            ;;
        rocm)
            echo "--device=/dev/kfd --device=/dev/dri --group-add video"
            ;;
        intel)
            echo "--device=/dev/dri --group-add video"
            ;;
        cpu)
            echo ""
            ;;
    esac
# Multi-Backend GPU Support Commands
build_cuda() {
    print_header "Building CUDA Docker Image"

    if [ "$1" == "--no-cache" ]; then
        docker build --no-cache -f Dockerfile.cuda -t "$IMAGE_NAME-cuda:$IMAGE_TAG" .
    else
        docker build -f Dockerfile.cuda -t "$IMAGE_NAME-cuda:$IMAGE_TAG" .
    fi

    print_success "CUDA image built: $IMAGE_NAME-cuda:$IMAGE_TAG"
}

build_rocm() {
    print_header "Building ROCm Docker Image"

    if [ "$1" == "--no-cache" ]; then
        docker build --no-cache -f Dockerfile.rocm -t "$IMAGE_NAME-rocm:$IMAGE_TAG" .
    else
        docker build -f Dockerfile.rocm -t "$IMAGE_NAME-rocm:$IMAGE_TAG" .
    fi

    print_success "ROCm image built: $IMAGE_NAME-rocm:$IMAGE_TAG"
}

build_intel() {
    print_header "Building Intel GPU Docker Image"

    if [ "$1" == "--no-cache" ]; then
        docker build --no-cache -f Dockerfile.intel -t "$IMAGE_NAME-intel:$IMAGE_TAG" .
    else
        docker build -f Dockerfile.intel -t "$IMAGE_NAME-intel:$IMAGE_TAG" .
    fi

    print_success "Intel GPU image built: $IMAGE_NAME-intel:$IMAGE_TAG"
}

build_all() {
    print_header "Building All GPU Backend Images"

    build_cuda "$1"
    build_rocm "$1"
    build_intel "$1"

    print_success "All GPU backend images built"
}

run_cuda() {
    print_header "Running CUDA Container"

    docker run -d \
        --name "$CONTAINER_NAME-cuda" \
        --gpus all \
        -p 3737:3737 \
        -p 8765:8765 \
        -p 8766:8766 \
        -p 8767:8767 \
        -v ~/.vectorchat:/root/.vectorchat \
        -v vectorchat-logs:/app/logs \
        "$IMAGE_NAME-cuda:$IMAGE_TAG"

    print_success "CUDA container started: $CONTAINER_NAME-cuda"
    print_success "Setup wizard available at: http://localhost:3737"
}

run_rocm() {
    print_header "Running ROCm Container"

    docker run -d \
        --name "$CONTAINER_NAME-rocm" \
        --device=/dev/kfd \
        --device=/dev/dri \
        --group-add video \
        -p 3738:3737 \
        -p 8765:8765 \
        -p 8766:8766 \
        -p 8767:8767 \
        -v ~/.vectorchat:/root/.vectorchat \
        -v vectorchat-logs:/app/logs \
        "$IMAGE_NAME-rocm:$IMAGE_TAG"

    print_success "ROCm container started: $CONTAINER_NAME-rocm"
    print_success "Setup wizard available at: http://localhost:3738"
}

run_intel() {
    print_header "Running Intel GPU Container"

    docker run -d \
        --name "$CONTAINER_NAME-intel" \
        --device=/dev/dri \
        --group-add video \
        -p 3739:3737 \
        -p 8765:8765 \
        -p 8766:8766 \
        -p 8767:8767 \
        -v ~/.vectorchat:/root/.vectorchat \
        -v vectorchat-logs:/app/logs \
        "$IMAGE_NAME-intel:$IMAGE_TAG"

    print_success "Intel GPU container started: $CONTAINER_NAME-intel"
    print_success "Setup wizard available at: http://localhost:3739"
}

test_cuda() {
    print_header "Running CUDA Pytest Suite"

    mkdir -p artifacts/test-reports-cuda

    docker run --rm \
        --name "${CONTAINER_NAME}-test-cuda" \
        --gpus all \
        -v ~/.vectorchat:/root/.vectorchat \
        -v "$(pwd)/artifacts/test-reports-cuda:/app/artifacts/test-reports" \
        "$IMAGE_NAME-cuda:$IMAGE_TAG" \
        pytest --maxfail=1 --disable-warnings \
            --json-report --json-report-file=/app/artifacts/test-reports/report.json \
            --html=/app/artifacts/test-reports/report.html --self-contained-html

    print_success "CUDA pytest suite completed"
    print_success "CUDA reports available in artifacts/test-reports-cuda"
}

test_rocm() {
    print_header "Running ROCm Pytest Suite"

    mkdir -p artifacts/test-reports-rocm

    docker run --rm \
        --name "${CONTAINER_NAME}-test-rocm" \
        --device=/dev/kfd \
        --device=/dev/dri \
        --group-add video \
        -v ~/.vectorchat:/root/.vectorchat \
        -v "$(pwd)/artifacts/test-reports-rocm:/app/artifacts/test-reports" \
        "$IMAGE_NAME-rocm:$IMAGE_TAG" \
        pytest --maxfail=1 --disable-warnings \
            --json-report --json-report-file=/app/artifacts/test-reports/report.json \
            --html=/app/artifacts/test-reports/report.html --self-contained-html

    print_success "ROCm pytest suite completed"
    print_success "ROCm reports available in artifacts/test-reports-rocm"
}

test_intel() {
    print_header "Running Intel GPU Pytest Suite"

    mkdir -p artifacts/test-reports-intel

    docker run --rm \
        --name "${CONTAINER_NAME}-test-intel" \
        --device=/dev/dri \
        --group-add video \
        -v ~/.vectorchat:/root/.vectorchat \
        -v "$(pwd)/artifacts/test-reports-intel:/app/artifacts/test-reports" \
        "$IMAGE_NAME-intel:$IMAGE_TAG" \
        pytest --maxfail=1 --disable-warnings \
            --json-report --json-report-file=/app/artifacts/test-reports/report.json \
            --html=/app/artifacts/test-reports/report.html --self-contained-html

    print_success "Intel GPU pytest suite completed"
    print_success "Intel GPU reports available in artifacts/test-reports-intel"
}

test_all() {
    print_header "Running All GPU Backend Pytest Suites"

    test_cuda
    test_rocm
    test_intel

    print_success "All GPU backend test suites completed"
    print_success "Reports available in artifacts/test-reports-*"
}

# Legacy build function (kept for backward compatibility)
build() {
    print_header "Building Docker Image (Legacy - Auto-detects GPU)"

    # Detect GPU backend
    GPU_BACKEND=$(detect_gpu_backend)

    # Check if runtime is available for non-CPU backends
    if ! check_gpu_runtime "$GPU_BACKEND"; then
        print_warning "Falling back to CPU backend"
        GPU_BACKEND="cpu"
    fi

    print_success "Using GPU backend: $GPU_BACKEND"

    if [ "$1" == "--no-cache" ]; then
        docker build --no-cache --build-arg GPU_BACKEND="$GPU_BACKEND" -f "$DOCKERFILE" -t "$IMAGE_NAME:$IMAGE_TAG" .
    else
        docker build --build-arg GPU_BACKEND="$GPU_BACKEND" -f "$DOCKERFILE" -t "$IMAGE_NAME:$IMAGE_TAG" .
    fi

    print_success "Image built: $IMAGE_NAME:$IMAGE_TAG (GPU backend: $GPU_BACKEND)"
}

run() {
    print_header "Running Container (Legacy - Auto-detects GPU)"

    # Detect GPU backend for runtime args
    GPU_BACKEND=$(detect_gpu_backend)
    GPU_ARGS=$(get_gpu_docker_args "$GPU_BACKEND")

    print_success "Using GPU backend: $GPU_BACKEND"

    docker run -d \
        --name "$CONTAINER_NAME" \
        -p 3737:3737 \
        -p 8765:8765 \
        -p 8766:8766 \
        -p 8767:8767 \
        -v ~/.vectorchat:/root/.vectorchat \
        -v vectorchat-logs:/app/logs \
        $GPU_ARGS \
        "$IMAGE_NAME:$IMAGE_TAG"

    print_success "Container started: $CONTAINER_NAME"
    print_success "Setup wizard available at: http://localhost:3737"
}

# Legacy test function (kept for backward compatibility)
test_suite() {
    print_header "Running Pytest Suite in Docker (Legacy - Auto-detects GPU)"

    mkdir -p artifacts/test-reports

    # Detect GPU backend for runtime args
    GPU_BACKEND=$(detect_gpu_backend)
    GPU_ARGS=$(get_gpu_docker_args "$GPU_BACKEND")

    print_success "Using GPU backend: $GPU_BACKEND"

    docker run --rm \
        --name "${CONTAINER_NAME}-test" \
        -v ~/.vectorchat:/root/.vectorchat \
        -v "$(pwd)/artifacts/test-reports:/app/artifacts/test-reports" \
        $GPU_ARGS \
        "$IMAGE_NAME:$IMAGE_TAG" \
        pytest --maxfail=1 --disable-warnings \
            --json-report --json-report-file=/app/artifacts/test-reports/report.json \
            --html=/app/artifacts/test-reports/report.html --self-contained-html

    print_success "Pytest suite completed"
    print_success "Reports available in artifacts/test-reports"
}

compose_up() {
    print_header "Starting with Docker Compose"
    
    docker-compose -f "$COMPOSE_FILE" up -d
    
    print_success "Services started"
    print_success "Setup wizard available at: http://localhost:3737"
}

compose_down() {
    print_header "Stopping Docker Compose Services"
    
    docker-compose -f "$COMPOSE_FILE" down
    
    print_success "Services stopped"
}

stop() {
    print_header "Stopping Container"
    
    if docker ps -a --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
        docker stop "$CONTAINER_NAME"
        print_success "Container stopped: $CONTAINER_NAME"
    else
        print_warning "Container not found: $CONTAINER_NAME"
    fi
}

start() {
    print_header "Starting Container"
    
    if docker ps -a --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
        docker start "$CONTAINER_NAME"
        print_success "Container started: $CONTAINER_NAME"
    else
        print_error "Container not found: $CONTAINER_NAME"
        exit 1
    fi
}

remove() {
    print_header "Removing Container"
    
    if docker ps -a --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
        docker rm -f "$CONTAINER_NAME"
        print_success "Container removed: $CONTAINER_NAME"
    else
        print_warning "Container not found: $CONTAINER_NAME"
    fi
}

logs() {
    print_header "Container Logs"
    
    if [ -n "$1" ]; then
        docker logs --tail="$1" -f "$CONTAINER_NAME"
    else
        docker logs -f "$CONTAINER_NAME"
    fi
}

shell() {
    print_header "Opening Container Shell"
    
    docker exec -it "$CONTAINER_NAME" /bin/bash
}

status() {
    print_header "Container Status"
    
    if docker ps --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
        print_success "Container is running"
        docker ps --filter "name=$CONTAINER_NAME" --format "table {{.Names}}\t{{.Status}}\t{{.Ports}}"
    else
        print_warning "Container is not running"
        if docker ps -a --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
            docker ps -a --filter "name=$CONTAINER_NAME" --format "table {{.Names}}\t{{.Status}}"
        else
            print_error "Container not found"
        fi
    fi
}

stats() {
    print_header "Container Resource Usage"
    
    docker stats "$CONTAINER_NAME" --no-stream
}

test_suite() {
    print_header "Running Pytest Suite in Docker"

    mkdir -p artifacts/test-reports

    # Detect GPU backend for runtime args
    GPU_BACKEND=$(detect_gpu_backend)
    GPU_ARGS=$(get_gpu_docker_args "$GPU_BACKEND")

    print_success "Using GPU backend: $GPU_BACKEND"

    docker run --rm \
        --name "${CONTAINER_NAME}-test" \
        -v ~/.vectorchat:/root/.vectorchat \
        -v "$(pwd)/artifacts/test-reports:/app/artifacts/test-reports" \
        $GPU_ARGS \
        "$IMAGE_NAME:$IMAGE_TAG" \
        pytest --maxfail=1 --disable-warnings \
            --json-report --json-report-file=/app/artifacts/test-reports/report.json \
            --html=/app/artifacts/test-reports/report.html --self-contained-html

    print_success "Pytest suite completed"
    print_success "Reports available in artifacts/test-reports"
}

health() {
    print_header "Container Health Check"
    
    if docker exec "$CONTAINER_NAME" curl -f http://localhost:3737/health > /dev/null 2>&1; then
        print_success "Container is healthy"
    else
        print_error "Container health check failed"
        exit 1
    fi
}

clean() {
    print_header "Cleaning Up Docker Resources"
    
    print_warning "Removing dangling images..."
    docker image prune -f
    
    print_warning "Removing stopped containers..."
    docker container prune -f
    
    print_warning "Removing unused volumes..."
    docker volume prune -f
    
    print_success "Cleanup complete"
}

push() {
    print_header "Pushing Image to Registry"
    
    if [ -z "$1" ]; then
        print_error "Registry URL required: ./docker-manage.sh push <registry-url>"
        exit 1
    fi
    
    REGISTRY="$1"
    docker tag "$IMAGE_NAME:$IMAGE_TAG" "$REGISTRY/$IMAGE_NAME:$IMAGE_TAG"
    docker push "$REGISTRY/$IMAGE_NAME:$IMAGE_TAG"
    
    print_success "Image pushed to: $REGISTRY/$IMAGE_NAME:$IMAGE_TAG"
}

pull() {
    print_header "Pulling Image from Registry"
    
    if [ -z "$1" ]; then
        print_error "Registry URL required: ./docker-manage.sh pull <registry-url>"
        exit 1
    fi
    
    REGISTRY="$1"
    docker pull "$REGISTRY/$IMAGE_NAME:$IMAGE_TAG"
    docker tag "$REGISTRY/$IMAGE_NAME:$IMAGE_TAG" "$IMAGE_NAME:$IMAGE_TAG"
    
    print_success "Image pulled from: $REGISTRY/$IMAGE_NAME:$IMAGE_TAG"
}

help() {
    cat << EOF
${BLUE}VectorChat Docker Management Script${NC}

${BLUE}Usage:${NC} ./docker-manage.sh [command] [options]

${BLUE}Commands:${NC}
  ${GREEN}build${NC}              Build Docker image (auto-detects GPU backend)
  ${GREEN}build --no-cache${NC}   Build without cache
  ${GREEN}build-cuda${NC}         Build CUDA-specific Docker image
  ${GREEN}build-rocm${NC}         Build ROCm-specific Docker image
  ${GREEN}build-intel${NC}        Build Intel GPU-specific Docker image
  ${GREEN}build-all${NC}          Build all GPU backend images
  ${GREEN}run${NC}                Run container (standalone, with GPU support)
  ${GREEN}run-cuda${NC}           Run CUDA container on port 3737
  ${GREEN}run-rocm${NC}           Run ROCm container on port 3738
  ${GREEN}run-intel${NC}          Run Intel GPU container on port 3739
  ${GREEN}compose-up${NC}         Start with docker-compose
  ${GREEN}compose-down${NC}       Stop docker-compose services
  ${GREEN}start${NC}              Start existing container
  ${GREEN}stop${NC}               Stop running container
  ${GREEN}remove${NC}             Remove container
  ${GREEN}logs${NC} [lines]       View container logs
  ${GREEN}shell${NC}              Open container shell
  ${GREEN}status${NC}             Show container status
  ${GREEN}stats${NC}              Show resource usage
  ${GREEN}test${NC}               Run pytest suite in Docker (with GPU support)
  ${GREEN}test-cuda${NC}          Run CUDA-specific pytest suite
  ${GREEN}test-rocm${NC}          Run ROCm-specific pytest suite
  ${GREEN}test-intel${NC}         Run Intel GPU-specific pytest suite
  ${GREEN}test-all${NC}           Run all GPU backend pytest suites
  ${GREEN}health${NC}             Check container health
  ${GREEN}clean${NC}              Clean up Docker resources
  ${GREEN}push${NC} <registry>    Push image to registry
  ${GREEN}pull${NC} <registry>    Pull image from registry
  ${GREEN}help${NC}               Show this help message

${BLUE}GPU Support:${NC}
  Automatically detects NVIDIA, AMD, Intel GPUs on host
  Builds with appropriate backend (CUDA, ROCm, Intel, or CPU)
  Adds runtime flags for GPU access (--gpus, --device, etc.)
  Test suite runs with GPU support when available

${BLUE}GPU Backend Ports:${NC}
  CUDA:     Port 3737 (primary)
  ROCm:     Port 3738 (AMD)
  Intel:    Port 3739 (Intel GPU)
  Legacy:   Port 3737 (auto-detect)

${BLUE}Examples:${NC}
  ./docker-manage.sh build-cuda
  ./docker-manage.sh run-cuda
  ./docker-manage.sh test-cuda
  ./docker-manage.sh build-all
  ./docker-manage.sh test-all
  ./docker-manage.sh logs 50
  ./docker-manage.sh push docker.io/myrepo
  ./docker-manage.sh health

${BLUE}Container Info:${NC}
  CUDA Image:    $IMAGE_NAME-cuda:$IMAGE_TAG
  ROCm Image:    $IMAGE_NAME-rocm:$IMAGE_TAG
  Intel Image:   $IMAGE_NAME-intel:$IMAGE_TAG
  Legacy Image:  $IMAGE_NAME:$IMAGE_TAG

${BLUE}Setup Wizards:${NC}
  CUDA:  http://localhost:3737
  ROCm:  http://localhost:3738
  Intel: http://localhost:3739

EOF
}

# Main
case "${1:-help}" in
    build)
        build "$2"
        ;;
    run)
        run
        ;;
    compose-up)
        compose_up
        ;;
    compose-down)
        compose_down
        ;;
    start)
        start
        ;;
    stop)
        stop
        ;;
    remove)
        remove
        ;;
    logs)
        logs "$2"
        ;;
    shell)
        shell
        ;;
    status)
        status
        ;;
    stats)
        stats
        ;;
    test)
        test_suite
        ;;
    health)
        health
        ;;
    clean)
        clean
        ;;
    push)
        push "$2"
        ;;
    pull)
        pull "$2"
        ;;
    help)
        help
        ;;
    # Multi-backend GPU commands
    build-cuda)
        build_cuda "$2"
        ;;
    build-rocm)
        build_rocm "$2"
        ;;
    build-intel)
        build_intel "$2"
        ;;
    build-all)
        build_all "$2"
        ;;
    run-cuda)
        run_cuda
        ;;
    run-rocm)
        run_rocm
        ;;
    run-intel)
        run_intel
        ;;
    test-cuda)
        test_cuda
        ;;
    test-rocm)
        test_rocm
        ;;
    test-intel)
        test_intel
        ;;
    test-all)
        test_all
        ;;
    *)
        print_error "Unknown command: $1"
        help
        exit 1
        ;;
esac

}
