#! /usr/bin/env python3
# -*- coding: UTF-8 -*-

import argparse
import collections
import os
from pathlib import Path
import re
import shutil
import sys
import tqdm

audiooexts = [
    "ac3",
    "al?ac",
    "aiff?c?",
    "amr",
    "at[39]",
    "au",
    "awb",
    "caf",
    "flac?",
    "gsm",
    "m4[abr]",
    "mp[123c]",
    "og[ag]",
    "opus",
    "wave?",
    "wv",
]

audiomatch = [re.compile(r"^.*\.%s$" % i) for i in audiooexts]


def iterfiles_dfs(base: Path):
    """Recursively yields all files in the tree rooted at base in depth-first order.

    Args:
        base (pathlib.Path): Starting directory of probe.

    Yields:
        pathlib.Path: files (excluding symbolic links) within this tree
    """
    try:
        for item in base.iterdir():
            try:
                if item.is_symlink() or item.is_mount():
                    continue
                if item.is_dir():
                    yield from iterfiles_dfs(item)
                elif item.is_file:
                    yield item
            except (PermissionError, OSError):
                continue
    except (PermissionError, OSError):
        pass  # we're not allowed to probe this directory

def iterfiles_bfs(base: Path):
    """Recursively yields all files in the tree rooted at base in breadth-first order

    Args:
        base (pathlib.Path): Starting directory of probe.

    Yields:
        pathlib.Path: files (excluding symbolic links) within this tree
    """
    dirs = collections.deque()
    dirs.append(base)
    while dirs:
        curdir = dirs.popleft()
        try:
            for item in curdir.iterdir():
                try:
                    if item.is_symlink() or item.is_mount():
                        continue
                    if item.is_dir():
                        dirs.append(item)
                    elif item.is_file():
                        yield item
                except (PermissionError, OSError):
                    continue
        except (PermissionError, OSError):
            pass  # we're not allowed to probe this directory


def audio_filter(file: Path) -> bool:
    """Indicate whether the passed file path is considered an audio file by extension

    This function is best used with filter on an iterator of files.
    The corresponding file name is matched against a set of common audio extensions.

    Args:
        file (pathlib.Path): The path pointing to the specified file

    Returns:
        bool: True if the file's extension suggests it is an audio file, False otherwise
    """
    return any(i.fullmatch(file.name) for i in audiomatch)


def main():
    parser = argparse.ArgumentParser(
        "audioscrape.py",
        description="""Scans a directory tree for audio files
        and copies any found to an output tree.
        Directory structure is preserved.""",
    )
    parser.add_argument(
        "source_path",
        type=Path,
        help="The root of the directory tree to scan for audio files.",
    )
    parser.add_argument(
        "dest_path",
        type=Path,
        help="""The root of the destination tree to place all found files.
        Will be created if does not exist.""",
    )
    parser.add_argument(
        "-b",
        "--breadth-first",
        action="store_const",
        dest="search",
        const=iterfiles_bfs,
        default=iterfiles_dfs,
        help="""Use a breadth-first search when probing the tree.
        A depth-first search is used if not specified.""",
    )
    args = parser.parse_args()
    iterfunc = args.search
    source = args.source_path
    dest = args.dest_path
    if not source.exists():
        parser.error(f"directory {source} does not exist")
    elif not source.is_dir():
        parser.error(f"{source} is not a directory")
    if dest.exists():
        if not dest.is_dir():
            parser.error(f"{dest} already exists and is not a directory")
        elif len(os.listdir(dest)) > 0:
            print(f"WARNING: directory {dest} is not empty.", file=sys.stderr)
    print(f"searching {source}...")
    files = list(filter(audio_filter, iterfunc(source)))
    if files:
        files.sort(key=lambda path: len(path.parts))
        print("Found %d audio files. Preparing destination tree..." % len(files))
        dest.mkdir(parents=True, exist_ok=True)
        parents = {file.parent.relative_to(source) for file in files}
        for parent in parents:
            (dest / parent).mkdir(parents=True, exist_ok=True)
        print("Copying...")
        for file in tqdm.tqdm(
            files, miniters=1, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"
        ):
            relfile = file.relative_to(source)
            destfile = dest / relfile
            shutil.copy2(file, destfile)
        print("Done.")
    else:
        print("No audio files were found.")
        sys.exit(0)


if __name__ == "__main__":
    main()
