#!/bin/bash

KEY_FILE="key.json"
BUCKET_NAME="sage-dist"
OBJECT_PATH="sage-precheck/latest"
LOCAL_FILE="sage-precheck"
ENDPOINT="https://storage.yandexcloud.net"

if [[ ! -f "$KEY_FILE" ]]; then
  echo "File $KEY_FILE not found"
  exit 1
fi

cmd="$1"

case "$cmd" in
  login)
    docker login --username json_key --password-stdin cr.yandex < "$KEY_FILE"
    ;;
    
  precheck-download)
    if ! command -v jq &> /dev/null || ! command -v openssl &> /dev/null; then
        echo "Error: jq or openssl not installed"
        exit 1
    fi

    SERVICE_ACCOUNT_ID=$(jq -r '.service_account_id' "$KEY_FILE")
    KEY_ID=$(jq -r '.id' "$KEY_FILE")
    
    TMP_KEY_FILE=$(mktemp)
    jq -r '.private_key' "$KEY_FILE" > "$TMP_KEY_FILE"

    NOW=$(date +%s)
    EXP=$(($NOW + 3600))

    b64_enc() { openssl enc -base64 -A | tr '+/' '-_' | tr -d '='; }

    HEADER_B64=$(echo -n "{\"alg\":\"PS256\",\"kid\":\"$KEY_ID\"}" | b64_enc)
    PAYLOAD_B64=$(echo -n "{\"aud\":\"https://iam.api.cloud.yandex.net/iam/v1/tokens\",\"iss\":\"$SERVICE_ACCOUNT_ID\",\"iat\":$NOW,\"exp\":$EXP}" | b64_enc)
    
    SIGNATURE_B64=$(echo -n "$HEADER_B64.$PAYLOAD_B64" | \
        openssl dgst -sha256 -sigopt rsa_padding_mode:pss -sigopt rsa_pss_saltlen:-1 -sign "$TMP_KEY_FILE" | \
        b64_enc)

    JWT="$HEADER_B64.$PAYLOAD_B64.$SIGNATURE_B64"
    rm "$TMP_KEY_FILE"

    RESP=$(curl -s -X POST -H 'Content-Type: application/json' \
        -d "$(jq -n --arg jwt "$JWT" '{"jwt": $jwt}')" \
        https://iam.api.cloud.yandex.net/iam/v1/tokens)

    IAM_TOKEN=$(echo "$RESP" | jq -r '.iamToken')

    if [ "$IAM_TOKEN" == "null" ] || [ -z "$IAM_TOKEN" ]; then
        echo "Error getting IAM token: $RESP"
        exit 1
    fi

    curl -s -X GET \
       --header "Authorization: Bearer ${IAM_TOKEN}" \
       --fail  \
       "${ENDPOINT}/${BUCKET_NAME}/${OBJECT_PATH}" \
       > "${LOCAL_FILE}"
       
    echo "$LOCAL_FILE downloaded"

    ;;
    
  *)
    echo "Usage: $0 {login|precheck-download}"
    exit 1
    ;;
esac