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

mimetypes.init()
load_dotenv()

api_key =       os.getenv('api_key')
id_key =        os.getenv('id_key')
passport_key =  os.getenv('passport_key')

def custom_ocr(image_path, card_type):

    extracted_field = {
        'name': '',
        'dob': '',
        'country': '',
        'gender': '',
        'document_number': '',
    }

    if card_type == "passport":
        response = Client(api_key).extract_document(passport_key, image_path)
    elif card_type == "id_card":
        response = Client(api_key).extract_document(id_key, image_path)

    dictionary = response.to_dict()
    print("dictionary: ", dictionary)

    name_fields = {'Last Name', 'First Name', 'Middle Name'}
    field_mapping = {
        'Document Number': 'document_number',
        'Date of Birth': 'dob',
        'Birth Date': 'dob',
        'State': 'country',
        'County': 'country',
        'Place of Birth': 'country',
        'Nationality': 'country',
        'Sex': 'gender'
    }

    for field in dictionary['formFields']:
        field_name = field['fieldName']

        if 'Name' in field_name and extracted_field['name'] == '':
            if field_name in name_fields:
                extracted_field['name'] = field['value']

        if field_name in field_mapping:
            field_key = field_mapping[field_name]
            if len(field['value']) > 0:
                extracted_field[field_key] = field['value']

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