#!/usr/bin/env python3

from cgitb import enable
import os, argparse, json
from pathlib import Path

parser = argparse.ArgumentParser()
parser.add_argument('inputfile', help="Image file to be processed")
parser.add_argument('outputfile', nargs="?", help="Output file (.json, optional)")
parser.add_argument('--credentials', nargs="?", help="Google application credential file (.json)") # credentials obtained from https://cloud.google.com/vision/docs/setup#sa-create
# parser.add_argument('--credentials', nargs="?", help="Google application credential file (.json)", default="atea-storage-cd63a39dfeb5.json")
parser.add_argument('--enable_image_labelling', action="store_true", help="") # enables image labelling / annotations
parser.add_argument('--enable_image_ocr', action="store_true", help="") # enables ocr for images
parser.add_argument('--enable_document_ocr', action="store_true", help="") # enables ocr for documents

args = parser.parse_args()

fileName = getattr(args, "inputfile")
outputFile = getattr(args, "outputfile")
gCreds = getattr(args, "credentials")
enable_image_labelling = getattr(args, "enable_image_labelling")
enable_image_ocr = getattr(args, "enable_image_ocr")
enable_document_ocr = getattr(args, "enable_document_ocr")

if (gCreds != None): # if creds not supplied, assume environment var is already set
    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = gCreds; 

p = Path(fileName)
if (outputFile == None): # replace file extension with .csv if output file not defined
    outputFile = p.with_suffix(".json")

def detect_text(path):
    # detects text/objects in image
    from google.cloud import vision
    from google.protobuf.json_format import MessageToDict
    import io

    print(" ---- Running Cloud Vision API ---- ")
    client = vision.ImageAnnotatorClient()

    with io.open(path, 'rb') as image_file:
        content = image_file.read()

    image = vision.Image(content=content)

    if (enable_image_labelling): 
        response = client.label_detection(image=image) # response obj
    elif (enable_image_ocr):
        response = client.text_detection(image=image) # response obj
    elif (enable_document_ocr):
        response = client.document_text_detection(image=image) # response obj

    # print(response)

    # https://stackoverflow.com/questions/52169264/vision-api-how-to-get-json-output 
    texts = MessageToDict(response._pb) # converts protobuf to dictionary (potentially a cleaner solution exists?)

    # https://stackoverflow.com/questions/16291358/python-saving-json-files-as-utf-8
    with open(outputFile, 'w', encoding='utf8') as f:
        json.dump(texts, f, ensure_ascii=False)

    print("output written to " + str(outputFile))

    if response.error.message:
        raise Exception(
            '{}\nFor more info on error messages, check: '
            'https://cloud.google.com/apis/design/errors'.format(
                response.error.message))

detect_text(fileName)