import fitz
import os
import io
import base64
from mimetypes import guess_type
from PIL import Image
import openai
from openai import AzureOpenAI

# Configure Azure OpenAI
openai.api_type = "azure"
api_base = 'API_BASE'
api_key="API_KEY"
deployment_name = 'gpt-4'
api_version = 'API_VERSION'

client = AzureOpenAI(
   api_key=api_key, 
   api_version=api_version,
   base_url=f"{api_base}openai/deployments/{deployment_name}",
)

def local_image_to_data_url(image_path):
   mime_type, _ = guess_type(image_path)
   if mime_type is None:
       mime_type = 'application/octet-stream'
   with open(image_path, "rb") as image_file:
       base64_encoded_data = base64.b64encode(image_file.read()).decode('utf-8')
   return f"data:{mime_type};base64,{base64_encoded_data}"

def generate_image_description(image_url, image_path):
   try:       
       prompt = """You are an image analysis expert. You will be provided with various types of images extracted from documents like research papers, technical blogs, and more.
      
Your task is to generate concise, accurate descriptions of the images without adding any information you are not confident about. Write only a short description.

Important Guidelines:
* Prioritize accuracy: If you are uncertain about any detail, state "None" instead of guessing.
* Avoid hallucinations: Do not add information that is not directly supported by the image.
* Consider context: If the image is a screenshot or contains text, incorporate that information into your description.
"""
       response = client.chat.completions.create(
           model=deployment_name,
           messages=[
               {"role": "system", "content": "You are a helpful assistant."},
               {"role": "user", "content": [
                   {"type": "text", "text": prompt},
                   {"type": "image_url", "image_url": {"url": image_url}}
               ]}
           ],
           max_tokens=1000
       )
       return response.model_dump()["choices"][0]["message"]["content"].strip()
   except Exception as e:
       print(f"Error generating description for {image_path}: {e}")
       return "No description available."

def parse_pdf(pdf_path, text_output_file="output.txt", image_output_dir="output_images"):
   pdf_document = fitz.open(pdf_path)
   os.makedirs(image_output_dir, exist_ok=True)
   with open(text_output_file, "w", encoding="utf-8") as text_file:
       for page_num in range(len(pdf_document)):
           page = pdf_document.load_page(page_num)
           elements = []
           blocks = page.get_text("dict")["blocks"]
          
           for block in blocks:
               if "lines" in block:
                   for line in block["lines"]:
                       for span in line["spans"]:
                           elements.append({"type": "text", "content": span["text"], "bbox": block["bbox"]})
          
           images = page.get_images(full=True)
           processed_images = set()
           for img_index, img in enumerate(images):
               xref = img[0]
               if xref in processed_images:
                   continue
               processed_images.add(xref)
               base_image = pdf_document.extract_image(xref)
               image_bytes = base_image["image"]
               image = Image.open(io.BytesIO(image_bytes))
               image_filename = f"page_{page_num + 1}_image_{img_index + 1}.png"
               image_path = os.path.join(image_output_dir, image_filename)
               image.convert("RGB").save(image_path)
               image_url = local_image_to_data_url(image_path)
               description = generate_image_description(image_url,image_path)
               img_rects = page.get_image_rects(xref)
               if img_rects:
                   img_rect = img_rects[0]
                   elements.append({"type": "image", "content": f"[[IMAGE_DATA_START]] Image URL: {image_path}, Description: {description} [[IMAGE_DATA_END]]", "bbox": (img_rect.x0, img_rect.y0, img_rect.x1, img_rect.y1)})
          
           elements.sort(key=lambda e: e["bbox"][1])
           last_line_was_blank = False
           for element in elements:
               if element["type"] == "text":
                   line = element["content"].strip()
                   if line:
                       text_file.write(line + "\n")
                       last_line_was_blank = False
                   elif not last_line_was_blank:
                       text_file.write("\n")
                       last_line_was_blank = True
               elif element["type"] == "image":
                   text_file.write(element["content"] + "\n")
                   last_line_was_blank = False
           if not last_line_was_blank:
               text_file.write("\n")
   print(f"Text and images saved: {text_output_file}, {image_output_dir}")

pdf_path = "AWS_Bedrock_Blog.pdf"
text_output = "AWS_Bedrock_Blog_1.txt"
image_output = "AWS_Bedrock_Blog_images_1"

parse_pdf(pdf_path, text_output, image_output)