from unstructured.partition.pdf import partition_pdf

output_path = "./content/"
file_path = output_path + 'attention.pdf'

def pdf_extraction(file_path):
    chunks = partition_pdf(
        filename = file_path,
        infer_table_structure = True,               
        strategy = "fast",                        

        extract_image_block_types = ["Image"],            
        extract_image_block_to_payload = True,      

        chunking_strategy = "by_title",             
        max_characters = 10000,                     
        combine_text_under_n_chars = 2000,          
        new_after_n_chars = 6000,                    
    )

    return chunks

def extract_tables_and_text(chunks):
    tables = []
    text = []

    for chunk in chunks:
        # Check if the chunk is actually an object with metadata
        if not hasattr(chunk, 'metadata'):
            continue
            
        if 'Table' in str(type(chunk)):
            tables.append(chunk)

        if 'CompositeElement' in str(type(chunk)): 
            text.append(chunk) 

    return tables, text        

def get_images_base64(chunks):
    images_b64 = []
    for chunk in chunks:
        if 'CompositeElement' in str(type(chunk)):
            chunk_els = chunk.metadata.orig_elements
            for el in chunk_els:     
                if 'Image' in str(type(el)):
                    images_b64.append(el.metadata.image_base64)

    return images_b64
 
def run_extraction_pipeline(file_path):
    raw_chunks = pdf_extraction(file_path)
    
    # NEW: Safety Filter - only keep items that are objects (have metadata)
    # This removes any stray strings or non-element items created by the partition
    valid_chunks = [chunk for chunk in raw_chunks if hasattr(chunk, 'metadata')]
    
    tables, text_chunks = extract_tables_and_text(valid_chunks)
    images = get_images_base64(valid_chunks)

    return tables, text_chunks, images