name: Build and Deploy

on:
  push:
    branches: [main]

env:
  REGISTRY: registry.rose-ash.com:5000
  IMAGE_CPU: celery-l1-server
  IMAGE_GPU: celery-l1-gpu-server

jobs:
  build-and-deploy:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4

      - name: Set up SSH
        env:
          SSH_KEY: ${{ secrets.DEPLOY_SSH_KEY }}
          DEPLOY_HOST: ${{ secrets.DEPLOY_HOST }}
        run: |
          mkdir -p ~/.ssh
          echo "$SSH_KEY" > ~/.ssh/id_rsa
          chmod 600 ~/.ssh/id_rsa
          ssh-keyscan -H "$DEPLOY_HOST" >> ~/.ssh/known_hosts 2>/dev/null || true

      - name: Sync code to server
        env:
          DEPLOY_HOST: ${{ secrets.DEPLOY_HOST }}
        run: |
          rsync -avz --delete \
            --exclude '.git' \
            --exclude '__pycache__' \
            --exclude '*.pyc' \
            --exclude '.pytest_cache' \
            --exclude 'venv' \
            ./ "root@$DEPLOY_HOST:/root/art-dag/celery/"

      - name: Build and push CPU image
        env:
          DEPLOY_HOST: ${{ secrets.DEPLOY_HOST }}
        run: |
          ssh "root@$DEPLOY_HOST" "
            cd /root/art-dag/celery
            docker build --build-arg CACHEBUST=\$(date +%s) -t ${{ env.REGISTRY }}/${{ env.IMAGE_CPU }}:latest -t ${{ env.REGISTRY }}/${{ env.IMAGE_CPU }}:${{ github.sha }} .
            docker push ${{ env.REGISTRY }}/${{ env.IMAGE_CPU }}:latest
            docker push ${{ env.REGISTRY }}/${{ env.IMAGE_CPU }}:${{ github.sha }}
          "

      - name: Build and push GPU image
        env:
          DEPLOY_HOST: ${{ secrets.DEPLOY_HOST }}
        run: |
          ssh "root@$DEPLOY_HOST" "
            cd /root/art-dag/celery
            docker build --build-arg CACHEBUST=\$(date +%s) -f Dockerfile.gpu -t ${{ env.REGISTRY }}/${{ env.IMAGE_GPU }}:latest -t ${{ env.REGISTRY }}/${{ env.IMAGE_GPU }}:${{ github.sha }} .
            docker push ${{ env.REGISTRY }}/${{ env.IMAGE_GPU }}:latest
            docker push ${{ env.REGISTRY }}/${{ env.IMAGE_GPU }}:${{ github.sha }}
          "

      - name: Deploy stack
        env:
          DEPLOY_HOST: ${{ secrets.DEPLOY_HOST }}
        run: |
          ssh "root@$DEPLOY_HOST" "
            cd /root/art-dag/celery
            docker stack deploy -c docker-compose.yml celery
            echo 'Waiting for services to update...'
            sleep 10
            docker stack services celery
          "

      - name: Deploy GPU worker
        env:
          GPU_HOST: ${{ secrets.GPU_HOST }}
          SSH_KEY: ${{ secrets.GPU_SSH_KEY }}
        if: ${{ env.GPU_HOST != '' }}
        run: |
          # Set up GPU SSH if different host
          if [ -n "$SSH_KEY" ]; then
            echo "$SSH_KEY" > ~/.ssh/gpu_key
            chmod 600 ~/.ssh/gpu_key
            ssh-keyscan -H "${GPU_HOST#*@}" >> ~/.ssh/known_hosts 2>/dev/null || true

            ssh -i ~/.ssh/gpu_key "$GPU_HOST" "
              docker pull ${{ env.REGISTRY }}/${{ env.IMAGE_GPU }}:latest
              docker stack deploy -c /root/art-dag/celery/docker-compose.yml celery || \
              docker service update --image ${{ env.REGISTRY }}/${{ env.IMAGE_GPU }}:latest celery_l1-gpu-worker
            "
          fi