import fitz
import os
import streamlit as st
from io import BytesIO

PDF_STORAGE_DIR = 'pdf_storage'

def parse_pdf(pdf_file_name, page_numbers, search_strings, rect_coords):
    doc = fitz.open(os.path.join(PDF_STORAGE_DIR, pdf_file_name))
    pages = [doc.load_page(i) for i in page_numbers]
    output = ""

    output += "### Extracted Text:\n"
    for page in pages:
        output += "### Page %d\n" % page.number
        output += page.get_text() + "\n\n"

    if search_strings:
        output += "\n### Search Results:\n"
        for search_string in search_strings:
            for page in pages:
                rect = page.search_for(search_string)
                if rect:
                    output += f"#### {search_string} - \n\n"
                    for r in rect:
                        r = list(map(lambda x: round(x,3), r))
                        output += f"##### {r} \n\n"
                else:
                    output += f"{search_string} not found\n"
            output += "\n"

    if rect_coords:
        x1, y1, x2, y2 = map(float, rect_coords.split(','))
        rect = fitz.Rect(x1, y1, x2, y2)
        output += "\n### Rectangle Text:\n"
        for page in pages:
            output += page.get_text("text", clip=rect) + "\n\n"
    return output

st.write("### Extract Text from PDF")
        
with st.container():
    pdf_files = os.listdir(PDF_STORAGE_DIR)
    
    pdf_file = st.selectbox("Select PDF file to delete", pdf_files)

    st.write("### Enter Page Numbers:")
    page_numbers = st.text_input("Enter page numbers (comma-separated)")

    st.write("### Enter Search Strings:")
    search_strings = st.text_input("Enter search strings (comma-separated)")

    # Rectangle Coordinates
    st.write("### Enter Rectangle Coordinates:")
    coord_method = st.selectbox("Select input method", ["Enter all at once", "Enter individually"])

    rect_coords = ""

    if coord_method == "Enter all at once":
        rect_coords = st.text_input("Enter rectangle coordinates (x1,y1,x2,y2)", placeholder="e.g., 10,20,30,40")
    else:
        col1, col2 = st.columns(2)
        
        with col1:
            x1 = st.number_input("X1 (top-left x-coordinate)", min_value=0, value=0)
            y1 = st.number_input("Y1 (top-left y-coordinate)", min_value=0, value=0)
        
        with col2:
            x2 = st.number_input("X2 (bottom-right x-coordinate)", min_value=0, value=500)
            y2 = st.number_input("Y2 (bottom-right y-coordinate)", min_value=0, value=500)
        
        rect_coords = f"{x1},{y1},{x2},{y2}"    
    
    if st.button("Extract Text"):
        if not page_numbers:
            page_numbers = [0]
        else:
            page_numbers = [int(i) for i in page_numbers.split(',')]
        
        if not search_strings:
            search_strings = []
        else:
            search_strings = [s.strip() for s in search_strings.split(',')]
        
        output = parse_pdf(pdf_file, page_numbers, search_strings, rect_coords)
        st.write(output)