#!/usr/bin/env python3

# Import Media!
#
# This script imports photos (or other documents) into date based directories based on 
# the date found in the media's exif data, falling back to the filesystem timestamp if needed.
#
# Copyright 2014-2023 Michael Moore <stuporglue@gmail.com>

import sys
import os
import signal
import mimetypes
import time
import shutil
import hashlib
import re
import subprocess
from os.path import expanduser
import traceback

class ImportMedia:
    ###
    ### SETTINGS
    ###

    def __init__(self):
        # The base directory where photos will be copied into
        self.destdir = expanduser("~") + '/Pictures/'

        # The directory structure that the pictures will be copied into
        self.dirformat = "%Y" + os.sep + "%m" + os.sep + "%d"

        # These mime types will be copied. Other mime types will be ignored
        self.supportedTypes = ['image','video']

        ########################################################################
        ########################################################################
        ########################################################################
        ###
        ### CODE
        ###
        ### You shouldn't need to modify anything below here if you're just using the program
        ###
        ########################################################################
        ########################################################################
        ########################################################################

        self.totalnew = 0
        self.totaldup = 0
        self.totalnonimage = 0
        self.totalsidecar = 0
        self.totalerror = []
        self.progress = 0
        self.totalprocessed = 0
        self.totalfiles = 0
        self.lastprogress = 0

        # m4v isn't in the db for some reason...
        mimetypes.add_type('video/mp4','.m4v',True)
        mimetypes.add_type('image/heic','.heic',True)
        mimetypes.add_type('image/dng','.dng',True)
        mimetypes.add_type('text/plain','.xmp',True)
        mimetypes.add_type('text/plain','.aae',True)

        # Handle Ctrl-c without doing a backtrace
        signal.signal(signal.SIGINT, self.sigint)
        signal.signal(signal.SIGUSR1, self.sigusr1)

        # Print USAGE statement if no args are given
        if len(sys.argv) < 2:
            print("USAGE: import_media.py some_file_or_directory [additional_files_or_directory...]")
            print("")
            print("Media will be put in " + self.destdir + ", organized by " + self.dirformat + ".")
            sys.stdout.flush()
            exit()

    def sigusr1(self, signal, frame):
        print('PROGRESS:' + str(self.progress)


    # Why did we need this? Maybe we don't? 
    def sigint(self, signal, frame):
        sys.exit(0)

    # When we have a file name collision, we will calculate an md5 and compare it. 
    # We are OK with md5 because it is fast and the odds of a collission between 
    # different files with the same exif date seems low. 
    def md5file(self,filename):
        f = open(filename,mode='rb')
        d = hashlib.md5()
        while True:
            data = f.read(8192)
            if not data:
                break;
            d.update(data)
        f.close()
        return d.hexdigest()

    # Try to copy a file. Handle duplicate names
    def copyFile(self, timestamp,filename):
        destpath =  self.destdir + os.sep + time.strftime(self.dirformat,timestamp) + os.sep + os.path.basename(filename)

        if not os.path.exists(os.path.dirname(destpath)):
            os.makedirs(os.path.dirname(destpath))

        # keep modifying destpath with incrementing numbers until we find
        # an unused number or find a duplicate file. If the file is a
        # duplicate, return instead of copying
        if os.path.exists(destpath):
            basename = os.path.basename(filename)
            filenamepart,extension = os.path.splitext(basename)
            counter = 1
            srchash = self.md5file(filename)

        fileexisted = False
        if os.path.exists(destpath):
            fileexisted = True 

        while os.path.exists(destpath):
            if self.md5file(destpath) == srchash:
                print("Identical file [" + destpath + "] already exists. Not re-copying.")
                if os.path.exists(filename + '.xmp'):
                    if not os.path.exists(destpath + '.xmp'):
                        shutil.copy2(filename + '.xmp',destpath + '.xmp')
                        print("\t\t *** Nice! Found an .xmp sidecar file");
                        self.totalnonimage -= 1
                        self.totalsidecar += 1
                    elif ( self.md5file(filename + '.xmp') != self.md5file(destpath + '.xmp') ):
                        print("\t\t *** .xmp sidecar file exists in source and destination but files are different. NOT OVERWRITING EXISTING .xmp!");

                srcAAE = re.sub("\.[^.]*$","",filename) + '.AAE'
                dstAAE = re.sub("\.[^.]*$","",destpath) + '.AAE'
                if os.path.exists(srcAAE):
                    if not os.path.exists(dstAAE):
                        shutil.copy2(srcAAE,dstAAE)
                        print("\t\t *** Nice! Found an .AAE sidecar file");
                        self.totalnonimage -= 1
                        self.totalsidecar += 1
                    elif ( self.md5file(srcAAE) != self.md5file(dstAAE) ):
                        print("\t\t *** .AAE sidecar file exists in source and destination but files are different. NOT OVERWRITING EXISTING .xmp!");
                
                return False
            destpath = self.destdir + os.sep + time.strftime(self.dirformat,timestamp) + os.sep + filenamepart + "_" + str(counter) + extension
            counter += 1

        shutil.copy2(filename,destpath)
        # Copy the xmp sidecar file, if it exists
        if os.path.exists(filename + '.xmp'):
            shutil.copy2(filename + '.xmp',destpath + '.xmp')
            print("\t\t *** Nice! Found an .xmp sidecar file");
            self.totalnonimage -= 1
            self.totalsidecar += 1

        srcAAE = re.sub("\.[^.]*$","",filename) + '.AAE'
        dstAAE = re.sub("\.[^.]*$","",destpath) + '.AAE'
        if os.path.exists(srcAAE):
            shutil.copy2(srcAAE,dstAAE)
            print("\t\t *** Nice! Found an .AAE sidecar file");
            self.totalnonimage -= 1
            self.totalsidecar += 1

        print("[" + destpath + "]")
        return True

    def supportedFile(self, filename):
        maintype,subtype= mimetypes.guess_type(filename) 
        if maintype is None:
            print("ERROR: \"" + filename + "\" --> No mimetype found")
            return False

        category,subtype = maintype.split('/')
        if category in self.supportedTypes:
            return True 

        return False

    # Check all files and send appropriate files off to get copied
    def tryCopy(self, filename):
        destPath = None
        print("\t* [" + filename + "] --> ",end="")
        sys.stdout.flush()

        try:
            # Try to get exif taken date
            maybetags = [
                    "DateTimeOriginal",
                    "CreateDate",
                    "CreationDate",
                    "DateCreated",
                    "TrackCreateDate",
                    "MediaCreateDate",
                    "GPSDateTime",
                    "ModifyDate",
                    "MediaModifyDate",
                    "FileModifyDate",
                    "TrackModifyDate",
                    "FileInodeChangeDate",
                    "FileAccessDate"
                    ] 

            for maybetag in maybetags:
                doutput = subprocess.check_output(["exiftool","-" + maybetag,"-extractEmbedded",filename],encoding='UTF-8')
                if (len(doutput) > 0 and '0000:00:00 00:00:00' not in doutput ):
                    break

            if (len(doutput) > 0 and '0000:00:00 00:00:00' not in doutput):

                # Strip the tag name off the front
                dreone = re.sub(r".*: ",'',doutput)
                
                # Trim the newline off the end
                draw = re.sub(r"\n.*",'',dreone)

                # Remove the timezone info off the end
                draw = re.sub(r"[+-][0-9]{1,2}:[0-9]{2}$",'',draw)
                draw = re.sub(r"Z$",'',draw)

                # Convert dashes or other delimiters to colons
                draw = re.sub(r"[^0-9 ]",':',draw)
                
                if len(draw) > 0:
                    sys.stdout.flush()
                    origTime = draw
                    timestamp = time.strptime(str(origTime),"%Y:%m:%d %H:%M:%S")
                else:
                    sys.stdout.flush()
                    create_date = os.stat(filename)[9]
                    timestamp = time.gmtime(create_date)
            else:
                return False
                create_date = os.stat(filename)[9]
                timestamp = time.gmtime(create_date)
        except Exception as e:
            print(e)
            traceback.print_exc()
            theerror = sys.exc_info()[0]
            sys.stderr.write(theerror)
            print("CAUGHT EXCEPTION: " + theerror)
            timestamp = time.gmtime(os.stat(filename)[9]) # [9] is st_ctime

        return self.copyFile(timestamp,filename)

    def processOne(self, filename):
        try:
            if not self.supportedFile(filename):
                self.totalnonimage += 1
            elif self.tryCopy(filename):
                self.totalnew += 1
            else:
                self.totaldup += 1

            self.totalprocessed += 1
            self.progress = (self.totalprocessed*100/self.totalfiles)
            if  int(self.progress) > int(self.lastprogress):
                self.lastprogress = self.progress
                print('PROGRESS:' + str(int(self.progress)))
                sys.stdout.flush()
        except Exception as e:
            sys.exit(str(e))
            sys.stderr.write(str(e))
            self.totalerror.append(filename)
            self.totalprocessed += 1
            if  int(self.progress) > int(self.lastprogress):
                self.lastprogress = self.progress
                print('PROGRESS:' + str(int(self.progress)))
                sys.stdout.flush()

    def process(self):
        print("Starting Import")
        sys.stdout.flush()

        # Do 2 loops. One to count files, one to import them
        # The count loop is so we can print a progress bar

        # Count loop
        print("Counting files...")
        sys.stdout.flush()
        for file_or_dir in sys.argv[1:]:
            if os.path.isfile(file_or_dir):
                self.totalfiles += 1
            elif os.path.isdir(file_or_dir):
                for (root, subFolders, files) in os.walk(file_or_dir):
                    self.totalfiles += len(files)


        print("Importing " + str(self.totalfiles) + "...")
        sys.stdout.flush()

        # Import loop
        if self.totalfiles > 0:
            self.lastprogress = self.totalprocessed*100/self.totalfiles
        else:
            self.lastprogress = 100
        print('PROGRESS:' + str(int(self.lastprogress)))
        sys.stdout.flush()


        for file_or_dir in sys.argv[1:]:
            sys.stdout.flush()
            if os.path.isfile(file_or_dir):
                self.processOne(file_or_dir)
            elif os.path.isdir(file_or_dir):
                print("\nConsidering " + file_or_dir + " and its subfolders\n")
                for (root, subFolders, files) in os.walk(file_or_dir):
                    for filename in files:
                        self.processOne(root + os.sep + filename)
        if len(self.totalerror):
            print("\nErrors were found when copying the following files:")
            print("\t* " + "\n\t* ".join(self.totalerror))
            sys.stdout.flush()

        print(str(self.totalnew) + " new files copied.")
        print(str(self.totalsidecar) + " sidecar files copied (xmp, AAE).")
        print(str(self.totaldup) + " duplicates not copied.")
        print(str(self.totalnonimage) + " non-image files not copied.")
        print(str(self.totalfiles) + " total files processed.")
        sys.stdout.flush()


im = ImportMedia()
im.process()
