#!/usr/bin/env python
# License: GPL v3 Copyright: 2021, Kovid Goyal <kovid at kovidgoyal.net>

import io
import os
import re

from calibre.constants import iswindows


def open_archive(path_or_stream, mode='r'):
    from py7zr import SevenZipFile
    return SevenZipFile(path_or_stream, mode=mode)


def names(path_or_stream):
    with open_archive(path_or_stream) as zf:
        return tuple(zf.getnames())


class DataSavingWriter(io.BytesIO):

    def close(self):
        return  # make this a no-op as we need to call getvalue() after close


class Writer:

    def __init__(self):
        self.outputs = {}

    def create(self, filename):
        b = self.outputs[filename] = DataSavingWriter()
        return b

    def asdatadict(self):
        return {k: v.getvalue() for k, v in self.outputs.items()}

    def __enter__(self):
        return self

    def __exit__(self, *a):
        for v in self.outputs.values():
            io.BytesIO.close(v)
        self.outputs.clear()


def read_file(archive, name):
    with Writer() as w:
        archive.extract(targets=[name], factory=w)
        for v in w.outputs.values():
            return v.getvalue()
    raise KeyError(f'No file named {name} in archive')


def extract_member(path_or_stream, match=None, name=None):
    if iswindows and name is not None:
        name = name.replace(os.sep, '/')
    if match is None:
        match = re.compile(r'\.(jpg|jpeg|gif|png)\s*$', re.I)

    def is_match(fname):
        if iswindows:
            fname = fname.replace(os.sep, '/')
        return (name is not None and fname == name) or \
               (match is not None and match.search(fname) is not None)

    with open_archive(path_or_stream) as ar:
        all_names = list(filter(is_match, ar.getnames()))
        if all_names:
            return all_names[0], read_file(ar, all_names[0])


def extract_cover_image(stream):
    pos = stream.tell()
    from calibre.libunzip import name_ok, sort_key
    all_names = sorted(names(stream), key=sort_key)
    stream.seek(pos)
    for name in all_names:
        if name_ok(name):
            return extract_member(stream, name=name, match=None)


def extract(path_or_stream, location):
    with open_archive(path_or_stream) as f:
        f.extract(location)


# Test {{{

def test_basic():
    from tempfile import TemporaryDirectory

    from calibre import CurrentDir

    tdata = {
        '1/sub-one': b'sub-one\n',
        '2/sub-two.txt': b'sub-two\n',
        'F\xfc\xdfe.txt': b'unicode\n',
        'max-compressed': b'max\n',
        'one.txt': b'one\n',
        'symlink': b'2/sub-two.txt',
        'uncompressed': b'uncompressed\n',
        '\u8bf6\u6bd4\u5c41.txt': b'chinese unicode\n'}

    def do_test():
        for name, data in tdata.items():
            if '/' in name:
                os.makedirs(os.path.dirname(name), exist_ok=True)
            with open(name, 'wb') as f:
                f.write(data)
        with open_archive(os.path.join('a.7z'), mode='w') as zf:
            for name in tdata:
                zf.write(name)
        with open_archive(os.path.join('a.7z')) as zf:
            if set(zf.getnames()) != set(tdata):
                raise ValueError('names not equal')
            with Writer() as w:
                zf.extractall(factory=w)
                read_data = w.asdatadict()
            if read_data != tdata:
                raise ValueError('data not equal')

        os.mkdir('ex')
        extract('a.7z', 'ex')
        for name, data in tdata.items():
            if name not in '1 2 symlink'.split():
                with open(os.path.join(tdir, 'ex', name), 'rb') as s:
                    if s.read() != data:
                        raise ValueError(f'Did not extract {name} properly')
        if extract_member('a.7z', name='one.txt')[1] != tdata['one.txt']:
            raise ValueError('extract_member failed')

    with TemporaryDirectory('test-7z') as tdir, CurrentDir(tdir):
        do_test()


if __name__ == '__main__':
    test_basic()
