# dobromirpopov/petals-ai :latest full
# dobromirpopov/petals-ai:full

FROM python:slim

# System dependencies
RUN apt-get update && apt-get install -y git && apt-get clean && rm -rf /var/lib/apt/lists/*

# Upgrade pip
RUN pip install --upgrade pip

# Add ENV for MODEL and HF_TOKEN
ENV MODEL_NAME=stabilityai/StableBeluga2
ENV HF_TOKEN=
ENV DHT_PEERS=
ENV ENV_PARAMS=

# pull petals and bake it into the image. comment for lite image
RUN git clone https://github.com/d-popov/health.petals.dev.git && cd /health.petals.dev && pip install -r requirements.txt

#    FORMATTED_PEERS=$$(echo $$DHT_PEERS | sed "s/[^ ]\+/\'&\'/g")\n\
#    FORMATTED_PEERS=$$(echo $$DHT_PEERS | sed "s/[^ ]\\+/\\'&\\'/g")\n\

# Create the entrypoint script inside the Dockerfile
RUN echo '#!/bin/sh\n\
echo "Script is starting..."\n\
echo "MODE is set to: $MODE"\n\
if [ "$MODE" = "monitor" ] && [ -n "$DHT_PEERS" ]; then\n\
    echo "DHT_PEERS is set to: $DHT_PEERS"\n\
    FORMATTED_PEERS=$DHT_PEERS\n\
    echo "Debug FORMATTED_PEERS: $FORMATTED_PEERS"\n\
##
#   # Convert space-separated string to comma-separated string\n\
#    TEMP_PEERS=$(echo $DHT_PEERS | tr " " ",")\n\
#    # Add single quotes around each element and make it Python list format\n\
#    FORMATTED_PEERS=$(echo $TEMP_PEERS | sed "s/\\([^,]*\\)/\'\\1\'/g")\n\
#    FORMATTED_PEERS="[$FORMATTED_PEERS]"\n\
##
    echo "monitoring PRIVATE swarm. Initial F peers set to $FORMATTED_PEERS"\n\
    echo "monitoring PRIVATE swarm. Initial peers set to $DHT_PEERS"\n\
    echo "monitoring PRIVATE swarm. Initial peers set to ['"$DHT_PEERS"']"\n\
    
    echo "INITIAL_PEERS = ['$DHT_PEERS']" >> /health.petals.dev/config.py\n\
    echo "FORMATTED_PEERS = $FORMATTED_PEERS" >> /health.petals.dev/config.py\n\
else \n\
    echo "monitor has no DHT peers. Monitoring public swarm."\n\
fi\n\

if [ ! -f "/initialized" ]; then\n\
    touch /initialized\n\
    pip install --upgrade petals\n\
    case "$MODE" in\n\
        "monitor")\n\
            # git clone https://github.com/d-popov/health.petals.dev.git && \
            # cd /health.petals.dev && pip install -r requirements.txt\n\
            ;;\n\
        "chat")\n\
            git clone https://github.com/petals-infra/chat.petals.dev.git && \
            cd /chat.petals.dev && pip install -r requirements.txt\n\
            ;;\n\
        "server" | "dht")\n\
            pip install git+https://github.com/bigscience-workshop/petals\n\
            ;;\n\
    esac\n\
fi\n\

if [ -n "$HF_TOKEN" ]; then\n\
    git config --global credential.helper store && \n\
    huggingface-cli login --token $HF_TOKEN --add-to-git-credential && \n\
    ENV_PARAMS="--use_auth_token $ENV_PARAMS"\n\
fi\n\
INITIAL_PEERS_ARG=""\n\
if [ -n "$DHT_PEERS" ]; then\n\
    INITIAL_PEERS_ARG="--initial_peers $DHT_PEERS"\n\
fi\n\
case "$MODE" in\n\
    "monitor")\n\
        cd /health.petals.dev\n\
        exec gunicorn app:app --bind 0.0.0.0:8484 --workers 4 --threads 10\n\
        ;;\n\
    "chat")\n\
        cd /chat.petals.dev\n\
        exec gunicorn app:app --bind 0.0.0.0:8484 --workers 4 --threads 10\n\
        ;;\n\
    "dht")\n\
        exec python -m petals.cli.run_dht --host_maddrs /ip4/0.0.0.0/tcp/31337 --identity_path bootstrap1.id $INITIAL_PEERS_ARG\n\
        ;;\n\
    "server" | *)\n\
        exec python -m petals.cli.run_server $ENV_PARAMS $MODEL_NAME $INITIAL_PEERS_ARG\n\
        ;;\n\
esac' > /entrypoint.sh && chmod +x /entrypoint.sh

# if using separate entrypoint file:
# COPY entrypoint.sh /entrypoint.sh
# RUN chmod +x /entrypoint.sh  TEST

ENTRYPOINT ["/entrypoint.sh"]