import gradio as gr
import fitz  # PyMuPDF
from PIL import Image, ImageDraw
from io import BytesIO
import pandas as pd
import os
import numpy as np
import google.generativeai as genai
import openai
import base64
import requests
import tempfile
import ast

gemini_api_key = os.getenv("GEMINI_API_KEY")
genai.configure(api_key=gemini_api_key)

openai_api_key = os.getenv("OPENAI_API_KEY")
openai.api_key = openai_api_key

import gradio as gr
import fitz  # PyMuPDF
from PIL import Image
from io import BytesIO
import pandas as pd
import numpy as np
import tempfile


# Define the model extraction functions
def extract_bounding_box_pymupdf(pdf_content):
    bounding_boxes = []
    pdf_file = fitz.open(stream=pdf_content, filetype="pdf")

    for page_index in range(len(pdf_file)):
        page_bbox = []
        page = pdf_file[page_index]
        image_list = page.get_images(full=True)

        for image_index, img in enumerate(page.get_images(full=True), start=1):
            rect = page.get_image_bbox(img[7])
            bbox = list(rect)
            page_bbox.append(bbox)
        bounding_boxes.append(page_bbox)

    pdf_file.close()  # Close the PDF file after use
    return bounding_boxes


def extract_bounding_boxes_gemini(api_key, images):
    # Placeholder for Gemini API integration
    bounding_boxes = [[(0, 0, 100, 100)]] * len(images)  # Dummy bounding boxes
    return bounding_boxes


def extract_bounding_box_gpt(api_key, pdf_content):
    # Placeholder for GPT-4 API integration
    bounding_boxes = [[(0, 0, 100, 100)]] * len(
        fitz.open(stream=pdf_content, filetype="pdf")
    )  # Dummy bounding boxes
    return bounding_boxes


def extract_images_and_tables(pdf_file, model_option):
    if isinstance(pdf_file, str):
        # If input is a file path (usually in testing or local execution)
        with open(pdf_file, "rb") as f:
            pdf_bytes = f.read()
    elif isinstance(pdf_file, bytes):
        # If input is bytes (from Gradio)
        pdf_bytes = pdf_file
    else:
        raise TypeError("Unsupported input type for pdf_file.")

    pdf_document = fitz.open(stream=pdf_bytes, filetype="pdf")

    images = []
    for page_index in range(len(pdf_document)):
        for img_index, img in enumerate(pdf_document.get_page_images(page_index)):
            xref = img[0]
            base_image = pdf_document.extract_image(xref)
            image_bytes = base_image["image"]
            image = Image.open(BytesIO(image_bytes))
            images.append(image)

    tables = []
    for page_num in range(len(pdf_document)):
        page = pdf_document.load_page(page_num)
        text = page.get_text("text")

        lines = [line.strip() for line in text.split("\n") if line.strip()]

        if any("," in line for line in lines):
            rows = [line.split(",") for line in lines]
            tables.extend(rows)

    table_content = ""
    if tables:
        max_columns = max(len(row) for row in tables)
        tables = [row + [""] * (max_columns - len(row)) for row in tables]
        df = pd.DataFrame(tables[1:], columns=tables[0])
        table_content = df.to_csv(index=False)

    pdf_document.close()

    if model_option == "PyMuPDF":
        bounding_boxes = extract_bounding_box_pymupdf(pdf_bytes)
    elif model_option == "Gemini":
        bounding_boxes = extract_bounding_boxes_gemini(
            "your_gemini_api_key_here", images
        )
    elif model_option == "GPT-4":
        bounding_boxes = extract_bounding_box_gpt("your_gpt4_api_key_here", pdf_bytes)
    else:
        bounding_boxes = []

    return images, table_content, bounding_boxes


def handle_model_selection(pdf_file, model_option):
    return extract_images_and_tables(pdf_file, model_option)


# Define the Gradio interface
interface = gr.Interface(
    fn=handle_model_selection,
    inputs=[
        gr.File(type="filepath", label="Upload PDF"),
        gr.Dropdown(
            label="Select Model",
            choices=["PyMuPDF", "Gemini", "GPT-4"],
            value="PyMuPDF",
        ),
    ],
    outputs=[
        gr.Gallery(label="Extracted Images"),
        gr.Textbox(label="Extracted Tables"),
        gr.JSON(label="Extracted Bounding Boxes"),
    ],
    title="PDF Image and Table Extractor",
    description="Upload a PDF to extract images and tables. Choose the model for extraction.",
)

interface.launch(share=True)