'''
shift_exif_timezone.py - Convert local timestamps in Exif to UTC timestamps

Timestamps in Exif records are not timezone-aware. They are naive format. So
there will be protocol mismatch; some camera writes "local" timestamp in Exif,
but some software recognizes it as "UTC" on the other hand. To resolve the
situation, this snippet converts timestamps in '0th' and 'Exif' IFD to UTC
timestamps.

Note: GPS always uses UTC timestamps, so timestamps in 'GPS' IFD are not
problematic.
'''
import argparse
import datetime
import pathlib
import re
import shutil
import subprocess

import piexif
import pytz
import tzlocal


def shift_timestamp(local_timestamp_obj, from_zone=None):
    '''
        Convert local zone `datetime.datetime` object to UTC zone
        `datetime.datetime` object.

        :param local_timestamp_obj: Naive `datetime.datetime` object that
            should be treated as local zone.
        :param from_zone: Timezone object for the given local timestamp. Can
            be None to use system locale instead.
        :retvap: Converted UTC-aware `datetime.datetime` object.
    '''
    if from_zone is None:
        from_zone = tzlocal.get_localzone()

    return from_zone.localize(local_timestamp_obj).astimezone(pytz.utc)


def shift_exif_timestamp(local_timestamp, from_zone=None):
    '''
        Convert local Exif timestamp bytes to UTC timestamp bytes.

        :param local_timestamp: Exif-format timestamp (yyyy:mm:dd hh:mm:ss),
            in bytes object.
        :param from_zone: Timezone object for the given local timestamp. Can
            be None to use system locale instead.
        :retvap: Converted UTC timestamp in byts object.
    '''
    FORMAT = '%Y:%m:%d %H:%M:%S'

    return shift_timestamp(
        datetime.datetime.strptime(local_timestamp.decode('utf-8'), FORMAT),
        from_zone=from_zone
    ).strftime(FORMAT).encode('utf-8')


def shift_mov_timestamp(local_timestamp, from_zone=None):
    '''
        Convert local MOV timestamp string to UTC timestamp string.

        :param local_timestamp: MOV-format timestamp (yyyy-mm-dd hh:mm:ss),
            in str object.
        :param from_zone: Timezone object for the given local timestamp. Can
            be None to use system locale instead.
        :retvap: Converted UTC timestamp in str object.
    '''
    FORMAT = '%Y-%m-%d %H:%M:%S'

    return shift_timestamp(
        datetime.datetime.strptime(local_timestamp, FORMAT),
        from_zone=from_zone
    ).strftime(FORMAT)


if __name__ == '__main__':
    # Parse command line arguments.
    parser = argparse.ArgumentParser(
        description=(
            'Convert local timestamps in Exif to UTC timestamps. '
            '"*.JPG" and "*.MOV" files are supported '
            '(ffmpeg is required to convert "*.MOV" files). '
            'Note: "1st" and "thumbnail" IFDs will be removed because they '
            'can produce corrupted output files.'
        )
    )
    parser.add_argument(
        '-z', '--zone',
        type=str,
        help=(
            'tz database timezone name of source timestamps. Omit to use '
            'system locale timezone.'
        )
    )
    parser.add_argument(
        'target_dir',
        type=str,
        help=(
            'Path to the directory that contains target *.JPG / *.MOV files.'
        )
    )
    parser.add_argument(
        'output_dir',
        type=str,
        nargs='?',
        help=(
            'Path to the output directory. If omitted, "output" dirctory '
            'will be created under the target directory.'
        )
    )

    args = parser.parse_args()

    # Wrap arguments with appropriate objects.
    if args.zone is not None:
        args.zone = pytz.timezone(args.zone)

    args.target_dir = pathlib.Path(args.target_dir).resolve()

    if not args.target_dir.is_dir():
        raise RuntimeError('Target directory is not found.')

    if args.output_dir is None:
        args.output_dir = args.target_dir / 'output'
    else:
        args.output_dir = pathlib.Path(args.output_dir).resolve()

    args.output_dir.mkdir(parents=False, exist_ok=True)

    print('Target:', str(args.target_dir))
    print('Output:', str(args.output_dir))

    # Convert *.JPG files.
    for in_path in args.target_dir.glob('*.jpg'):
        out_path = args.output_dir / in_path.name
        print('Processing:', in_path.name)

        exif = piexif.load(str(in_path))

        for ifd_name, ifd_value in (
            ('0th', piexif.ImageIFD.DateTime),
            ('Exif', piexif.ExifIFD.DateTimeOriginal),
            ('Exif', piexif.ExifIFD.DateTimeDigitized),
        ):
            if ifd_name in exif and ifd_value in exif[ifd_name]:
                exif[ifd_name][ifd_value] = shift_exif_timestamp(
                    exif[ifd_name][ifd_value],
                    args.zone
                )

        del exif['1st']
        del exif['thumbnail']

        shutil.copy(in_path, out_path)
        piexif.insert(piexif.dump(exif), str(out_path))

        print('  OK')

    # Convert *.MOV files
    for in_path in args.target_dir.glob('*.mov'):
        out_path = args.output_dir / in_path.name
        print('Processing:', in_path.name)

        r = subprocess.run(
            ['ffprobe', str(in_path)],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            check=True
        )
        found_stamp = None

        for line in r.stderr.decode('utf-8').split('\n'):
            if line.strip().startswith('creation_time'):
                found_stamp = line.split(':', 1)[1].strip()
                break

        if found_stamp is not None:
            found_stamp = shift_mov_timestamp(found_stamp, args.zone)

            subprocess.run(
                [
                    'ffmpeg',
                    '-i',
                    str(in_path),
                    '-acodec',
                    'copy',
                    '-vcodec',
                    'copy',
                    '-metadata',
                    'creation_time=%s' % found_stamp,
                    str(out_path),
                ],
                check=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE
            )
        else:
            shutil.copy(in_path, out_path)

        print('  OK')