import numpy as np
import gradio as gr
from fastapi import FastAPI, File, UploadFile
from app.extract_country.country_dictionary import find_country
from app.extract_dates.date_engine import date_extractor
from app.extract_gender.gender_extractor import gender_extract
from app.extract_identity_number.doc_number_extractor import doc_number
from app.ocr_engine.ocr import OCR
from paddleocr import PaddleOCR
from app.layoutLM_api.api import custom_ocr
import cv2

def ocr(image_input, webcam_input):

    file_path = 'app/images/idcards/input.jpg'

    print("image_input : ", type(image_input))
    print("Webcam Image : ", type(webcam_input))

    if image_input is not None and image_input.any():
        print("In input image")
        cv2.imwrite(file_path, image_input)
    else:
        print("In webcam")
        cv2.imwrite(file_path, webcam_input)
    
    # print(type(input_img))
    ocr = PaddleOCR(use_angle_cls=True, lang='en') # need to run only once to download and load model into memory

    
    
    # with open(file_path, "wb+") as file_object:
    #     file_object.write(input_img.file.read())
    
    dictionary = custom_ocr(file_path)

    if '' in list(dictionary.values()):
        print("Missing value found in Dic")
        
        extract_text = []
        result = ocr.ocr(file_path, cls=True)

        extract_text = [line[1][0] for res in result for line in res]
        print("extract_text", extract_text)

        if len(dictionary['gender']) == 0:
            print("Gender Missing")
            gender_found, gender, ocr_list = gender_extract(extract_text)
            
            dictionary["gender"] = gender if gender_found else None

        if len(dictionary['dob']) == 0:
            print("Dob Missing")
            dob_found, dob, ocr_list = date_extractor(extract_text)

            dictionary["dob"] = dob if dob_found else None

        if len(dictionary['country']) == 0:
            print("Country Missing")
            country_found, country, ocr_list = find_country(extract_text)
            dictionary["country"] = country if country_found else None

        if len(dictionary['document_number']) == 0:
            print("document Number missing")

            document_number_found, document_number, ocr_list = doc_number(extract_text)
            dictionary["document_number"] = document_number if document_number_found else None


    print("Updated Dict ",dictionary)
    response = {"Status" : 200, "OCR" : dictionary}  
    return response

# webcam_checkbox = gr.inputs.Checkbox(label="Use webcam", optional=True)

# Define the input objects
image_input = gr.inputs.Image(label="Upload Image")
webcam_input = gr.inputs.Image(label="Webcam", source="webcam")

# Create the Gradio interface
interface = gr.Interface(fn=ocr, inputs=[image_input, webcam_input], outputs="json")
interface.launch(share=True)