import streamlit as st
import fitz  # PyMuPDF
from PIL import Image, ImageDraw, ImageFont
import difflib
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import letter
import tempfile
import os
from io import BytesIO

# Load a standard font for annotations
font_path = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
font_size = 10  # Reduced size for better CPU performance
font = ImageFont.truetype(font_path, font_size)

def load_and_compare_documents(file1, file2):
    file1_content = file1.read()
    file2_content = file2.read()

    # Convert PDFs to images
    images1 = pdf_to_images(file1_content)
    images2 = pdf_to_images(file2_content)

    # Perform text comparison
    text1 = extract_text_from_pdf(file1_content)
    text2 = extract_text_from_pdf(file2_content)
    differences = compare_texts(text1, text2)

    # Annotate differences on images
    marked_images_1, marked_images_2 = annotate_images(images1, images2, differences)

    # Generate the output PDF
    pdf_buffer = create_pdf_with_side_by_side(marked_images_1, marked_images_2, differences)
    overall_summary = generate_overall_summary(differences)

    return pdf_buffer, overall_summary

def pdf_to_images(file_content):
    """Converts each page of the PDF into a medium-resolution image for CPU efficiency."""
    images = []
    pdf_document = fitz.open(stream=file_content, filetype="pdf")
    for page in pdf_document:
        pix = page.get_pixmap(dpi=200)  # Lower DPI for faster processing
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        images.append(img)
    pdf_document.close()
    return images

def extract_text_from_pdf(file_content):
    """Extracts text content from the PDF."""
    pdf_document = fitz.open(stream=file_content, filetype="pdf")
    text = [page.get_text("text") for page in pdf_document]
    pdf_document.close()
    return text

def compare_texts(text1, text2):
    """Compares text content and generates differences."""
    differences = []
    position_counter = 1

    for page_num, (page1, page2) in enumerate(zip(text1, text2), start=1):
        diff = list(difflib.ndiff(page1.splitlines(), page2.splitlines()))
        page_diffs = []

        for change in diff:
            if change.startswith("+ ") or change.startswith("- "):
                change_type = "Added" if change.startswith("+ ") else "Deleted"
                value = change[2:]
                description = f"'{value}' {change_type.lower()} at position {position_counter}"
                page_diffs.append({"type": change_type, "value": value, "index": position_counter, "description": description})
                position_counter += 1

        differences.append({"page": page_num, "differences": page_diffs})
    return differences

def annotate_images(images1, images2, differences):
    """Annotates differences on the images."""
    marked_images_1 = {}
    marked_images_2 = {}

    for page_num, (img1, img2) in enumerate(zip(images1, images2), start=1):
        draw1 = ImageDraw.Draw(img1)
        draw2 = ImageDraw.Draw(img2)

        for diff in differences[page_num - 1]["differences"]:
            value = diff["value"]
            position = diff["index"]

            # Annotate differences
            draw2.text((10, 10 + position * 15), f"{position}: {value}", fill="blue", font=font)
            draw2.rectangle([10, 10 + position * 15, 200, 30 + position * 15], outline="red", width=2)

        marked_images_1[page_num] = img1
        marked_images_2[page_num] = img2

    return marked_images_1, marked_images_2

def create_pdf_with_side_by_side(marked_images_1, marked_images_2, differences):
    """Generates a PDF with side-by-side comparison and observation summary."""
    pdf_buffer = BytesIO()
    c = canvas.Canvas(pdf_buffer, pagesize=(letter[0] * 2, letter[1]))

    for page_num, img1 in marked_images_1.items():
        img2 = marked_images_2[page_num]

        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_img_file1:
            img1.save(temp_img_file1, format="PNG")
            temp_img_path1 = temp_img_file1.name

        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_img_file2:
            img2.save(temp_img_file2, format="PNG")
            temp_img_path2 = temp_img_file2.name

        c.drawImage(temp_img_path1, 0, 0, width=letter[0], height=letter[1])
        c.drawImage(temp_img_path2, letter[0], 0, width=letter[0], height=letter[1])
        c.showPage()

        os.remove(temp_img_path1)
        os.remove(temp_img_path2)

        y_position = 750
        c.drawString(10, y_position, f"Observation Summary for Page {page_num}:")
        y_position -= 20

        for diff in differences[page_num - 1]["differences"]:
            text = diff["description"]
            c.drawString(10, y_position, text)
            y_position -= 15
            if y_position < 50:
                c.showPage()
                y_position = 750

        c.showPage()

    c.save()
    pdf_buffer.seek(0)
    return pdf_buffer

def generate_overall_summary(differences):
    """Generates a summary of additions and deletions."""
    total_additions = sum(len([d for d in diff["differences"] if d["type"] == "Added"]) for diff in differences)
    total_deletions = sum(len([d for d in diff["differences"] if d["type"] == "Deleted"]) for diff in differences)
    return {"total_additions": total_additions, "total_deletions": total_deletions}

def main():
    st.title("PDF Comparison Tool with Annotations (CPU-Optimized)")
    st.write("Upload PDFs for side-by-side comparison and observation summary.")

    file1 = st.file_uploader("Document 1 (PDF only)", type=["pdf"])
    file2 = st.file_uploader("Document 2 (PDF only)", type=["pdf"])

    if st.button("Compare Documents") and file1 and file2:
        pdf_buffer, overall_summary = load_and_compare_documents(file1, file2)

        st.subheader("Overall Summary")
        for key, value in overall_summary.items():
            st.write(f"{key.replace('_', ' ').capitalize()}: {value}")

        st.subheader("Download Comparison PDF")
        st.download_button("Download PDF", data=pdf_buffer, file_name="comparison_with_summary.pdf", mime="application/pdf")

if __name__ == "__main__":
    main()