# -*- coding: utf-8 -*-
# pyknic/lib/io/tar.py
#
# Copyright (C) 2025-2026 the pyknic authors and contributors
# <see AUTHORS file>
#
# This file is part of pyknic.
#
# pyknic is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# pyknic is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with pyknic.  If not, see <http://www.gnu.org/licenses/>.

import dataclasses
import datetime
import grp
import io
import os
import pathlib
import pwd
import tarfile
import typing

from abc import ABCMeta, abstractmethod

from pyknic.lib.verify import verify_value
from pyknic.lib.io import IOGenerator, IOProducer
from pyknic.lib.io.aio_wrapper import IOThrottler, cag
from pyknic.lib.io.aligner import Aligner


MAX_DYNAMIC_FILE_SIZE = (2 ** 10) ** 7  # This is a maximum number of bytes (1 zebibyte) that may be saved "dynamically"
# Dynamically means that a size of data is unknown at archiving startup.
# Units:
#   kibi  Ki 2^10
#   mebi  Mi 2^20
#   gibi  Gi 2^30
#   tebi  Ti 2^40
#   pebi  Pi 2^50
#   exbi  Ei 2^60
#   zebi  Zi 2^70
#   yobi  Yi 2^80
#   robi  Ri 2^90
#   quebi Qi 2^100
#
# Note. The maximum volume size for EXT4 is 1 exbibyte


class StaticTarEntryProto(metaclass=ABCMeta):
    """This class describes tar entry which size is known before start
    """

    @abstractmethod
    def entry(self) -> IOGenerator:
        """Yield bytes that describe tar header and data.
        """
        raise NotImplementedError('This method is abstract')


class DynamicTarEntryProto(metaclass=ABCMeta):
    """This class describes tar entry which size is unknown

    :note: please note the MAX_DYNAMIC_FILE_SIZE
    """

    @abstractmethod
    def entry(self, file_obj: typing.IO[bytes]) -> IOGenerator:
        """Yield bytes that describe tar header and data.

        :param file_obj: file object to which data is written. Header and content must be yielded, but this object
        may be used for patching header only
        """
        raise NotImplementedError('This method is abstract')


class _TarPaddingGenerator:
    """This class calculates required padding and write it
    """

    @verify_value(aligned_to=lambda x: x > 0)
    def __init__(self, aligned_to: int, extra_padding: bool = False):
        """Create new padding helper

        :param aligned_to: number of bytes required for padding
        :param extra_padding: whether extra padding block should be appended at the end
        """
        self.__aligned_to = aligned_to if aligned_to else tarfile.BLOCKSIZE
        self.__extra_padding = extra_padding
        self.__counter = 0

    def _write(self, data: IOProducer) -> IOGenerator:
        """This processor writes data and required padding

        :param data: data to write
        """
        for block in data:
            self.__counter += len(block)
            yield block

        _, extra_bytes = divmod(self.__counter, self.__aligned_to)
        delta = (self.__aligned_to - extra_bytes) if extra_bytes else 0
        if self.__extra_padding:
            delta += self.__aligned_to
        yield (tarfile.NUL * delta) if delta else b''


class _BaseTarInnerFileGenerator(_TarPaddingGenerator):
    """This class helps to save a custom entry inside a tar archive.
    """

    def __init__(self) -> None:
        """Create a new file writer
        """
        _TarPaddingGenerator.__init__(self, tarfile.BLOCKSIZE)

    # noinspection PyMethodMayBeStatic
    def __pax_headers_padding(self, tar_info: tarfile.TarInfo, expected_size: int) -> tarfile.TarInfo:
        """Fill comment section in PAX header in order to comply limitations

        :param tar_info: original header to update
        :param expected_size: expected size of tar header
        """

        if expected_size % tarfile.BLOCKSIZE != 0:
            raise ValueError('Expected tar header is not aligned with the blocksize')

        if len(tar_info.tobuf()) > expected_size:
            raise ValueError('Original tar header is more than expected size')

        comment_chunk = 'pyknic-' + ('-c' * 60) + ';'  # 128 bytes
        tar_info.pax_headers['comment'] = comment_chunk  # type: ignore[index]
        while len(tar_info.tobuf()) < expected_size:
            tar_info.pax_headers['comment'] += comment_chunk  # type: ignore[index]

        if len(tar_info.tobuf()) != expected_size:
            raise ValueError('Unable to align tar header')

        return tar_info

    def _custom_tar_info(
        self,
        filename: str,
        /,
        size: typing.Optional[int] = None,
        mtime: typing.Optional[int] = None,
        mode: typing.Optional[int] = None,
        uid: typing.Optional[int] = None,
        gid: typing.Optional[int] = None,
        uname: typing.Optional[str] = None,
        gname: typing.Optional[str] = None,
        expected_header_size: typing.Optional[int] = None,
    ) -> tarfile.TarInfo:
        """Generate meta-info for a single file inside archive.

        :param filename: filename of an inner file.
        :param size: size of file (zero by default).
        :param mtime: file modification time (is now by default).
        :param mode: file mode (0440 by default).
        :param uid: file UID (current process user is used by default).
        :param gid: file GID (current process group is used by default).
        :param uname: username related to uid.
        :param gname: group name related to gid.
        :param expected_header_size: if specified then it is number of bytes that tar header should have.
        (result may have extra comment message in PAX header in order to achieve expectations)
        """
        if filename.startswith('/'):
            raise ValueError('Tar entries should not have absolute path because of possible tar bomb')

        tar_info = tarfile.TarInfo(name=filename)
        if size is not None:
            tar_info.size = size
        tar_info.mtime = mtime if mtime is not None else int(datetime.datetime.now().timestamp())
        tar_info.mode = mode if mode is not None else int('440', base=8)
        tar_info.type = tarfile.REGTYPE
        tar_info.uid = uid if uid is not None else os.getuid()
        tar_info.gid = gid if gid is not None else os.getgid()
        tar_info.uname = uname if uname is not None else pwd.getpwuid(tar_info.uid).pw_name
        tar_info.gname = gname if gname is not None else grp.getgrgid(tar_info.gid).gr_name

        if expected_header_size is not None and len(tar_info.tobuf()) != expected_header_size:
            tar_info = self.__pax_headers_padding(tar_info, expected_header_size)

        return tar_info

    # noinspection PyMethodMayBeStatic
    def _tar_info_by_file(self, filename: typing.Union[str, pathlib.Path]) -> tarfile.TarInfo:
        """Retrieve meta-info by a real file

        :param filename: file name which meta-info should be retrieved
        """
        # TODO: filename related to some directory?!
        tar_file = tarfile.TarFile('/dev/zero', mode='w')  # TODO: it seams bogus
        tar_info = tar_file.gettarinfo(filename)

        if tar_info.name.startswith('/'):
            raise ValueError('Tar entries should not have absolute path because of possible tar bomb')

        return tar_info


class TarInnerFileGenerator(_BaseTarInnerFileGenerator, StaticTarEntryProto):
    """This class helps to save a static file inside a tar archive.
    """

    def __init__(self, source: typing.Union[str, pathlib.Path]):
        """Create a new inner file generator

        :param source: a path to a file to archive
        """
        _BaseTarInnerFileGenerator.__init__(self)
        self.__source = source

    def entry(self) -> IOGenerator:
        """ :meth:`.StaticTarEntryProto.entry` implementation
        """

        def data_generator() -> IOGenerator:
            tar_info = self._tar_info_by_file(self.__source)
            yield tar_info.tobuf()

            with open(self.__source, 'rb') as f:
                # TODO: consider checking that file modification time hasn't changed!

                for chunk in IOThrottler.sync_reader(f):
                    yield chunk

        yield from self._write(data_generator())


class TarInnerGenerator(_BaseTarInnerFileGenerator, StaticTarEntryProto):
    """This class helps to save a static data (with known size) from generator inside a tar archive.
    """

    def __init__(
        self,
        source: IOProducer,
        data_size: int,
        filename: str,
        /,
        mtime: typing.Optional[int] = None,
        mode: typing.Optional[int] = None,
        uid: typing.Optional[int] = None,
        gid: typing.Optional[int] = None,
        uname: typing.Optional[str] = None,
        gname: typing.Optional[str] = None,
    ):
        """Create a new inner file generator

        :param source: a data to save
        :param filename: name of an inner file
        :param data_size: number of bytes that source has. It must be exact number of bytes (not more, not less)

        :param mtime: file modification time (is now by default)
        :param mode: file mode (0440 by default)
        :param uid: file UID (current process user is used by default)
        :param gid: file GID (current process group is used by default)
        :param uname: username related to uid
        :param gname: group name related to gid        """
        _BaseTarInnerFileGenerator.__init__(self)

        if filename.startswith('/'):
            raise ValueError('File path must be relative')

        self.__source = source
        self.__data_size = data_size
        self.__filename = filename
        self.__mtime = mtime
        self.__mode = mode
        self.__uid = uid
        self.__gid = gid
        self.__uname = uname
        self.__gname = gname

    def entry(self) -> IOGenerator:
        """ :meth:`.StaticTarEntryProto.entry` implementation
        """

        def data_generator() -> IOGenerator:
            tar_info = self._custom_tar_info(
                self.__filename,
                size=self.__data_size,
                mtime=self.__mtime,
                mode=self.__mode,
                uid=self.__uid,
                gid=self.__gid,
                uname=self.__uname,
                gname=self.__gname,
            )
            yield tar_info.tobuf()

            calculated_size = 0
            for data in self.__source:
                calculated_size += len(data)
                if calculated_size > self.__data_size:
                    raise ValueError(
                        'Original data has more bytes than expected. '
                        f'Expected -- {self.__data_size}, found -- {calculated_size}'
                    )

                yield data

            if calculated_size != self.__data_size:
                raise ValueError(
                    'Incorrect number of bytes in original data. '
                    f'Expected -- {self.__data_size}, found -- {calculated_size}'
                )

        yield from self._write(data_generator())


class TarInnerDynamicGenerator(_BaseTarInnerFileGenerator, DynamicTarEntryProto):
    """This class helps to save a dynamic data from generator (which size is unknown) inside a tar archive.
    """

    def __init__(
        self,
        source: IOProducer,
        filename: str,
        /,
        mtime: typing.Optional[int] = None,
        mode: typing.Optional[int] = None,
        uid: typing.Optional[int] = None,
        gid: typing.Optional[int] = None,
        uname: typing.Optional[str] = None,
        gname: typing.Optional[str] = None,
    ):
        """Create a new inner file generator

        :param source: a data to save
        :param filename: name of an inner file

        :param mtime: file modification time (is now by default)
        :param mode: file mode (0440 by default)
        :param uid: file UID (current process user is used by default)
        :param gid: file GID (current process group is used by default)
        :param uname: username related to uid
        :param gname: group name related to gid        """
        _BaseTarInnerFileGenerator.__init__(self)

        if filename.startswith('/'):
            raise ValueError('File path must be relative')

        self.__source = source
        self.__filename = filename
        self.__mtime = mtime
        self.__mode = mode
        self.__uid = uid
        self.__gid = gid
        self.__uname = uname
        self.__gname = gname

    @verify_value(file_obj=lambda x: x.seekable())
    def entry(self, file_obj: typing.IO[bytes]) -> IOGenerator:
        """ :meth:`.DynamicTarEntryProto.entry` implementation
        """

        def data_generator() -> IOGenerator:
            start_pos = file_obj.tell()

            tar_info = self._custom_tar_info(
                self.__filename,
                size=MAX_DYNAMIC_FILE_SIZE,
                mtime=self.__mtime,
                mode=self.__mode,
                uid=self.__uid,
                gid=self.__gid,
                uname=self.__uname,
                gname=self.__gname,
            )
            binary_tar_info = tar_info.tobuf()
            yield binary_tar_info

            data_size = 0
            for chunk in self.__source:
                data_size += len(chunk)

                if data_size > MAX_DYNAMIC_FILE_SIZE:
                    raise ValueError('Are you saving the universe?!')

                yield chunk

            end_pos = file_obj.tell()

            patched_tar_info = self._custom_tar_info(
                self.__filename,
                size=data_size,
                mtime=self.__mtime,
                mode=self.__mode,
                uid=self.__uid,
                gid=self.__gid,
                uname=self.__uname,
                gname=self.__gname,
                expected_header_size=len(binary_tar_info)
            )

            patched_binary_tar_info = patched_tar_info.tobuf()
            assert(len(binary_tar_info) == len(patched_binary_tar_info))

            file_obj.seek(start_pos, os.SEEK_SET)
            file_obj.write(patched_binary_tar_info)
            file_obj.seek(end_pos, os.SEEK_SET)

        yield from self._write(data_generator())


@dataclasses.dataclass
class _FilthyTarEntry:
    """This class helps to track tar entries while reading. Using it in other modules should be avoided. Better choose
    :class:`.DecentTarEntry`
    """

    source: IOGenerator           # source of the data
    binary_info: bytes            # original tar header (may be larger than tarfile.BLOCKSIZE when it is combined
    # from several extended entries)
    parsed_info: tarfile.TarInfo  # tar header

    def data(self) -> IOGenerator:
        """Read data related to this entry
        """
        counter = self.parsed_info.size

        while counter > 0:
            next_chunk = next(self.source)
            assert(len(next_chunk) == tarfile.BLOCKSIZE)
            counter -= tarfile.BLOCKSIZE
            yield next_chunk


class DecentTarEntry:
    """This class helps to read tar entry in a safe way.
    """

    def __init__(self, tar_info: tarfile.TarInfo, data_generator: IOGenerator) -> None:
        """Create new entry.

        :param tar_info: information about inner entry
        :param data_generator: content of inner entry that must be read before the request for the next entry.
        """
        self.__tar_info = tar_info
        self.__data_generator = data_generator
        self.__wasted = False

    def is_wasted(self) -> bool:
        """Return True if data fully read.
        """
        return self.__wasted

    def tar_info(self) -> tarfile.TarInfo:
        """Return information about inner entry
        """
        return self.__tar_info

    def data(self) -> IOGenerator:
        """Read data related to this entry
        """
        counter = self.__tar_info.size

        while counter > 0:
            next_chunk = next(self.__data_generator)
            assert(len(next_chunk) == tarfile.BLOCKSIZE)

            if counter >= tarfile.BLOCKSIZE:
                yield next_chunk
            else:
                yield next_chunk[:counter]
            counter -= tarfile.BLOCKSIZE

        self.__wasted = True


class _TarReader:
    """This class helps to read data from a tar archive. Since it scans through all the data it may be unacceptable
    slowly. Use it on you own risk
    """

    # noinspection PyMethodMayBeStatic
    def _is_regular_entry(self, entry: tarfile.TarInfo) -> bool:
        """Return True if entry is FS-representable entry like file, directory or some named socket.

        :param entry: entry to check
        """

        regular_entries = [
            tarfile.REGTYPE,
            tarfile.AREGTYPE,
            tarfile.LNKTYPE,
            tarfile.SYMTYPE,
            tarfile.CHRTYPE,
            tarfile.BLKTYPE,
            tarfile.DIRTYPE,
            tarfile.FIFOTYPE,

            # there is tarfile.CONTTYPE also, but what the fuck is this?!
        ]

        return entry.type in regular_entries

    # noinspection PyMethodMayBeStatic
    def _is_extended_head_entry(self, entry: tarfile.TarInfo) -> bool:
        """Return True if entry is a special one dedicated to overcome original TAR limits (like path limitations
        or size limitations)

        :param entry: entry to check
        """
        extended_head_types = [
            tarfile.GNUTYPE_LONGNAME,  # GNU format
            tarfile.GNUTYPE_LONGLINK,  # GNU format
            tarfile.GNUTYPE_SPARSE,    # GNU format
            tarfile.XHDTYPE            # PAX format

            # tarfile.XGLTYPE -- this is PAX format, but wtf?!
            # tarfile.SOLARIS_XHDTYPE -- some old shit
        ]

        return entry.type in extended_head_types

    # noinspection PyMethodMayBeStatic
    def _iterate_raw_entries(self, source: IOProducer) -> typing.Generator[_FilthyTarEntry, None, None]:
        """Iterate over all entries in source. Extended headers will be split and yielded independently
        """

        zero_block = b'\x00' * tarfile.BLOCKSIZE
        aligner = Aligner(tarfile.BLOCKSIZE, strict_mode=True)

        next_entry_size = 0
        aligner_generator = aligner.iterate_data(source)
        for chunk in aligner_generator:
            next_entry_size -= len(chunk)

            if chunk == zero_block:
                continue

            head = tarfile.TarInfo.frombuf(chunk, tarfile.ENCODING, 'strict')  # read the
            # https://docs.python.org/3/library/codecs.html#error-handlers
            yield _FilthyTarEntry(aligner_generator, chunk, head)

    def _iterate_extended_entries(self, source: IOProducer) -> typing.Generator[_FilthyTarEntry, None, None]:
        """Iterate over all entries in source. Extended headers will be combined and only "regular" entries will be
        yielded.
        """
        raw_entries_generator = self._iterate_raw_entries(source)
        for tar_entry in raw_entries_generator:
            tar_head = tar_entry.parsed_info

            if self._is_regular_entry(tar_head):
                yield tar_entry
                continue

            if self._is_extended_head_entry(tar_head):

                extended_head_data = tar_entry.binary_info
                for chunk in tar_entry.data():
                    extended_head_data += chunk

                next_entry = next(raw_entries_generator)
                while not self._is_regular_entry(next_entry.parsed_info):
                    if not self._is_extended_head_entry(next_entry.parsed_info):
                        raise ValueError(
                            f'Unknown tar entry spotted -- {str(tar_head.type)} (file -- {str(tar_head.name)})'
                        )

                    extended_head_data += next_entry.binary_info
                    for chunk in next_entry.data():
                        extended_head_data += chunk

                    next_entry = next(raw_entries_generator)

                extended_head_data += next_entry.binary_info

                combined_head = tarfile.TarFile(fileobj=io.BytesIO(extended_head_data)).next()
                assert(combined_head is not None)
                yield _FilthyTarEntry(next_entry.source, extended_head_data, combined_head)
                continue

            raise ValueError(f'Unknown tar entry spotted -- {str(tar_head.type)} (file -- {str(tar_head.name)})')

    def iterate_entries(self, data: IOProducer) -> typing.Generator[DecentTarEntry, None, None]:
        """Yield information about inner files
        """
        aligner = Aligner(tarfile.BLOCKSIZE, strict_mode=True)

        for filthy_entry in self._iterate_extended_entries(aligner.iterate_data(data)):
            previous_entry = DecentTarEntry(filthy_entry.parsed_info, filthy_entry.source)
            yield previous_entry

            if not previous_entry.is_wasted():
                raise RuntimeError('You must read data from previous entry before requesting a next one')


class TarArchive(_TarPaddingGenerator):
    """This class helps to save a single tar archive or helps to read it
    """

    def __init__(self) -> None:
        """Create a new tar archive writer
        """
        _TarPaddingGenerator.__init__(self, aligned_to=tarfile.RECORDSIZE, extra_padding=True)

    def static_archive(self, sources: typing.Iterable[StaticTarEntryProto]) -> IOGenerator:
        """Yield data that may be saved as a tar archive.

        :param sources: archive entries
        """

        def entries() -> IOGenerator:
            for s in sources:
                yield from s.entry()

        yield from self._write(entries())

    @verify_value(destination=lambda x: x.seekable())
    async def dynamic_archive(
        self,
        destination: typing.IO[bytes],
        sources: typing.Iterable[typing.Union[StaticTarEntryProto, DynamicTarEntryProto]],
        write_throttling: typing.Optional[int] = None
    ) -> None:
        """Save data as a tar archive.

        :param destination: destination file
        :param sources: archive entries
        :param write_throttling: whether to throttle at saving (number of bytes per second)
        """
        destination.seek(0)
        destination.truncate()

        def entries() -> IOGenerator:
            for s in sources:
                if isinstance(s, StaticTarEntryProto):
                    yield from s.entry()
                else:
                    assert(isinstance(s, DynamicTarEntryProto))
                    yield from s.entry(destination)

        await cag(IOThrottler.async_writer(self._write(entries()), destination, throttling=write_throttling))

    @staticmethod
    @verify_value(source=lambda x: x.seekable())
    def extract(source: typing.IO[bytes], filename: str) -> IOGenerator:
        """Extract data from an archive

        :param source: tar-data
        :param filename: name of a file to retrieve from archive
        """

        start_pos = source.tell()

        with tarfile.open(fileobj=source) as tar:
            tar_info = tar.getmember(filename)
            file_size = tar_info.size
            source.seek(start_pos + tar_info.offset_data, os.SEEK_SET)

            if file_size:  # not empty file
                yield from IOThrottler.sync_reader(source, read_size=file_size)

        source.seek(start_pos, os.SEEK_SET)

    @staticmethod
    def entries(source: IOProducer) -> typing.Generator[DecentTarEntry, None, None]:
        """Read data and yield information about inner files. Entries are returned one by one. Each entry must be
        read with the :meth:`.DecentTarEntry.data` method before the request for the next entry.

        :param source: tar data to read
        """
        yield from _TarReader().iterate_entries(source)

    @staticmethod
    def inner_descriptors(source: IOProducer) -> typing.Generator[tarfile.TarInfo, None, None]:
        """Yield information about inner files. Entries are returned one by one

        :param source: tar data to read
        """

        for i in _TarReader().iterate_entries(source):
            yield i.tar_info()
            for _ in i.data():
                pass
