#!/usr/bin/env python3

from __future__ import annotations
from concurrent.futures import ProcessPoolExecutor, as_completed

import argparse
import subprocess
import tempfile
from pathlib import Path
import time

RED = "\033[91m"
GREEN = "\033[92m"
YELLOW = "\033[93m"
BLUE = "\033[94m"
RESET = "\033[0m"
logfile=Path(f"media_{time.strftime("%Y%m%d_%H%M%S")}.log")

def log_message(message):
    with open(logfile, "a") as log:
        log.write(f"{message}\n")
    print(message)

def get_duration_ms(media: Path) -> int:
    duration = subprocess.check_output(
        [
            "ffprobe",
            "-v",
            "error",
            "-show_entries",
            "format=duration",
            "-of",
            "default=noprint_wrappers=1:nokey=1",
            str(media),
        ],
        text=True,
    ).strip()

    return int(float(duration) * 1000)

def ffmpeg_escape(path: Path) -> str:
    return str(path.resolve()).replace("'", r"'\''")

def create_m4b(root_path: Path, book_dir: Path) -> None:
    media_files = sorted(list(book_dir.glob("*.mp3"))
        + list(book_dir.glob("*.m4a"))
        + list(book_dir.glob("*.m4b"))
        + list(book_dir.glob("*.opus")))
    if not media_files:
        return

    output_file_final = book_dir / f"{book_dir.name}.m4b"

    if output_file_final.exists():
        log_message(f"{RED}[SKIP]{RESET} {output_file_final}")
        log_message("Book already present")
        return

    log_message(f"[CREATE] {output_file_final}")

    with (tempfile.TemporaryDirectory(dir=root_path.parent) as tmp):
        tmpdir = Path(tmp)

        output_file_tmp = tmpdir / f"{book_dir.name}.m4b"
        concat_file = tmpdir / "concat.txt"
        metadata_file = tmpdir / "chapters.txt"

        current_ms = 0

        with concat_file.open("w", encoding="utf-8") as cf, \
            metadata_file.open("w", encoding="utf-8") as mf:
            mf.write(";FFMETADATA1\n")

            for media in media_files:
                cf.write(f"file '{ffmpeg_escape(media)}'\n")

                duration_ms = get_duration_ms(media)

                mf.write("[CHAPTER]\n")
                mf.write("TIMEBASE=1/1000\n")
                mf.write(f"START={current_ms}\n")
                mf.write(f"END={current_ms + duration_ms}\n")
                mf.write(f"title={media.stem}\n")

                current_ms += duration_ms
        cmd_backup = [ "ffmpeg","-y","-f","concat","-safe","0",
                "-i",str(concat_file),"-i",str(metadata_file),
                # cover from first media
                "-i",str(media_files[0]),
                # audio from concat
                "-map","0:a",
                # cover image from first media
                "-map","2:v?",
               # copy artist/album/composer/etc from first media file
               "-map_metadata", "2",
               # copy chapters from ffmetadata
               "-map_chapters", "1",

               "-c:a","copy",
                "-c:a", "aac",
                "-b:a", "128k",
                # copy cover image without re-encoding
                "-c:v","copy","-disposition:v:0","attached_pic",
                str(output_file_tmp),
            ]
        cmd_default = ["ffmpeg", "-y", "-f", "concat", "-safe", "0",
               "-i", str(concat_file), "-i", str(metadata_file),
               # cover from first media
               "-i", str(media_files[0]),
               # audio from concat
               "-map", "0:a",
               # cover image from first media
               "-map", "2:v?",

               # copy artist/album/composer/etc from first media file
               "-map_metadata", "2",
               # copy chapters from ffmetadata
                "-map_chapters", "1",
               "-c:a", "copy",
               # copy cover image without re-encoding
               "-c:v", "copy", "-disposition:v:0", "attached_pic",
               str(output_file_tmp),
               ]
        with logfile.open("a", encoding="utf-8") as log:
            result = subprocess.run(
                cmd_default,
                stdout=log,
                stderr=subprocess.STDOUT,
                text=True,
            )
            if result.returncode == 234:
                log_message(f"{YELLOW}[WARNING]{RESET} - {book_dir.name} -  {result.returncode} retrying with aac encoder")
                result = subprocess.run(
                    cmd_backup,
                    stdout=log,
                    stderr=subprocess.STDOUT,
                    text=True,
                )
        if result.returncode == 0:
            output_file_tmp.rename(output_file_final)
            for media in media_files:
                media.unlink()
            log_message(f"{GREEN}[SUCCESS]{RESET} - {book_dir.name} - delete media files")
        else:
            log_message(f"{RED}[ERROR]{RESET} - {book_dir.name} - Code: {result.returncode}, Stderr: {result.stderr}")

def process_book(args):
    target, book_dir = args
    try:
        create_m4b(target, book_dir)
        return book_dir, None
    except Exception as exc:
        return book_dir, str(exc)

def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "path",
        type=Path,
        help="Book folder or root folder",
    )

    args = parser.parse_args()

    target = args.path.resolve()

    if not target.exists():
        raise SystemExit(f"Path does not exist: {target}")

    # Direct book folder
    if any(list(target.glob("*.mp3"))
        + list(target.glob("*.m4a"))
        + list(target.glob("*.m4b"))
        + list(target.glob("*.opus"))):
        create_m4b(target.parent,target)
        return

    book_dir_set, library_size = find_book_dirs(target)

    for book_dir in book_dir_set:
        log_message(f"Found Media: {book_dir}")

    print("Proceed? y/n")
    name = input()
    if name != "y":
        log_message(f"Aborted")
        exit(0)

    # Root folder containing many books
    with ProcessPoolExecutor(max_workers=8) as executor:
        for book_dir, error in executor.map(
                process_book,
                ((target, book_dir) for book_dir in book_dir_set),
        ):
            if error:
                log_message(f"{RED}[ERROR]{RESET} {book_dir}: {error}")


if __name__ == "__main__":
    main()