import os
import argparse
import re
import logging
from PyPDF2 import PdfFileReader, PdfFileWriter
import random
import re
import uuid
from PIL import Image
import img2pdf
import io

image_extns = ['jpg', 'jpeg', 'tiff', 'png', 'gif']


def pdf_global(pdfs, args, pdf_writer):
    cnt = 0

    for pdf in pdfs:

        if cnt == args.count:
            break

        if args.start_string == None and args.end_string == None and args.contains == None:
            pdf_reader = PdfFileReader(pdf)
            pdf_writer = read_pdf(pdf_reader, pdf_writer)

            cnt += 1

        # With end or start strings
        else:
            # Starts

            if args.start_string is not None:
                #print('DEBUG: Entered StartsWith Portion')

                if re.search(r'^'+re.escape((args.start_string))+r'[\w+\s+\d+]+', pdf[:-4], re.IGNORECASE):
                    pdf_reader = PdfFileReader(pdf)
                    pdf_writer = read_pdf(pdf_reader, pdf_writer)

                    cnt += 1

            if args.end_string is not None:
                #print('DEBUG: Entered EndsWith Portion')

                if re.search(r'[\d+\w+\s+]'+re.escape(str(args.end_string))+r'$', pdf[:-4], re.IGNORECASE):
                    pdf_reader = PdfFileReader(pdf)
                    pdf_writer = read_pdf(pdf_reader, pdf_writer)
                    cnt += 1

            if args.contains is not None:
                #print('DEBUG: Entered Contains Portion')

                if re.search(re.escape(str(args.contains)), pdf[:-4], re.IGNORECASE):
                    pdf_reader = PdfFileReader(pdf)
                    pdf_writer = read_pdf(pdf_reader, pdf_writer)
                    cnt += 1

    return pdf_writer


def image_global(imgs, args, pdf_writer, count=-1):

    cnt = 0
    for img in imgs:

        if cnt == args.count:
            break
        if args.start_string == None and args.end_string == None and args.contains == None:
            pdf_writer = read_img(img, pdf_writer)
            cnt += 1

        # With end or start strings
        else:
            # Starts

            if args.start_string is not None:
                #print('DEBUG: Entered StartsWith Portion')

                if re.search(r'^'+re.escape((args.start_string))+r'[\w+\s+\d+]+', img[:-4], re.IGNORECASE):
                    pdf_writer = read_pdf(img, pdf_writer)
                    cnt += 1

            if args.end_string is not None:
                #print('DEBUG: Entered EndsWith Portion')

                if re.search(r'[\d+\w+\s+]'+re.escape(str(args.end_string))+r'$', img[:-4], re.IGNORECASE):
                    pdf_writer = read_pdf(img, pdf_writer)
                    cnt += 1

            if args.contains is not None:
                #print('DEBUG: Entered Contains Portion')

                if re.search(re.escape(str(args.contains)), img[:-4], re.IGNORECASE):
                    pdf_writer = read_pdf(img, pdf_writer)
                    cnt += 1

    return pdf_writer


def read_img(image, pdf_writer):
    try:
        pdf_bytes = img2pdf.convert(image)
        pdf_reader = PdfFileReader(io.BytesIO(pdf_bytes))
        pdf_writer.addPage(pdf_reader.getPage(0))
    except:
        print('Image cannot be added to PDF.')

    return pdf_writer


def read_pdf(pdf_reader, pdf_writer):
    for page in range(pdf_reader.getNumPages()):
        pdf_writer.addPage(pdf_reader.getPage(page))

    return pdf_writer


def main():

    parser = argparse.ArgumentParser(
        description='A package which merges all PDFs from a folder into a single PDF.')
    parser.add_argument('-c', '--count', action='store', dest='count', type=int,
                        help='Total number of PDFs to be merged. Default is all.')  # fixed count, all

    parser.add_argument('-o', '--order', action='store', dest='order', type=int, default=0,
                        help='Order in which PDFs need to be merged. 0 is Ascending, 1 is descending and 2 is shuffle.')  # ascending, descending, shuffle

    parser.add_argument('-s', '--startswith', action='store', dest='start_string',
                        help='Only merge PDFs whose names startswith provided strings.')  # default None
    parser.add_argument('-e', '--endswith', action='store', dest='end_string',
                        help='Only merge PDFs whose names startswith provided strings.')  # default None
    parser.add_argument('-cn', '--contains', action='store', dest='contains',
                        help='Only merge PDFs whose names contains provided strings.')  # default None
    parser.add_argument('-i', '--imageonly', action='store', dest='imageonly', default=0, type=int,
                        help='Only merge Images. Please avoid images with Transparency. Default is 0 (No images)')  # default None
    parser.add_argument('-ip', '--image_pdf', action='store', dest='image_pdf', default=0, type=int,
                        help='Merge Both PDFs and Images. Default is Only PDF')  # default Only PDF
    parser.add_argument('-f', '--filename', action='store', dest='filename_string', default=str(uuid.uuid4()
                                                                                                ).split('-')[0]+'.pdf', help='Filename of the merged PDF. Default is randomly generated names.')

    args = parser.parse_args()
    print(args)

    os.chdir(os.getcwd())

    pdf_writer = PdfFileWriter()

    pdfs = sorted([x for x in os.listdir() if ('.pdf' in x.lower())])
    imgs = sorted([x for x in os.listdir() if (
        x.lower().split('.')[-1] in image_extns)])

    # Ascending
    if args.order == 0:
        if args.imageonly == 1:
            pdf_writer = image_global(imgs, args, pdf_writer)

        elif args.image_pdf == 1:
            pdf_writer = pdf_global(pdfs, args, pdf_writer)
            pdf_writer = image_global(imgs, args, pdf_writer)

        else:
            pdf_writer = pdf_global(pdfs, args, pdf_writer)

        with open(args.filename_string, 'wb') as out:
            pdf_writer.write(out)

    # Descending
    if args.order == 1:
        pdfs = pdfs[::-1]
        imgs = imgs[::-1]

        if args.imageonly == 1:
            pdf_writer = image_global(imgs, args, pdf_writer)

        elif args.image_pdf == 1:
            pdf_writer = pdf_global(pdfs, args, pdf_writer)
            pdf_writer = image_global(imgs, args, pdf_writer)
        else:
            pdf_writer = pdf_global(pdfs, args, pdf_writer)

        with open(args.filename_string, 'wb') as out:
            pdf_writer.write(out)

    # Shuffle
    if args.order == 2:
        pdf_choice = []
        img_choice = []

        count = 0
        for count in range(len(pdfs)):
            choice = random.choice(pdfs)
            pdf_choice.append(choice)
            pdfs.remove(choice)

        count = 0
        for count in range(len(imgs)):
            choice = random.choice(imgs)
            img_choice.append(choice)
            imgs.remove(choice)

        if args.imageonly == 1:
            pdf_writer = image_global(img_choice, args, pdf_writer)

        elif args.image_pdf == 1:
            pdf_writer = pdf_global(pdf_choice, args, pdf_writer)
            pdf_writer = image_global(img_choice, args, pdf_writer)

        else:
            pdf_writer = pdf_global(pdf_choice, args, pdf_writer)

        with open(args.filename_string, 'wb') as out:
            pdf_writer.write(out)


if __name__ == '__main__':
    main()
