import zipfile
import tarfile
from pathlib import Path


def _get_archive_writer(output_path):
    """
    Detects and returns the appropriate archive writer

    output_path is the pathlib.Path of the archive to write
    """
    if not output_path.suffixes:
        raise ValueError('Output name has no suffix: %s' % output_path.name)
    elif output_path.suffixes[-1].lower() == '.zip':
        archive_root = Path(output_path.stem)
        output_archive = zipfile.ZipFile(str(output_path), 'w', zipfile.ZIP_DEFLATED)

        def add_func(in_path, arc_path):
            """Add files to zip archive"""
            if in_path.is_dir():
                for sub_path in in_path.rglob('*'):
                    output_archive.write(
                        str(sub_path), str(arc_path / sub_path.relative_to(in_path)))
            else:
                output_archive.write(str(in_path), str(arc_path))
    elif '.tar' in output_path.name.lower():
        if len(output_path.suffixes) >= 2 and output_path.suffixes[-2].lower() == '.tar':
            tar_mode = 'w:%s' % output_path.suffixes[-1][1:]
            archive_root = Path(output_path.with_suffix('').stem)
        elif output_path.suffixes[-1].lower() == '.tar':
            tar_mode = 'w'
            archive_root = Path(output_path.stem)
        else:
            raise ValueError('Could not detect tar format for output: %s' % output_path.name)
        output_archive = tarfile.open(str(output_path), tar_mode)
        add_func = lambda in_path, arc_path: output_archive.add(str(in_path), str(arc_path))
    else:
        raise ValueError('Unknown archive extension with name: %s' % output_path.name)
    return output_archive, add_func, archive_root