"""
Unpack conda packages without using a temporary file.
"""

from __future__ import annotations

import bz2
import os
import os.path
import tarfile
import zipfile
from collections.abc import Generator
from enum import Enum

# acquire umask taking advantage of import system lock, instead of possibly in
# multiple threads at once.
UMASK = os.umask(0)
os.umask(UMASK)

try:
    import compression.zstd as zstd  # Python 3.14+
except ImportError:
    try:
        import backports.zstd as zstd
    except ImportError:
        import warnings

        warnings.warn(
            "zstd module could not be imported. Running without .conda support."
        )

        zstd = None


class CondaComponent(Enum):
    pkg = "pkg"
    info = "info"

    def __str__(self):
        return self.value


class TarfileNoSameOwner(tarfile.TarFile):
    def __init__(self, *args, umask: int | None = None, **kwargs):
        """Open an (uncompressed) tar archive `name'. `mode' is either 'r' to
        read from an existing archive, 'a' to append data to an existing
        file or 'w' to create a new file overwriting an existing one. `mode'
        defaults to 'r'.
        If `fileobj' is given, it is used for reading or writing data. If it
        can be determined, `mode' is overridden by `fileobj's mode.
        `fileobj' is not closed, when TarFile is closed.
        """
        super().__init__(*args, **kwargs)
        self.umask = umask if umask is not None else UMASK

    def chown(self, tarinfo, targetpath, numeric_owner):
        """
        Override chown to be a no-op, since we don't want to preserve ownership
        here. (tarfile.TarFile only lets us toggle all of (chown, chmod, mtime))
        """
        return

    def chmod(self, tarinfo, targetpath):
        """
        Set file permissions of targetpath according to tarinfo, respecting
        umask.
        """
        if tarinfo.mode is None:
            return

        try:
            os.chmod(targetpath, tarinfo.mode & (-1 & (~self.umask)))
        except OSError as e:
            raise tarfile.ExtractError("could not change mode") from e


def tar_generator(
    fileobj, tarfile_open=TarfileNoSameOwner.open, closefd=False, *, encoding="utf-8"
) -> Generator[tuple[tarfile.TarFile, tarfile.TarInfo]]:
    """
    Yield (tar, member) from fileobj.

    Args:
        fileobj: file-like object

        encoding: "utf-8" passed to TarFile.open(); can be changed for testing.
    """
    # tarfile will not close fileobj because _extfileobj is True
    # caller should take care to close files all the way back to the http request...
    try:
        with tarfile_open(fileobj=fileobj, mode="r|", encoding=encoding) as tar:
            for member in tar:
                yield tar, member
    finally:
        if closefd:
            fileobj.close()


def stream_conda_info(
    filename, fileobj=None
) -> Generator[tuple[tarfile.TarFile, tarfile.TarInfo]]:
    """
    Yield members from conda's embedded info/ tarball.

    For .tar.bz2 packages, yield all members.

    Yields (tar, member) tuples. You must only use the current member to
    prevent tar seeks and scans.

    To extract to disk, it's possible to call ``tar.extractall(path)`` on the
    first result and then ignore the rest of this generator. ``extractall`` takes
    care of some directory permissions/mtime issues, compared to ``extract`` or
    writing out the file objects yourself.
    """
    component = "info"
    return stream_conda_component(filename, fileobj, component)


def stream_conda_component(
    filename,
    fileobj=None,
    component: CondaComponent | str = CondaComponent.pkg,
    *,
    encoding="utf-8",
    zf: zipfile.ZipFile | None = None,
) -> Generator[tuple[tarfile.TarFile, tarfile.TarInfo]]:
    """
    Yield members from .conda's embedded {component}- tarball. "info" or "pkg".

    For .tar.bz2 packages, yield all members.

    Yields (tar, member) tuples. You must only use the current member to
    prevent tar seeks and scans.

    To extract to disk, it's possible to call ``tar.extractall(path)`` on the
    first result and then ignore the rest of this generator. ``extractall`` takes
    care of some directory permissions/mtime issues, compared to ``extract`` or
    writing out the file objects yourself.

    Args:
        filename: path or file-like object to the .conda or .tar.bz2 file

        component: "pkg" or "info" component to extract

        encoding: "utf-8" passed to TarFile.open(); can be changed for testing.

        zf: optional pre-opened ``zipfile.ZipFile`` for .conda archives.
            When provided, reuses the already-opened ZipFile instead of
            parsing the central directory again. Callers that stream both
            the ``pkg`` and ``info`` components can open the zip once and
            pass it in both times, avoiding duplicate parsing. See #173.
    """
    if str(filename).endswith(".conda"):
        if zstd is None:
            raise RuntimeError(
                "Cannot unpack `.conda` without `backports.zstd` or `compression.zstd`"
            )

        if zf is None:
            zf = zipfile.ZipFile(fileobj or filename)
        stem, _, _ = os.path.basename(filename).rpartition(".")
        component_name = f"{component}-{stem}"
        component_filename = [
            info for info in zf.infolist() if info.filename.startswith(component_name)
        ]
        if not component_filename:
            raise LookupError(f"didn't find {component_name} component in {filename}")
        assert len(component_filename) == 1
        reader = zstd.open(zf.open(component_filename[0]))
    elif str(filename).endswith(".tar.bz2"):
        reader = bz2.open(fileobj or filename, mode="rb")
    else:
        raise ValueError("unsupported file extension")
    return tar_generator(reader, closefd=fileobj is None, encoding=encoding)