from flask import Response
import io
import os
import tarfile

# 4MiB chuns
CHUNK_SIZE = 4_194_304


def header_chunk(filename: str, filepath: str) -> bytes:
    bytes = io.BytesIO()
    stat = os.stat(filepath)

    # Create dummy tar to extract tar header for file
    with tarfile.open(fileobj=bytes, mode='w') as buffer:
        tar_info = tarfile.TarInfo(filepath)
        tar_info.name = filename
        tar_info.size = stat.st_size
        buffer.addfile(tar_info)

    # TODO if we were able to build this chunk without tarfile, it would avoid
    # whole file copy in memory
    return bytes.getvalue()[:512]


class TarSender:
    def __init__(self):
        self.files: dict[str, str] = {}

    def add_file(self, filename: str, filepath: str):
        self.files[filename] = filepath

    def response(self):
        def generate():
            for name, file in self.files.items():
                yield header_chunk(name, file)

                bytes_sent = 0

                with open(file, 'rb') as f:
                    while True:
                        bytes = f.read(CHUNK_SIZE)

                        if len(bytes) == 0:
                            break

                        bytes_sent += len(bytes)
                        yield bytes

                    # Because tar use records of 512 bytes, we need to pad the
                    # file with zeroes to fill the last chunk
                    yield b'\x00' * (512 - bytes_sent % 512)

        return Response(
            generate(),
            mimetype='application/x-tar',
            headers={'Content-Disposition': 'attachment; filename="archive.tar"'}
        )