#!/bin/bash

set -e

root=$(cd $(dirname ${0})/..; pwd)
build=${root}/build
node_modules=${root}/node_modules
src=${root}/src
third_party=${root}/third_party
tools=${root}/tools

identifier=tensorflow
virtualenv=${third_party}/virtualenv/${identifier}

if [ $(which python3) ]; then
    python="python3"
else
    python="python"
fi

bold() {
    echo "$(tty -s && tput bold)$1$(tty -s && tput sgr0)" 
}

git_sync () {
    mkdir -p "${third_party}"
    if [ -d "${third_party}/${1}" ]; then
        git -C "${third_party}/${1}" fetch --quiet -p
        git -C "${third_party}/${1}" reset --quiet --hard origin/master
    else
        git -C "${third_party}" clone --quiet --recursive ${2}
    fi
}

clean() {
    bold "tf clean"
    rm -rf ${virtualenv}
    rm -rf ${third_party}/${identifier}
}

sync() {
    bold "tf sync"
    git_sync tensorflow https://github.com/tensorflow/tensorflow.git
}

install() {
    bold "tf install"
    case "$(uname)" in
        "Linux")
            [ -n "$(which protoc)" ] || sudo apt install -y protobuf-compiler libprotoc-dev
            ;;
        "Darwin")
            brew list protobuf > /dev/null 2>&1 || brew install protobuf > /dev/null
            ;;
    esac
    [ -n "$(python3 -m pip list --format columns --disable-pip-version-check | grep -w virtualenv)" ] || ${python} -m pip install --force-reinstall --user --quiet virtualenv
    [ -d "${virtualenv}" ] || virtualenv --quiet -p ${python} ${virtualenv}
    source ${virtualenv}/bin/activate
    ${python} -m pip install --quiet protobuf
    deactivate
}

schema() {
    bold "tf schema"
    ${node_modules}/protobufjs/bin/pbjs -t static-module -w closure --no-encode --no-delimited --no-comments --no-convert --no-verify --no-create --keep-case --decode-text -r tf -o ${src}/tf-proto.js \
        ${third_party}/${identifier}/tensorflow/core/protobuf/saved_model.proto \
        ${third_party}/${identifier}/tensorflow/core/protobuf/meta_graph.proto \
        ${third_party}/${identifier}/tensorflow/core/protobuf/saver.proto \
        ${third_party}/${identifier}/tensorflow/core/framework/graph.proto \
        ${third_party}/${identifier}/tensorflow/core/framework/op_def.proto \
        ${third_party}/${identifier}/tensorflow/core/framework/tensor_shape.proto \
        ${third_party}/${identifier}/tensorflow/core/framework/types.proto \
        ${third_party}/${identifier}/tensorflow/core/framework/node_def.proto \
        ${third_party}/${identifier}/tensorflow/core/framework/versions.proto \
        ${third_party}/${identifier}/tensorflow/core/framework/function.proto \
        ${third_party}/${identifier}/tensorflow/core/framework/attr_value.proto \
        ${third_party}/${identifier}/tensorflow/core/framework/tensor.proto \
        ${third_party}/${identifier}/tensorflow/core/framework/variable.proto \
        ${third_party}/${identifier}/tensorflow/core/framework/resource_handle.proto \
        ${third_party}/${identifier}/tensorflow/core/protobuf/saved_object_graph.proto \
        ${third_party}/${identifier}/tensorflow/core/protobuf/trackable_object_graph.proto \
        ${third_party}/${identifier}/tensorflow/core/protobuf/struct.proto
}

metadata() {
    bold "tf metadata"
    source ${virtualenv}/bin/activate
    pushd ${tools} > /dev/null
    protoc --proto_path ${third_party}/${identifier} ${third_party}/${identifier}/tensorflow/core/framework/attr_value.proto --python_out=${tools}
    protoc --proto_path ${third_party}/${identifier} ${third_party}/${identifier}/tensorflow/core/framework/tensor.proto --python_out=${tools}
    protoc --proto_path ${third_party}/${identifier} ${third_party}/${identifier}/tensorflow/core/framework/types.proto --python_out=${tools}
    protoc --proto_path ${third_party}/${identifier} ${third_party}/${identifier}/tensorflow/core/framework/tensor_shape.proto --python_out=${tools}
    protoc --proto_path ${third_party}/${identifier} ${third_party}/${identifier}/tensorflow/core/framework/resource_handle.proto --python_out=${tools}
    protoc --proto_path ${third_party}/${identifier} ${third_party}/${identifier}/tensorflow/core/framework/api_def.proto --python_out=${tools}
    protoc --proto_path ${third_party}/${identifier} ${third_party}/${identifier}/tensorflow/core/framework/op_def.proto --python_out=${tools}
    touch ${tools}/tensorflow/__init__.py
    touch ${tools}/tensorflow/core/__init__.py
    touch ${tools}/tensorflow/core/framework/__init__.py
    ${python} tf-script.py
    rm -rf ${tools}/tensorflow
    popd > /dev/null
    deactivate
}

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