import argparse
import glob
import os
import shutil
from pathlib import Path

import fitz  # PyMuPDF, imported as fitz for backward compatibility reasons
import numpy as np
from tqdm import tqdm
import cv2 
from pdf2image import convert_from_path


def convert_pdf2image(filename, out_dir, dpi=300, reserve_parent_dir=False, is_get_first_page=False):
    """Convert and save to disk

    Args:
        filename (_type_): _description_
        out_dir (_type_): _description_
        dpi (int, optional): _description_. Defaults to 300.
        reserve_parent_dir (bool, optional): _description_. Defaults to False.
    """
    out_dir = Path(out_dir)
    filename = Path(filename)
    filename_str = str(filename)

    if reserve_parent_dir:
        parent_dir = filename.parent.name
        out_dir = out_dir / parent_dir

    if not out_dir.exists():
        out_dir.mkdir(parents=True)
    if ".txt" in str(filename).lower():
        return 
    if ".pdf" not in str(filename).lower():
        shutil.copy(filename, out_dir)
        return

    
    # doc = fitz.open(filename_str)  # open document
    # # zoom = 2  # zoom factor, standard: 72 dpi
    
    # # dpi = 300
    # zoom = dpi // 72
    # magnify = fitz.Matrix(zoom, zoom) 
    # for idx, page in enumerate(doc):
    #     pix = page.get_pixmap(matrix=magnify)  # render page to an image

    #     outpath = out_dir / Path(os.path.splitext(os.path.basename(filename))[0] + "_" + str(idx+1) + ".jpg")
    #     pix.pil_save(outpath)
    try:
        imgs = pdf_to_image(pdf=filename_str, is_get_first_page=is_get_first_page, dpi=dpi)
    except:
        print("Use v2: ", filename_str)
        imgs = pdf_to_image_v2(pdf=filename_str, is_get_first_page=is_get_first_page, dpi=dpi)
        print("Len img: ", len(imgs))
    for idx, img in enumerate(imgs):
        outpath = str(out_dir / Path(os.path.splitext(os.path.basename(filename))[0] + "_" + str(idx+1) + ".jpg"))
        
        cv2.imwrite(img=img, filename=outpath)
    

def pdf_to_image_v2(pdf, dpi=300, is_get_first_page=False, max_page=1000):
    """_summary_

    Args:
        pdf (_type_): _description_
        dpi (int, optional): _description_. Defaults to 300.
        is_get_first_page (bool, optional): _description_. Defaults to False.
        max_page (int, optional): _description_. Defaults to 1000.

    Raises:
        NotImplementedError: _description_

    Returns:
        _type_: _description_
    """
    if isinstance(pdf, str):
        if not os.path.exists(pdf):
            print(f"Not found pdf path at {pdf}")
            return []
        imgs = convert_from_path(pdf, dpi=dpi)  # PILLOW 
    else:
        raise NotImplementedError(f"Not yet implement for {type(pdf)} type !!!")
    
    
    # zoom = dpi // 72
    # magnify = fitz.Matrix(zoom, zoom) 
    cv_imgs = []
    
    for idx, img in enumerate(imgs):
        img = img.convert("RGB")
        cv_img = np.array(img)
        cv_img = cv_img[:, :, ::-1].copy()
        cv_imgs.append(cv_img)
        if is_get_first_page or idx >= max_page:
            break
    return cv_imgs


def pdf_to_image(pdf, dpi=300, is_get_first_page=False, max_page=1000):
    """_summary_

    Args:
        pdf (_type_): _description_
        dpi (int, optional): _description_. Defaults to 300.
        is_get_first_page (bool, optional): _description_. Defaults to False.
        max_page (int, optional): _description_. Defaults to 1000.

    Raises:
        NotImplementedError: _description_

    Returns:
        _type_: _description_
    """
    if isinstance(pdf, str):
        if not os.path.exists(pdf):
            print(f"Not found pdf path at {pdf}")
            return []
        doc = fitz.open(pdf)  # open document
    elif isinstance(pdf, bytes):
        doc = fitz.open(stream=pdf, filetype='pdf')
    else:
        raise NotImplementedError(f"Not yet implement for {type(pdf)} type !!!")
    
    
    zoom = dpi // 72
    magnify = fitz.Matrix(zoom, zoom) 
    imgs = []
    
    for idx, page in enumerate(doc):
        pix = page.get_pixmap(matrix=magnify)  # render page to an image
        
        im = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.h, pix.w, pix.n)
        im = np.ascontiguousarray(im[..., [2, 1, 0]])  # rgb to bgr
        imgs.append(im)
        if is_get_first_page or idx >= max_page:
            break
    return imgs


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("--pdf_dir", type=str)
    parser.add_argument("--out_dir", type=str)
    parser.add_argument("--reserve_parent_dir", action='store_true')
    args = parser.parse_args()

    paths =  glob.glob(args.pdf_dir + "/*") \
            + glob.glob(args.pdf_dir + "/*/*")  \
            + glob.glob(args.pdf_dir + "/*/*/*") 
            # + glob.glob(args.pdf_dir + "/*")
    print(f"Total pdf paths in {args.pdf_dir}: {len(paths)} ")

    error_pdfs = []
    for path in tqdm(paths):
        path = str(path)
        try:
            convert_pdf2image(path, args.out_dir, reserve_parent_dir=args.reserve_parent_dir)
        except Exception as err:
            print(err, path)
            error_pdfs.append(path)
            continue
    print("Total error pdfs: ", len(error_pdfs))
    print(error_pdfs)