#!/usr/bin/env python3

# 3rd-party sw: cjxl (libjxl)

import itertools
import subprocess
import sys
from pathlib import Path
from typing import Final

source_dir: Final = Path(sys.argv[1])
if len(sys.argv) > 2:
    backup_dir: Final = Path(sys.argv[2])


def check_source_against_convert(source_dir: Path) -> None:
    for jxl in source_dir.rglob('*.jxl'):
        has_original = jxl.with_suffix('.jpg').exists() or jxl.with_suffix('.JPG').exists()
        if not has_original:
            print(f'Source image for {jxl} has been lost')


def check_source_against_backup(source_dir: Path, backup_dir: Path) -> None:
    pics_src = { p.relative_to(source_dir) for p in source_dir.rglob('*.JPG') }
    pics_src |= { p.relative_to(source_dir) for p in source_dir.rglob('*.jpg') }

    pics_bak = { p.relative_to(backup_dir) for p in backup_dir.rglob('*.JPG') }
    pics_bak |= { p.relative_to(backup_dir) for p in backup_dir.rglob('*.jpg') }

    diff = pics_src - pics_bak
    if diff:
        print('Source contains the following additional images:\n', '\n'.join(map(str, diff)))
    diff = pics_bak - pics_src
    if diff:
        print('Backup contains the following additional images:\n', '\n'.join(map(str, diff)))


def check_successful_conversion(source_dir: Path) -> None:
    pics_src = itertools.chain(source_dir.rglob('*.jpg'), source_dir.rglob('*.JPG'))

    for jpg in pics_src:
        jxl = jpg.with_suffix('.jxl')
        if not jxl.exists():
            print(f'Converted counterpart for {jpg} does not exist')


def convert(source_dir: Path) -> None:
    pics_src_it = itertools.chain(source_dir.rglob('*.jpg'), source_dir.rglob('*.JPG'))
    # set offset to start from image #offset (e.g. in case of a previous error)
    # e.g. offset = 5 will start with converting 5th image
    offset = 0
    pics_src = list(sorted(pics_src_it))[offset:]
    for i, jpg in enumerate(pics_src):
        jxl = jpg.with_suffix('.jxl')
        cmd = ['cjxl', '-v', '--lossless_jpeg=1', '--effort=9', '--brotli_effort=11', jpg, jxl]
        try:
            subprocess.run(cmd, check=True, capture_output=True, encoding='utf-8') # type: ignore[arg-type]
        except subprocess.CalledProcessError as ex:
            print(f'ERR processing file #{i + offset} {jpg}: {ex}')
            print(f'cjxl says: {ex.stderr}\n{ex.stdout}')
            sys.exit(ex.returncode)
        except Exception as ex:
            print(f'ERR unexpected: {ex}')
            sys.exit(1)
        print(f'Converted #{i + offset} {jpg}')

# check_source_against_convert(source_dir)
# check_source_against_backup(source_dir, second_dir)
convert(source_dir)
# check_successful_conversion(source_dir)