import pytesseract
from PIL import Image
from pdf2image import convert_from_path
import os 
from tqdm import tqdm
import streamlit




def extract_images_from_pdf(filename: str):
    ''' Takes as input a pdf file (its name) and convert each page into an image'''
    
    root = './HummingBird_prototype/processed_files/' + filename[:-4] + '/'

    # if the file is a pdf, convert it into images
    if filename.endswith('.pdf'):
        streamlit.write("Extracting images from pdf...")
        convert_from_path(root+filename, output_folder=root, fmt='png', output_file=filename[:1]+'page')
        streamlit.write("Extraction done")

    # if not a pdf, save a pdf copy (it will be used for highlighting later [UPDATE: higlighting does not work on images converted to pdf])
    elif filename.endswith('.png') or filename.endswith('.jpg'):
        image = Image.open(root+filename)
        image.convert('RGB').save(root+filename[:-4]+'.pdf')



def ocr_core(filename: str):
    ''' Takes as input a pdf file (its name) and run the OCR on each image created in
    the previous function
    '''

    streamlit.write("Performing OCR...")
    root = './HummingBird_prototype/processed_files/' + filename[:-4] + '/'
    for file in tqdm(os.listdir(root)):
        if file.endswith('.png') or file.endswith('.jpg'):
            result = pytesseract.image_to_string(Image.open(root+file))
            with open(root+file[:-4]+'-ocr.txt', 'w') as f:
                f.write(result)
    streamlit.write("OCR done")


if __name__ == '__main__':
    ocr_core('1.pdf')