from transformers import pipeline
import torch
import gradio as gr
import os
from PIL import Image
import pytesseract
import tempfile
import shutil
from pdf2image import convert_from_path
import json
import spaces

model_name = "deepset/roberta-base-squad2"

# Vision QnA requires: PyTesseract for OCR. Tesseract executable needs to be installed separately.
# sudo apt install tesseract-ocr (https://tesseract-ocr.github.io/tessdoc/Installation.html)

vision_qna = pipeline("document-question-answering", model="impira/layoutlm-document-qa", device=0)

def load_file(file_input, encoding = 'utf-8'):  
  if not os.path.exists(file_input):
    raise FileNotFoundError(f"The file does not exist.")

  with open(file_input, 'r', encoding=encoding) as file:
    try:
      content = file.read()
    except UnicodeDecodeError:
      # If a UnicodeDecodeError occurs, try reading with 'latin1' encoding
      with open(file_input, 'r', encoding='latin1') as file:
        content = file.read()

  return content

def save_image(file):
    try:
      temp_dir = tempfile.mkdtemp()
      file_path = os.path.join(temp_dir, os.path.basename(file.name))
      # Copy the file from the temporary Gradio directory to our temporary directory
      shutil.copyfile(file.name, file_path)
      # when working with saving image files through Gradio, 
      # using `shutil.copyfile` to handle `NamedString` objects for file uploads is the correct approach

      return file_path
    except Exception as e:
      print(e)


def save_pdf(file):
  temp_dir = tempfile.mkdtemp()
  pdf_path = os.path.join(temp_dir, os.path.basename(file.name))
  
  # Copy the file from the temporary Gradio directory to our temporary directory
  shutil.copyfile(file.name, pdf_path)
  
  # Convert PDF to images
  images = convert_from_path(pdf_path)
  
  image_paths = []
  for i, img in enumerate(images):
      image_path = os.path.join(temp_dir, f'page_{i}.png')
      img.save(image_path, 'PNG')
      image_paths.append(image_path)
  
  # print(image_paths)
  return image_paths


def qna_image_content(content, question):
    result = vision_qna(content, question)
    page_result = get_page_result(result, question)
    return get_result_json([page_result])
    
 
  

def get_page_result(raw_result, question):
    # print(raw_result[0]['answer'], raw_result[0]['score'])
    page_result = {
        "result": {
            "question": question,
            "answer": raw_result[0]['answer']
        },
        "confidence": raw_result[0]['score']
    }
    return page_result
    
def get_result_json(page_results):
    result = {
        "pages": page_results
    }
    return json.dumps(result)
def qna_pdf_content(image_paths, question, short_circuit_threshold):
  all_page_results = []
  try:
      for image_path in image_paths:
          result = vision_qna(image=image_path, question=question)
        
          page_result = get_page_result(result, question)
          all_page_results.append(page_result)
          if result[0]['score'] >= short_circuit_threshold:
              break
      return get_result_json(all_page_results)
  except Exception as e:
      return f"An error occurred during processing: {e}"

@spaces.GPU
def predict(doc, question, short_circuit_threshold):
 
  file_input = doc    
  if file_input is not None and question is not None:
    print(f"File type: {type(file_input)}")
    print(f"File name: {file_input.name}")

    filename_lower = file_input.name.lower()
    is_pdf = filename_lower.endswith('.pdf') or filename_lower.endswith("blob")
    if is_pdf:
        try:
            image_paths = save_pdf(file_input)
            
            if not question:
                return "Please provide a question."
            result = qna_pdf_content(image_paths, question, short_circuit_threshold)
            # print(result)
            return result
        except Exception as e:
            return f"An error occurred during vision processing: {e}" 

    else:
      try:
        img_file_path = save_image(file_input)
        
        if not question:
            return "Please provide a question."
        result = qna_image_content(img_file_path, question)
        # print(result)
        return result

      except Exception as e:
        return f"An error occurred during vision processing: {e}" 

  else:
      raise gr.Error("Provide both a document and a question")

gr.close_all()

with gr.Blocks() as demo:
    gr.Markdown("# LayoutLM Document QnA") 
    gr.Markdown("This app performs visual document QnA for docs in PDF or image formats.")

    file_input = gr.File(label="File Upload", file_types=['png', 'jpeg', 'jpg', 'bmp', 'gif', 'pdf'])

    question = gr.Textbox(label="Question", placeholder="Enter your question here...")
    threshold_input = gr.Number(label="Short-circuit threshold", value=0.97, minimum=0, maximum=1)
    output = gr.JSON(label="Answer")


    button = gr.Button("Run model")
    button.click(predict, inputs=[file_input, question, threshold_input], 
                outputs=output)
demo.launch()

# print(qna_image_content("https://gradientflow.com/wp-content/uploads/2023/10/newsletter87-RAG-simple.png", "What is the step prior to embedding?"))