import requests
import mimetypes
import requests
import mimetypes
from butler import Client
import os
from dotenv import load_dotenv

mimetypes.init()
load_dotenv()

api_key = os.getenv('api_key')
queue_id = os.getenv('queue_id')


def custom_ocr(image_path):
    
    extracted_field = {
        'name':             '',
        'dob':              '',
        'country':          '',
        'gender':           '',
        'document_number':  '',
    }

    # PUT THIS IN ENV FILE
    # Make sure to add your API Key to the auth headers
    

    # Response is a strongly typed object
    response = Client(api_key).extract_document(queue_id, image_path)
    dictionary = response.to_dict()
    print("dictionary : ", dictionary)

    for field in dictionary['formFields']:

        if ('Name' in field['fieldName']) and (extracted_field['name'] == ''):
            
            if field['fieldName'] == 'Last Name':
                extracted_field['name'] = field['value']
            elif field['fieldName'] == 'First Name':
                extracted_field['name'] = field['value']
            elif field['fieldName'] == 'Middle Name':
                extracted_field['name'] = field['value']
        
        if field['fieldName'] == 'Document Number':
            extracted_field['document_number'] = field['value']

        if field['fieldName'] == 'Date of Birth':
            extracted_field['dob'] = field['value']

        if (field['fieldName'] == 'State' or field['fieldName'] == 'County' or field['fieldName'] == 'Place of Birth' ) and (len(field['value']) > 0):
            extracted_field['country'] = field['value']

    print("MODEL EXTRACTED FIELDS : ", extracted_field)
    return extracted_field