#!/bin/bash

set -e
pushd $(cd $(dirname ${0})/..; pwd) > /dev/null

case "${OSTYPE}" in
    msys*) python="winpty python";;
    *) python="python";;
esac

venv() {
    env_dir=./third_party/env/tensorflow
    [ -d "${env_dir}" ] || ${python} -m venv ${env_dir}
    case "${OSTYPE}" in
        msys*) source ${env_dir}/Scripts/activate;;
        *) source ${env_dir}/bin/activate;;
    esac
    ${python} -m pip install --quiet --upgrade pip
}

clean() {
    echo "tf clean"
    rm -rf "./third_party/env/tensorflow"
    rm -rf "./third_party/source/tensorflow"
    rm -rf "./third_party/source/tflite-support"
}

sync() {
    echo "tf sync"
    [ -d "./third_party/source/tensorflow" ] || git clone --quiet https://github.com/tensorflow/tensorflow.git "./third_party/source/tensorflow"
    git -C "./third_party/source/tensorflow" pull --quiet --prune
    mkdir -p "./third_party/source/tflite-support/tensorflow_lite_support/metadata"
    curl --silent --location --output "./third_party/source/tflite-support/tensorflow_lite_support/metadata/metadata_schema.fbs" "https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/metadata_schema.fbs?raw=true"
}

install() {
    echo "tf install"
    if [ "$(uname -sm)" = "Darwin arm64" ]; then
        echo "- arm64 [skip]"
        return
    fi
    venv
    ${python} -m pip install --quiet --upgrade wheel
    ${python} -m pip install --quiet --upgrade tf-nightly
    deactivate
}

schema() {
    echo "tf schema"
    [[ $(grep -U $'\x0D' ./source/tf-proto.js) ]] && crlf=1
    node ./tools/protoc.js --text --root tf --out ./source/tf-proto.js --path ./third_party/source/tensorflow tensorflow/core/protobuf/saved_model.proto tensorflow/core/protobuf/tensor_bundle.proto tensorflow/core/util/saved_tensor_slice.proto tensorflow/core/util/event.proto tensorflow/core/protobuf/config.proto tensorflow/python/keras/protobuf/saved_metadata.proto tensorflow/core/util/memmapped_file_system.proto
    if [[ -n ${crlf} ]]; then
        unix2dos --quiet --newfile ./source/tf-proto.js ./source/tf-proto.js
    fi
    echo "tflite schema"
    [[ $(grep -U $'\x0D' ./source/tflite-schema.js) ]] && crlf=1
    temp1=$(mktemp -d)
    temp2=$(mktemp -d)
    node ./tools/flatc.js --text --root tflite --out ./source/tflite-schema.js ./third_party/source/tensorflow/tensorflow/lite/schema/schema.fbs
    node ./tools/flatc.js --root tflite --out ${temp1}/tflite-metadata-schema.js ./third_party/source/tflite-support/tensorflow_lite_support/metadata/metadata_schema.fbs
    sed "s/var \$root = flatbuffers.get('tflite');//g" < ${temp1}/tflite-metadata-schema.js >> ${temp2}/tflite-metadata-schema.js
    cat ${temp2}/tflite-metadata-schema.js >> ./source/tflite-schema.js
    rm -rf ${temp1} ${temp2}
    if [[ -n ${crlf} ]]; then
        unix2dos --quiet --newfile ./source/tflite-schema.js ./source/tflite-schema.js
    fi
}

metadata() {
    echo "tflite metadata"
    export TF_CPP_MIN_LOG_LEVEL=2
    if [[ $(grep -U $'\x0D' ./source/tflite-metadata.json) ]]; then crlf=1; else crlf=; fi
    node ./tools/tflite_script.js
    if [[ -n ${crlf} ]]; then
        unix2dos --quiet --newfile ./source/tflite-metadata.json ./source/tflite-metadata.json
    fi
    echo "tf metadata"
    if [ "$(uname -sm)" = "Darwin arm64" ]; then
        echo "- arm64 [skip]"
        return
    fi
    venv
    if [[ $(grep -U $'\x0D' ./source/tf-metadata.json) ]]; then crlf=1; else crlf=; fi
    export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
    ${python} ./tools/tf_script.py metadata
    if [[ -n ${crlf} ]]; then
        unix2dos --quiet --newfile ./source/tf-metadata.json ./source/tf-metadata.json
    fi
    deactivate
    echo "keras metadata"
    venv
    if [[ $(grep -U $'\x0D' ./source/keras-metadata.json) ]]; then crlf=1; else crlf=; fi
    ${python} ./tools/keras_script.py metadata
    if [[ -n ${crlf} ]]; then
        unix2dos --quiet --newfile ./source/keras-metadata.json ./source/keras-metadata.json
    fi
    deactivate
}

while [ "$#" != 0 ]; do
    command="$1" && shift
    case "${command}" in
        "clean") clean;;
        "sync") sync;;
        "install") install;;
        "schema") schema;;
        "metadata") metadata;;
    esac
done