name: Deploy to SageMaker

on:
  push:
    branches:
      - main

jobs:
  deploy:
    runs-on: ubuntu-latest

    steps:
      - name: Checkout code
        uses: actions/checkout@v2

      - name: Setup Python
        uses: actions/setup-python@v3
        with:
          python-version: '3.10'

      - name: Install dependencies
        run: |
          python -m pip install --upgrade pip
          pip install --no-cache-dir --upgrade pip
          pip install --no-cache-dir torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu121
          pip install --no-cache-dir -r requirements.txt

      - name: Login to AWS
        uses: aws-actions/configure-aws-credentials@v2
        with:
          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
          aws-region: ${{ secrets.AWS_REGION }}

      - name: Login to Amazon ECR
        id: login-ecr
        uses: aws-actions/amazon-ecr-login@v1

      - name: Build and push Docker image
        env:
          ECR_REGISTRY: ${{ steps.login-ecr.outputs.registry }}
          ECR_REPOSITORY: customer-support-chatbot
          IMAGE_TAG: latest
        run: |
          docker build \
            --build-arg AWS_ACCESS_KEY_ID=${{ secrets.AWS_ACCESS_KEY_ID }} \
            --build-arg AWS_SECRET_ACCESS_KEY=${{ secrets.AWS_SECRET_ACCESS_KEY }} \
            --build-arg AWS_DEFAULT_REGION=${{ secrets.AWS_REGION }} \
            -t $ECR_REGISTRY/$ECR_REPOSITORY:$IMAGE_TAG .
          docker push $ECR_REGISTRY/$ECR_REPOSITORY:$IMAGE_TAG

      - name: Deploy model to SageMaker
        run: |
          python src/deploy_sagemaker.py \
            --account_id ${{ secrets.ACCOUNT_ID }} \
            --region ${{ secrets.AWS_REGION }} \
            --role_arn ${{ secrets.SAGEMAKER_ROLE_ARN }} \
            --ecr_repo_name "customer-support-chatbot" \
            --endpoint_name "customer-support-chatbot"