import gradio as gr
from dotenv import load_dotenv
import os 
import torch
import warnings
from PIL import Image
from util import file_helper
from inference.ocr import prepare_batch_for_inference
from inference.inference_handler import handle

os.system('sudo apt install -y -q tesseract-ocr')
os.system('sudo apt install -y -q libtesseract-dev')

load_dotenv()

def get_model():
    model_dir = "tmp"
    model_filename= 'receipt.pth'
    full_path = os.path.join(model_dir, model_filename)
    if os.path.isfile(full_path):
        return full_path
    
    return file_helper.download_gdrive(os.getenv('MODEL_ID'), model_dir, model_filename)

def run_inference(model_path, images_path):
    try:
        inference_batch = prepare_batch_for_inference(images_path)
        context = {"model_dir": model_path}
        print('handle....')
        return handle(inference_batch,context)
    except Exception as err:        
        print('err...', err)


def pretrained_model(model_path):
    if model_path.endswith('.pth'):
        layoutlmv3_model = torch.load(model_path)
        model_path = 'tmp/pre_trained_layoutlmv3'
        layoutlmv3_model.save_pretrained(model_path)
    return model_path


def run(img_path):    
    print('img path: ', img_path)
    model_path = get_model()
    model_path = pretrained_model(model_path)
    results = run_inference(model_path, [img_path])

    if len(results) == 0:
        return Image.open(img_path), {'status': 'error'}
    
    # img_name = os.path.basename(img_path)
    # ext =  os.path.splitext(img_name)[1]
    # img_name = img_name[:img_name.find('.')]    

    # return Image.open(os.path.join('tmp', f'{img_name}_inference.jpg')) 

    print(results[0])
    return Image.open(results[0]['image_path']), results[0]


gr.Markdown('Upload Foto Wajah Kamu (Pastikan hanya terdapat SATU wajah pada)')
iface = gr.Interface(fn=run, inputs=gr.Image(type="filepath"), outputs=["image","json"])
iface.launch()