import gradio as gr
from util import download
import os 
from dotenv import load_dotenv
from photo_search import PhotoSearch
from PIL import Image

load_dotenv()
# tree_path = download.gdrive(os.getenv('TREE_GDRIVE_ID'), dir='temp')
# data_path = download.gdrive(os.getenv('DATA_GDRIVE_ID'), dir='temp')
tree_path= 'temp/trees.ann'
data_path = 'temp/result_torch.npz'

ai = PhotoSearch(tree_path, data_path)

def search_photo(image):    
    try:
        found_ids = ai.search(image)
        if len(found_ids) == 0:
            return '''
                Foto Serupa Tidak Ditemukan, Tips: 
                - Gunakan foto terupdate,
                - Jika kamu punya foto saat lari, lebih baik gunakan foto itu,
                - Jangan menggunakan foto blur,
                - Jangan upload foto yang menggunakan aksesoris yang tidak kamu gunakan saat lari (semisal kacamata)
                '''
                
        result = f"{len(found_ids)} foto kamu berhasil ditemukan!"
    except Exception as e: 
        result = f"Gagal: {e}"
    return result 


gr_inputs = [
  gr.Image(type='pil')
]

gr.Markdown('Upload Foto Wajah Kamu Ya (Pastikan hanya terdapat SATU wajah pada foto yang kamu upload)')
iface = gr.Interface(fn=search_photo, inputs=gr_inputs, outputs="markdown")
iface.launch()