FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04 AS base
RUN apt-get update && \
    apt-get upgrade -y && \
    apt install -y \
    python3-pip \
    git \
    jq \
    python3-venv \
    wget \
    gnupg \
    vim \
    curl

RUN wget -q -O - https://dl-ssl.google.com/linux/linux_signing_key.pub | apt-key add -
RUN echo "deb http://dl.google.com/linux/chrome/deb/ stable main" >> /etc/apt/sources.list.d/google-chrome.list
RUN apt-get update && \
    apt-get install -y \
    google-chrome-stable

RUN rm -rf /var/lib/apt/lists/*

RUN curl https://sh.rustup.rs -sSf | bash -s -- -y
ENV PATH="/root/.cargo/bin:${PATH}"

# Create directory structure
RUN mkdir -p \
    /root/sia \
    /root/sia/scripts \
    /root/data/iterations \
    /root/data/user \
    /root/data/tasks \
    /root/data/environment \
    /root/models \
    /root/desktop \
    /root/venvs

# ITB tool setup
FROM base AS itb-env
COPY ./scripts/setup_binaries.py /root/sia/scripts/
COPY ./tools/itb/setup.py /root/sia/tools/itb/setup.py
RUN python3 /root/sia/scripts/setup_binaries.py /root/sia/tools/itb/setup.py
RUN python3 -m venv /root/venvs/itb
RUN --mount=type=cache,target=/root/.cache/pip \
    /root/venvs/itb/bin/pip install -e /root/sia/tools/itb/

# Train tool setup
FROM base AS train-env
COPY ./scripts/setup_binaries.py /root/sia/scripts/
COPY ./tools/train/setup.py /root/sia/tools/train/setup.py
RUN python3 /root/sia/scripts/setup_binaries.py /root/sia/tools/train/setup.py
RUN python3 -m venv /root/venvs/train
RUN --mount=type=cache,target=/root/.cache/pip \
    /root/venvs/train/bin/pip install -e /root/sia/tools/train/

# SIA core setup
FROM base AS sia-env
COPY ./scripts/setup_binaries.py /root/sia/scripts/
COPY ./setup.py /root/sia/setup.py
COPY ./lib /root/sia/lib
RUN python3 /root/sia/scripts/setup_binaries.py /root/sia/setup.py
RUN python3 -m venv /root/venvs/sia

COPY action_schema.xsd /root/sia/action_schema.xsd
RUN mkdir -p /root/sia/lib/xml_schema_validator/src/

RUN --mount=type=cache,target=/root/.cache/pip \
    /root/venvs/sia/bin/pip install -e /root/sia/

# Web frontend build
FROM node:20-alpine AS web-build
WORKDIR /app
COPY web/package*.json ./
RUN npm install
RUN rm -rf /root/.npm/_cacache
COPY web .
RUN npm run build

# Final image
FROM base

COPY --from=itb-env /root/venvs/itb /root/venvs/itb
COPY --from=train-env /root/venvs/train /root/venvs/train
COPY --from=sia-env /root/venvs/sia /root/venvs/sia
COPY --from=web-build /app/dist /root/static/

RUN echo 'source /root/sia/scripts/add_venvs_to_path.sh' >> /root/.bashrc
RUN echo 'source /root/sia/scripts/add_venvs_to_path.sh' >> /etc/profile

WORKDIR /root/desktop

ENTRYPOINT ["/bin/bash", "-c"]
CMD ["/root/sia/scripts/install.sh; /bin/bash -c /root/sia/scripts/restart.sh"]