#!/usr/bin/env bash

if [ "$TRUEPIC_DEBUG" = "2" ]; then
  set -xeo pipefail
else
  set -eo pipefail
fi

debug_echo() {
    if [ -n "$TRUEPIC_DEBUG" ]; then
        echo "$@"
    fi
}

MEDIA_FILE=$(readlink -f "$1")
OUTPUT_FILE=$2
shift
shift

TRUEPIC_CLI=/home/user/app/truepic
STEG_SCRIPTS=/home/user/app/scripts/
PRIVATE_KEY=/home/user/.truepic/truepic/private.key

filename=$(basename "${MEDIA_FILE}")
extension="${filename##*.}"
if [ "${extension}" = "jpg" ] || [ "${extension}" = "jpeg" ]; then
  mime_type="image/jpeg"
else
  if [ "${extension}" = "png" ]; then
    mime_type="image/png"
  else
    echo "Unsupported file extension: ${extension}"
    exit 1
  fi
fi

debug_echo -n "Signing media..."
signed_no_watermark=$(mktemp).${extension}
${TRUEPIC_CLI} sign --profile truepic $MEDIA_FILE "$@" --output ${signed_no_watermark} > /dev/null 2>&1
debug_echo " --> ${signed_no_watermark}"

debug_echo
debug_echo -n "Extracting manifest..."
no_watermark_manifest=$(mktemp).bin
${TRUEPIC_CLI} manifest extract ${signed_no_watermark} --output ${no_watermark_manifest} > /dev/null 2>&1
debug_echo " --> ${no_watermark_manifest}"

debug_echo
debug_echo -n "Creating watermark signature..."
verification_json=$(${TRUEPIC_CLI} verify ${signed_no_watermark})
if echo "${verification_json}" | jq -e '.manifest_store[0].assertions."c2pa.thumbnail.claim.jpeg"' >/dev/null; then
  thumbnail_key="c2pa.thumbnail.claim.jpeg"
else
  if echo "${verification_json}" | jq -e '.manifest_store[0].assertions."c2pa.thumbnail.claim.png"' >/dev/null; then
    thumbnail_key="c2pa.thumbnail.claim.png"
  else
    echo "Couldn't find thumbnail assertion in the C2PA manifest."
    exit 1
  fi
fi
thumbnail_hash=$(
  echo "${verification_json}" | \
  jq -r '.manifest_store[0].assertions."'${thumbnail_key}'"[0].thumbnail_id'
)
timestamp=$(
  echo "${verification_json}" | \
  jq -r '.manifest_store[0].trusted_timestamp.timestamp'
)
debug_echo -n " ${thumbnail_hash}|${timestamp} ..."
watermark_signature=$(openssl dgst -sha256 -sign ${PRIVATE_KEY} <(echo "${thumbnail_hash}|${timestamp}") | base64 | tr -d '\n')
debug_echo " ${watermark_signature}"

debug_echo
debug_echo -n "Uploading signed media to steg.ai..."
media_id=$(${STEG_SCRIPTS}/upload.sh ${signed_no_watermark} ${mime_type})
debug_echo " --> media_id=${media_id}"
rm -f ${signed_no_watermark}

debug_echo
debug_echo -n "Uploading manifest to steg.ai..."
manifest_id=$(${STEG_SCRIPTS}/upload.sh ${no_watermark_manifest} "application/cbor")
debug_echo " --> media_id=${manifest_id}"

debug_echo
debug_echo -n "Watermarking media..."
encode_response=$(
  curl -s https://api.steg.ai/encode_image_async \
    -H "x-api-key: ${STEG_AI_API_KEY}" \
    --data-raw '{
        "media_id": "'${media_id}'",
        "method": 0,
        "owner": "Truepic",
        "custom": "{\"manifest_id\":\"'${manifest_id}'\",\"watermark_signature\": \"'${watermark_signature}'\"}"
      }'
)
request_id=$(echo "$encode_response" | jq -r '.data.request_id')

if [ -z "$request_id" ] || [ "$request_id" = "null" ]; then
  debug_echo
  echo "No request_id"
  rm -f ${no_watermark_manifest}
  exit 1;
fi

watermark_id=$(echo "$encode_response" | jq -r '.data.encode_media_id')

status_response=""
watermarking_status=""
while [ "$watermarking_status" != "Completed." ]; do
  sleep 1
  debug_echo -n ".."
  status_response=$(
    curl -s https://api.steg.ai/media_status?request_id=${request_id} \
      -H "x-api-key: ${STEG_AI_API_KEY}"
  )
  watermarking_status=$(echo "${status_response}" | jq -r '.data.status')
done

download_url=$(echo "${status_response}" | jq -r '.data.media_data.media_url')
debug_echo " --> media_id=${watermark_id}"

debug_echo
debug_echo -n "Downloading watermarked media..."
watermarked_image=$(mktemp).${extension}
curl -s -o ${watermarked_image} "$download_url"
debug_echo " --> ${watermarked_image}"

debug_echo
debug_echo -n "Re-signing the watermarked media..."
${TRUEPIC_CLI} sign --profile steg ${watermarked_image} \
    --ingredient-manifest-store ${no_watermark_manifest} \
    --output "${OUTPUT_FILE}" \
    --assertions-inline '{
        "assertions": [
          {
            "label": "c2pa.actions",
            "data": {
              "actions": [
                {
                  "action": "ai.steg.watermark",
                  "when": "@now",
                  "softwareAgent": "steg.ai",
                  "parameters": {
                    "description": "An imperceptible digital watermark was added to the file, which can be used to retrieve information later."
                  }
                }
              ]
            }
          }
        ]
      }' > /dev/null 2>&1
debug_echo " --> ${OUTPUT_FILE}"
rm -f ${no_watermark_manifest}

debug_echo
debug_echo -n "Extracting new manifest..."
with_watermark_manifest=$(mktemp).bin
${TRUEPIC_CLI} manifest extract "${OUTPUT_FILE}" --output ${with_watermark_manifest} > /dev/null 2>&1
debug_echo " --> ${with_watermark_manifest}"

debug_echo
debug_echo -n "Uploading new manifest to steg.ai..."
new_manifest_id=$(${STEG_SCRIPTS}/upload.sh ${with_watermark_manifest} "application/cbor")
debug_echo " --> media_id=${new_manifest_id}"
rm -f ${with_watermark_manifest}

debug_echo
debug_echo -n "Updating media with new manifest ID... "
update_result=$(
  curl -s https://api.steg.ai/asset \
    -X POST \
    -H "x-api-key: ${STEG_AI_API_KEY}" \
    --data-raw '{
        "media_id" : "'${watermark_id}'",
        "custom": "{\"manifest_id\":\"'${new_manifest_id}'\",\"watermark_signature\":\"'${watermark_signature}'\",\"original_id\":\"'${watermark_id}'\"}"
      }'
)

debug_echo ${update_result} | jq -r '.message'

debug_echo
debug_echo -n "Deleting un-watermarked image (${media_id}) from steg.ai... "
delete_result=$(
  curl -s https://api.steg.ai/asset \
    -X DELETE \
    -H "x-api-key: ${STEG_AI_API_KEY}" \
    --data-raw '{
        "media_id" : "'${media_id}'"
      }'
)

if [ -n "$TRUEPIC_DEBUG" ]; then echo ${delete_result} | jq -r '.message'; fi

debug_echo
debug_echo -n "Deleting old manifest (${manifest_id}) from steg.ai... "
delete_result=$(
  curl -s https://api.steg.ai/asset \
    -X DELETE \
    -H "x-api-key: ${STEG_AI_API_KEY}" \
    --data-raw '{
        "media_id" : "'${manifest_id}'"
      }'
)

if [ -n "$TRUEPIC_DEBUG" ]; then echo ${delete_result} | jq -r '.message'; fi