# Stage 1: Pull Python from the official Python Docker image
FROM python:3.9-slim AS python-stage

# OPTIONAL: Stage 2: Pull CUDA 11.8
FROM nvidia/cuda:11.8.0-base-ubuntu22.04 AS cuda-stage

# Stage 3: Start from windows-local base image
FROM windowsarena/windows-local:latest

# Define build argument for profiling mode (default is off)
ARG PROFILE_MODE=false

# Install necessary packages and clean up in one layer to reduce image size
RUN apt-get update && apt-get install -y \
    dos2unix \
    libgl1 \
    libevdev-dev \
    python3-dev \
    build-essential \
    tesseract-ocr \
    && rm -rf /var/lib/apt/lists/*

# OPTIONAL: Copy CUDA libraries and binaries from CUDA stage
COPY --from=cuda-stage /usr/local/cuda-11.8 /usr/local/cuda-11.8
ENV PATH="/usr/local/cuda-11.8/bin:${PATH}"
ENV LD_LIBRARY_PATH="/usr/local/cuda-11.8/lib64:${LD_LIBRARY_PATH}"

# Copy Python binaries and libraries from the Python stage
COPY --from=python-stage /usr/local /usr/local

# Ensure scripts are callable
RUN ln -s /usr/local/bin/python /usr/bin/python && \
    ln -s /usr/local/bin/pip /usr/bin/pip

# Copy client application and requirements
COPY src/win-arena-container/client/requirements.txt /client/requirements.txt

# Install Python dependencies from requirements.txt
# Combine all Python package installations into one RUN statement to minimize layering
RUN pip install --no-cache-dir -r /client/requirements.txt && \
    pip install --no-cache-dir requests easyocr tenacity onnxruntime pyOpenSSL sentencepiece anytree

# Download model weights
RUN mkdir -p /models/groundingDINO/ && \
    wget -q https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth -O /models/groundingDINO/groundingdino_swint_ogc.pth && \
    apt-get update && apt-get install -y git git-lfs && git lfs install && \
    git clone https://huggingface.co/microsoft/OmniParser /models/omni

# Conditionally run the disk usage profiling if PROFILE_MODE is true
RUN if [ "$PROFILE_MODE" = "true" ]; then \
        du -ah / | sort -rh | head -n 50 > /du_output.txt; \
    fi