import fitz  # PyMuPDF
import os

keywords = [
	#textile
    "geotextile",
    "geo-textile",
    "geo textile",
    "geofabric",
    "geotex",
    "non woven",
    "non-woven",
    "nonwoven",
    "Monofilament",
    "textile",
    "synthetic fibre fabric",
    "filter fabric",
    "ASTM D-4632",
    "ASTM D4632",
       
    #grid
    "geogrid",
    "geo grid",
    "bi axial",
    "bi-axial",
    "biaxial",
    "fiberglass grid",
    "glasgrid",
    "swamp grid",
    "combigrid",
    "Miragrid",    
    
    #companies
    "titan enviro",
    "mirafi",
    "tencate",
    "layfield",
    "solmax",
    "Propex",
    "armtec",
    "terrafix",
    "nilex",
    "Coletanche",
    "agru america",
    
    #liner
    "BGM",
    "GCL",
    "BentoGard",
    "bentonite",
    "bentoliner",
    "geosynthetic clay liner",
    "geomembrane",
    "geo membrane",
    "EPDM liner",
    "PVC liner",
    "LLDPE liner",
    "HDPE liner",
    "RPE liner",
    "40mil LLDPE",
    "40mil HDPE",
    "60mil HDPE",
    "HAZGARD",
    "XR-5",
    "XR-3",
    "GEOFLEX",
    "secondary containment",
    "concrete protective liner",
    "Studliner",
    "Sure-Grip",
    "Sure Grip",
    "ultra grip",
    "HDPE pipe liner",
    "ClosureTurf",

	#linerdrainage
    "geonet",
    "geocomposite",
    "Biplanar",
    "triplanar",
    "hydranet",
    "composite grid",
    "composite geo",
    
    #cells
    "geocell",
    "geo cell",
    "Geoweb",
    "Tough Cell",
    
    #specialty
    "stormwater system",
    "storm water system",
    "Stormwater Storage",
    "StormTank",
    "baffle curtain",
    "concrete canvas",
    "Geosynthetic Cementitious Composite Mat",
    
    #insulation
    "hi-40",
    "foamular 400",
    "rigid insulation",
    
    #erosion
    "erosion control",
    "erosion control blanket",
    "ECB",
    "stenlog",
    "wattle",
    "sediment log",
    "ditch check",
    "bio-log",
    "silt fence", 
    "silt curtain",
    "sediment control",
    "turbidity curtain",
    "safety boom",
    
    "cable concrete",
    "shoreflex",
    "flexamat",
    "Concrete Matting",
    "articulated concrete",
    "articulating concrete",
    "cable concrete",
    
    "TRM",
    "turf reinforcement",
    "pyramat",
    "landlok",
    "Armormax",
    "Fabrinet",
    
    #drainage
    "drain tile",
    "subdrain",
    "weeping tile",
    "drainage tile",
    "frenchdrain",
    "french drain",
    
    "dual wall",
    "culvert",
    "corrugated steel",
    "CSP",
    
    "multiflow",
    "multi-flow",
    "multi flow",
    
    "coffer dam",
    "cofferdam",
    "Aquadam",
    
    "Dewatering Bag",
    "Silt Bag",
    "Geotube",
    
    #misc    
    "detectable warning tile",
    "warning tile",
    "detectable tile",
    "detectable warning systems",
    "ADA tile",
    "armor-tile",
    "armortile",
    "armor tile",
]

def extract_relevant_pages(input_pdf, output_pdf, keywords):
    doc = fitz.open(input_pdf)
    new_doc = fitz.open()
    untrimmed_doc = fitz.open()  # For creating the untrimmed version with highlights
    
    for page_num in range(len(doc)):
        page = doc[page_num]
        text = page.get_text("text")
        lower_text = text.lower()
        
        matched = False
        for keyword in keywords:
            keyword_lower = keyword.lower()
            if keyword_lower in lower_text:
                matched = True
                text_instances = page.search_for(keyword)
                for inst in text_instances:
                    page.add_highlight_annot(inst)
        
        # Add to untrimmed version regardless of match
        untrimmed_doc.insert_pdf(doc, from_page=page_num, to_page=page_num)
        
        # Add to trimmed version only if keyword is matched
        if matched:
            new_doc.insert_pdf(doc, from_page=page_num, to_page=page_num)
    
    # Save the trimmed version (with matched keywords only)
    if len(new_doc) > 0:
        new_doc.save(output_pdf)
        print(f"Trimmed PDF saved as: {output_pdf}")
    
    # Save the untrimmed version (with highlights on all pages)
    if len(untrimmed_doc) > 0:
        untrimmed_output_pdf = output_pdf.replace("trimmed", "untrimmed")
        untrimmed_doc.save(untrimmed_output_pdf)
        print(f"Untrimmed PDF saved as: {untrimmed_output_pdf}")
    
    new_doc.close()
    untrimmed_doc.close()
    doc.close()

def process_all_pdfs(folder, keywords):
    trimmed_folder = os.path.join(folder, "trimmed")
    untrimmed_folder = os.path.join(folder, "untrimmed")
    os.makedirs(trimmed_folder, exist_ok=True)
    os.makedirs(untrimmed_folder, exist_ok=True)

    for file in os.listdir(folder):
        if file.lower().endswith(".pdf"):
            input_pdf = os.path.join(folder, file)
            output_pdf = os.path.join(trimmed_folder, f"{file}")
            print(f"Processing: {file}")
            extract_relevant_pages(input_pdf, output_pdf, keywords)

if __name__ == "__main__":
    folder = os.getcwd()  # Uses current directory
    process_all_pdfs(folder, keywords)