"""
PDF Text/Image Classifier
--------------------------

Separates PDFs into:

text_pdfs/
    PDFs with extractable text

image_pdfs/
    Scanned/image-only PDFs requiring OCR


Usage:

python split_pdf_types.py

"""

from pathlib import Path
import shutil

import fitz
from tqdm import tqdm


# -----------------------------
# CONFIG
# -----------------------------

INPUT_FOLDER = Path("./pdfs")

OUTPUT_FOLDER = Path("./classified_pdfs")

TEXT_FOLDER = OUTPUT_FOLDER / "text_pdfs"

IMAGE_FOLDER = OUTPUT_FOLDER / "image_pdfs"


# Minimum characters per page
# below this -> considered scanned

MIN_CHARS_PER_PAGE = 30



# -----------------------------
# CLASSIFICATION
# -----------------------------

def classify_pdf(pdf_path):

    """
    Returns:

    TEXT
    IMAGE

    """

    try:

        doc = fitz.open(pdf_path)


        total_chars = 0
        total_pages = len(doc)

        image_pages = 0


        for page in doc:


            text = page.get_text().strip()


            total_chars += len(text)



            # Count pages containing images

            if len(page.get_images(full=True)) > 0:
                image_pages += 1



        avg_chars = total_chars / max(total_pages,1)



        #
        # Rules
        #

        # No text layer
        if avg_chars < MIN_CHARS_PER_PAGE:

            return "IMAGE"



        # Mostly image pages + little text
        if image_pages / total_pages > 0.7 and avg_chars < 100:

            return "IMAGE"



        return "TEXT"



    except Exception as e:

        print(
            f"Error reading {pdf_path.name}: {e}"
        )

        return "IMAGE"





# -----------------------------
# MAIN
# -----------------------------

def main():


    TEXT_FOLDER.mkdir(
        parents=True,
        exist_ok=True
    )


    IMAGE_FOLDER.mkdir(
        parents=True,
        exist_ok=True
    )


    pdfs=list(
        INPUT_FOLDER.glob("*.pdf")
    )


    print(
        f"Found {len(pdfs)} PDFs"
    )


    stats={
        "TEXT":0,
        "IMAGE":0
    }



    for pdf in tqdm(pdfs):


        category = classify_pdf(pdf)


        if category=="TEXT":

            destination = (
                TEXT_FOLDER /
                pdf.name
            )


        else:

            destination = (
                IMAGE_FOLDER /
                pdf.name
            )



        shutil.copy2(
            pdf,
            destination
        )


        stats[category]+=1




    print("\nCompleted")

    print(
        f"Text PDFs  : {stats['TEXT']}"
    )

    print(
        f"Image PDFs : {stats['IMAGE']}"
    )

    print("\nOutput:")
    print(
        TEXT_FOLDER
    )
    print(
        IMAGE_FOLDER
    )



if __name__=="__main__":
    main()