import os
import json
import fitz
import streamlit as st
from PIL import Image
import io
import pandas as pd
import pickle
import zipfile
import tempfile

def extract_text_images(
        pdf_path: str, output_folder: str,
        minimum_font_size: int,
        mode: str = 'headerwise',
        header_font_sizes: list[float] = None,
        tolerance: float = 0.01,
        extraction_type: str = 'both',
        headers_are_capital: bool = False
        ) -> str:
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    extraction_data = []
    current_header = None
    current_header_content = []

    def add_current_header_content() -> None:
        nonlocal current_header, current_header_content
        if current_header:
            extraction_data.append({
                'header': current_header,
                'content': current_header_content
            })
            current_header_content = []
        current_header = None

    def is_header_font_size(font_size: float) -> bool:
        return any(
            abs(font_size - header_font_size) <= tolerance
            for header_font_size in header_font_sizes
        )

    def is_bold(font: str) -> bool:
        return 'bold' in font.lower()

    pdf_document = fitz.open(pdf_path)

    for page_number in range(pdf_document.page_count):
        page = pdf_document.load_page(page_number)
        elements = []

        if extraction_type in ('text', 'both'):
            text_blocks = page.get_text("dict")["blocks"]
            lines = {}

            for block in text_blocks:
                if block["type"] == 0:  # Text block
                    for line in block["lines"]:
                        for span in line["spans"]:
                            font_size = span["size"]
                            top = span["bbox"][1]
                            font = span["font"]

                            if font_size < minimum_font_size:
                                continue

                            if top not in lines:
                                lines[top] = []
                            lines[top].append(span)

            for top in sorted(lines.keys()):
                line = lines[top]
                line_text = " ".join([span['text'] for span in line])
                line_font_size = line[0]['size']
                font = line[0]['font']

                if headers_are_capital:
                    line_text_is_header = line_text.isupper()
                else:
                    line_text_is_header = True

                elements.append({
                    'type': 'text',
                    'font_size': line_font_size,
                    'page': page_number + 1,
                    'content': line_text,
                    'x0': line[0]['bbox'][0],
                    'top': top,
                    'font': font,
                    'is_header': line_text_is_header
                })

        if extraction_type in ('images', 'both'):
            image_list = page.get_images(full=True)

            for img_index, img in enumerate(image_list):
                xref = img[0]
                base_image = pdf_document.extract_image(xref)
                image_bytes = base_image["image"]
                image_filename = os.path.join(
                    output_folder,
                    f"page_{page_number + 1}_img_{img_index + 1}.png"
                )

                with open(image_filename, "wb") as img_file:
                    img_file.write(image_bytes)

                img_rect = page.get_image_bbox(img)
                elements.append({
                    'type': 'image',
                    'page': page_number + 1,
                    'path': image_filename,
                    'x0': img_rect.x0,
                    'top': img_rect.y0
                })

        elements.sort(key=lambda e: (e['top'], e['x0']))

        if mode == 'headerwise':
            for element in elements:
                if element['type'] == 'text' and element['is_header'] and is_header_font_size(element['font_size']) and is_bold(element['font']):
                    add_current_header_content()
                    current_header = element['content']
                elif element['type'] == 'text':
                    if current_header_content and current_header_content[-1]['type'] == 'text':
                        current_header_content[-1]['content'] += " " + element['content']
                    else:
                        current_header_content.append({
                            'type': 'text',
                            'content': element['content']
                        })
                elif element['type'] == 'image':
                    current_header_content.append({
                        'type': 'image',
                        'path': element['path']
                    })

        if mode == 'headerwise':
            add_current_header_content()

    pdf_document.close()

    json_output_path = os.path.join(output_folder, 'extraction_data.json')
    with open(json_output_path, 'w', encoding='utf-8') as json_file:
        json.dump(extraction_data, json_file, ensure_ascii=False, indent=4)

    # Save to XLSX
    df = pd.json_normalize(extraction_data, sep='_')
    xlsx_output_path = os.path.join(output_folder, 'extraction_data.xlsx')
    df.to_excel(xlsx_output_path, index=False)

    # Save to Pickle
    pickle_output_path = os.path.join(output_folder, 'extraction_data.pkl')
    with open(pickle_output_path, 'wb') as pickle_file:
        pickle.dump(extraction_data, pickle_file)

    # Create ZIP file
    zip_output_path = os.path.join(output_folder, 'extraction_data.zip')
    with zipfile.ZipFile(zip_output_path, 'w') as zipf:
        zipf.write(json_output_path, os.path.basename(json_output_path))
        zipf.write(xlsx_output_path, os.path.basename(xlsx_output_path))
        zipf.write(pickle_output_path, os.path.basename(pickle_output_path))
        if extraction_type in ('images', 'both'):
            for root, _, files in os.walk(output_folder):
                for file in files:
                    if file.endswith('.png'):
                        zipf.write(os.path.join(root, file), file)

    return json_output_path, xlsx_output_path, pickle_output_path, zip_output_path

def render_pdf_page_as_image(pdf_path: str, page_number: int, zoom: float = 2.0) -> io.BytesIO:
    # Render PDF page as an image
    pdf_document = fitz.open(pdf_path)
    page = pdf_document.load_page(page_number - 1)  # Page number is zero-indexed in fitz
    pix = page.get_pixmap(matrix=fitz.Matrix(zoom, zoom))
    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
    img_bytes = io.BytesIO()
    img.save(img_bytes, format="PNG")
    img_bytes.seek(0)
    pdf_document.close()
    return img_bytes

# Streamlit UI

st.markdown("<h1 style='text-align: center; color: blue;'>PDF DATA SNATCHER: HEADERWISE</h1>", unsafe_allow_html=True)
st.markdown("<h3 style='text-align: center;color: brown;'>Extract valuable text and images from PDFs effortlessly and Convert PDFs into editable text and high-quality images </h3>", unsafe_allow_html=True)

with st.expander("Click here for more information"):
    st.write("""
        **This application allows you to extract text and images from PDF files Headerswise. You can choose to extract only text, only images, or both, and the extracted data can be downloaded in JSON or XLSX format. Additionally, if you choose to extract images, you can download a ZIP file containing both the images and the JSON data.**
        - **What is different about this app?**    
        - 1. The sequence of text and images will get maintained as per its order in pdf file
        - 2. You have options to extract entities from pdf
        - 3. You can download data in JSON or XLSX format
        - **PDF Preview:** You can preview a few pages of the uploaded PDF in the sidebar.
        - **Extraction Type:** Choose whether to extract text, images, or both.
        - **Minimum Font Size:** Set a threshold for the font size; text below this size will be ignored during extraction.
        - **Output:** Download the extracted data as a JSON file, an Excel file, or a ZIP file (if images are included).
        - *AUTHOR : CHINMAY BHALERAO*
        """)

    st.sidebar.markdown(
        """
        <div style="background-color: lightgray; padding: 2px; border-radius: 2px; text-align: center;">
        <h2 style="color: blue; margin: 0;">PDF PREVIEW</h2>
        </div>
        """, unsafe_allow_html=True)

# Upload PDF file
uploaded_file = st.file_uploader("Upload a PDF file", type=["pdf"])

if uploaded_file is not None:
    # Create a temporary directory
    with tempfile.TemporaryDirectory() as temp_dir:
        # Save uploaded file to a temporary path
        temp_pdf_path = os.path.join(temp_dir, "temp_uploaded_file.pdf")
        with open(temp_pdf_path, "wb") as f:
            f.write(uploaded_file.getbuffer())

        # Number of pages to preview
        num_pages_to_preview = st.sidebar.slider("Number of Pages to Preview", min_value=1, max_value=5, value=1)

        # Generate and display thumbnails for selected number of pages
        st.sidebar.write("Preview of Uploaded PDF:")
        for page_number in range(1, num_pages_to_preview + 1):
            thumbnail_bytes = render_pdf_page_as_image(temp_pdf_path, page_number=page_number)
            st.sidebar.image(thumbnail_bytes, caption=f"Page {page_number}", width=300)

        # Extraction type
        st.info("You can select **only text** or **only images** or **text and images both** to extract form pdf")

        extraction_type = st.radio("Extraction Type", options=['text', 'images', 'both'])

        # Headers are capital
        headers_are_capital = st.checkbox("Are Headers in Capital Letters?", value=False)

        # Minimum font size
        st.info("Minimum font size is the size below which size, the text will get ignored for extraction")

        minimum_font_size = st.slider("Minimum Font Size", min_value=8, max_value=20, value=10)

        # Header font sizes
        header_font_sizes_input = st.text_input(
            "Header Font Sizes (comma-separated, e.g., 10, 12.5, 14.75)", value="16.0")
        header_font_sizes = [float(size.strip()) for size in header_font_sizes_input.split(',') if size.strip().replace('.', '', 1).isdigit()]

        if st.button("Start Extraction"):
            # Run extraction
            json_output_path, xlsx_output_path, pickle_output_path, zip_output_path = extract_text_images(
                pdf_path=temp_pdf_path,
                output_folder=temp_dir,
                minimum_font_size=minimum_font_size,
                mode='headerwise',  # Pagewise mode has been removed
                header_font_sizes=header_font_sizes,
                extraction_type=extraction_type,
                headers_are_capital=headers_are_capital
            )

            st.success("Extraction complete!")

            # Display download options
            with open(json_output_path, 'r', encoding='utf-8') as json_file:
                extracted_data = json.load(json_file)
            st.json(extracted_data)  # Display JSON data in Streamlit

            with open(json_output_path, "rb") as file:
                st.download_button(
                    label="Download JSON",
                    data=file,
                    file_name="extraction_data.json",
                    mime="application/json"
                )

            with open(xlsx_output_path, "rb") as file:
                st.download_button(
                    label="Download XLSX",
                    data=file,
                    file_name="extraction_data.xlsx",
                    mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
                )

            with open(pickle_output_path, "rb") as file:
                st.download_button(
                    label="Download Pickle",
                    data=file,
                    file_name="extraction_data.pkl",
                    mime="application/octet-stream"
                )

            if extraction_type in ('images', 'both'):
                with open(zip_output_path, "rb") as file:
                    st.download_button(
                        label="Download ZIP",
                        data=file,
                        file_name="extraction_data.zip",
                        mime="application/zip"
                    )