"""Tarball venv building utilities for registry packages.

This module provides functionality to build compressed tarball venvs from
registry packages for upload to S3/MinIO. Tarballs are used by executors
to install and execute registry actions.
"""

from __future__ import annotations

import asyncio
import hashlib
import os
import re
import shutil
import stat
import subprocess
import sysconfig
import tarfile
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING
from urllib.parse import urlparse

import aiofiles
import tracecat_registry

from tracecat import config
from tracecat.logger import logger
from tracecat.storage import blob

if TYPE_CHECKING:
    from tracecat.ssh import SshEnv


class TarballBuildError(Exception):
    """Raised when tarball building fails."""


@dataclass
class TarballVenvBuildResult:
    """Result of building a compressed tarball venv."""

    tarball_path: Path
    tarball_name: str
    content_hash: str
    compressed_size_bytes: int
    squashfs_path: Path | None = None
    squashfs_size_bytes: int | None = None


def _squashfs_path_for(tarball_path: Path) -> Path:
    """Return the sibling SquashFS path for a gzip tarball path."""
    if tarball_path.name.endswith(".tar.gz"):
        squashfs_name = tarball_path.name.removesuffix(".tar.gz") + ".squashfs"
        return tarball_path.with_name(squashfs_name)
    return tarball_path.with_name(f"{tarball_path.name}.squashfs")


def _squashfs_key_for(key: str) -> str:
    """Return the sibling SquashFS key for a gzip tarball S3 key."""
    if key.endswith(".tar.gz"):
        return key.removesuffix(".tar.gz") + ".squashfs"
    return f"{key}.squashfs"


def parse_s3_uri(uri: str) -> tuple[str, str]:
    """Parse an S3 URI into bucket and key."""
    parsed = urlparse(uri)
    if parsed.scheme != "s3" or not parsed.netloc or not parsed.path.strip("/"):
        raise ValueError(f"Invalid S3 URI: {uri}")
    return parsed.netloc, parsed.path.lstrip("/")


def get_squashfs_sidecar_key(tarball_key: str) -> str:
    """Return the sibling SquashFS key for a tarball S3 key."""
    return _squashfs_key_for(tarball_key)


def _compute_file_hash(file_path: Path) -> str:
    """Compute SHA256 hash of a file."""
    sha256_hash = hashlib.sha256()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            sha256_hash.update(chunk)
    return sha256_hash.hexdigest()


def _copy_squashfs_entry(
    path: Path,
    dest: Path,
    *,
    preserve_symlinks: bool,
) -> None:
    """Stage one artifact entry for SquashFS."""

    def _normalized_file_mode(mode: int) -> int:
        mode = stat.S_IMODE(mode)
        normalized = mode | 0o444
        if mode & 0o111:
            normalized |= 0o111
        return normalized

    def _normalized_dir_mode(mode: int) -> int:
        return stat.S_IMODE(mode) | 0o555

    def _remove_existing_entry() -> None:
        if dest.is_symlink() or dest.is_file():
            dest.unlink()
        elif dest.is_dir():
            shutil.rmtree(dest)

    if path.is_symlink():
        if not preserve_symlinks:
            logger.info(
                "Skipping link entry while staging SquashFS",
                path=str(path),
                link_target=os.readlink(path),
            )
            return
        dest.parent.mkdir(parents=True, exist_ok=True)
        _remove_existing_entry()
        dest.symlink_to(os.readlink(path))
        return

    if path.is_dir():
        if dest.is_symlink() or dest.is_file():
            dest.unlink()
        dest.mkdir(parents=True, exist_ok=True)
        for child in path.iterdir():
            _copy_squashfs_entry(
                child,
                dest / child.name,
                preserve_symlinks=preserve_symlinks,
            )
        dest.chmod(_normalized_dir_mode(path.stat().st_mode))
        return

    if path.is_file():
        dest.parent.mkdir(parents=True, exist_ok=True)
        _remove_existing_entry()
        source_mode = stat.S_IMODE(path.stat().st_mode)
        normalized_mode = _normalized_file_mode(source_mode)
        if normalized_mode == source_mode:
            try:
                os.link(path, dest)
            except OSError:
                shutil.copy2(path, dest)
        else:
            shutil.copy2(path, dest)
            dest.chmod(normalized_mode)
        return

    logger.info("Skipping non-file entry while staging SquashFS", path=str(path))


def _create_squashfs_image(
    squashfs_path: Path,
    entries: list[tuple[Path, str]],
    *,
    preserve_symlinks: bool = True,
) -> bool:
    """Create a SquashFS image from path entries when mksquashfs is available."""
    if not config.TRACECAT__REGISTRY_SYNC_SQUASHFS_ENABLED:
        return False

    mksquashfs = shutil.which("mksquashfs")
    if mksquashfs is None:
        logger.warning("Skipping SquashFS sidecar build; mksquashfs is not installed")
        return False

    with tempfile.TemporaryDirectory(prefix="tracecat_squashfs_root_") as staging:
        staging_dir = Path(staging)
        # TemporaryDirectory creates a 0700 root, which would make the mounted
        # SquashFS unreadable to non-root executor processes after -all-root.
        staging_dir.chmod(0o755)
        for path, arcname in entries:
            _copy_squashfs_entry(
                path,
                staging_dir / arcname,
                preserve_symlinks=preserve_symlinks,
            )

        cmd = [
            mksquashfs,
            str(staging_dir),
            str(squashfs_path),
            "-noappend",
            "-comp",
            "gzip",
            "-no-xattrs",
            "-all-root",
            "-processors",
            str(config.TRACECAT__REGISTRY_SYNC_SQUASHFS_PROCESSORS),
            "-mem",
            config.TRACECAT__REGISTRY_SYNC_SQUASHFS_MEM,
        ]
        result = subprocess_run(cmd)
        if result.returncode != 0:
            error = result.stderr.strip() or result.stdout.strip()
            raise TarballBuildError(f"Failed to build SquashFS sidecar: {error}")

    return True


def _create_squashfs_sidecar(
    squashfs_path: Path,
    entries: list[tuple[Path, str]],
    *,
    preserve_symlinks: bool = True,
) -> None:
    """Create a SquashFS sidecar when enabled, otherwise leave no file behind."""
    if not _create_squashfs_image(
        squashfs_path,
        entries,
        preserve_symlinks=preserve_symlinks,
    ):
        squashfs_path.unlink(missing_ok=True)


def subprocess_run(cmd: list[str]) -> subprocess.CompletedProcess[str]:
    """Run a subprocess command.

    Kept as a small wrapper so tests can monkeypatch SquashFS creation without
    shelling out.
    """
    return subprocess.run(cmd, capture_output=True, text=True, check=False)


def _slugify_origin(origin: str) -> str:
    """Convert a repository origin to a safe slug for S3 keys."""
    # Remove protocol prefix
    slug = (
        origin.replace("git+ssh://", "").replace("https://", "").replace("http://", "")
    )
    # Replace non-alphanumeric characters with underscores
    slug = re.sub(r"[^a-zA-Z0-9_-]", "_", slug)
    # Remove consecutive underscores
    slug = re.sub(r"_+", "_", slug)
    # Remove leading/trailing underscores
    slug = slug.strip("_")
    return slug[:100]  # Limit length


def get_builtin_registry_source_path() -> Path:
    """Get the path to the builtin tracecat_registry package source.

    Checks in order:
    1. Config path (TRACECAT__BUILTIN_REGISTRY_SOURCE_PATH)
    2. Relative to installed package (editable/source install)

    Returns:
        Path to the package directory containing pyproject.toml.

    Raises:
        TarballBuildError: If package path cannot be determined.
    """
    # Check configured path first (Docker default: /app/packages/tracecat-registry)
    config_path = Path(config.TRACECAT__BUILTIN_REGISTRY_SOURCE_PATH)
    if (config_path / "pyproject.toml").exists():
        return config_path

    # Fall back to development: look relative to installed package
    package_path = Path(tracecat_registry.__file__).parent.parent
    pyproject = package_path / "pyproject.toml"
    if not pyproject.exists():
        package_path = package_path.parent
        pyproject = package_path / "pyproject.toml"

    if not pyproject.exists():
        raise TarballBuildError(
            "Cannot find pyproject.toml for tracecat_registry. "
            "Set TRACECAT__BUILTIN_REGISTRY_SOURCE_PATH or use source installation."
        )

    return package_path


def get_installed_site_packages_paths() -> list[Path]:
    """Get site-packages directories for the current interpreter.

    Returns both `purelib` and `platlib` when they differ.
    """
    site_packages_paths: list[Path] = []
    seen_paths: set[Path] = set()

    for path_name in ("purelib", "platlib"):
        site_packages_str = sysconfig.get_path(path_name)
        if site_packages_str is None:
            raise TarballBuildError(
                f"Could not resolve {path_name} path from current interpreter."
            )

        site_packages = Path(site_packages_str)
        if not site_packages.exists():
            raise TarballBuildError(
                f"Resolved {path_name} path does not exist: {site_packages}"
            )

        resolved_site_packages = site_packages.resolve()
        if resolved_site_packages in seen_paths:
            continue

        seen_paths.add(resolved_site_packages)
        site_packages_paths.append(site_packages)

    return site_packages_paths


async def build_tarball_venv_from_installed_environment(
    *,
    package_name: str,
    package_dir: Path,
    output_dir: Path | None = None,
) -> TarballVenvBuildResult:
    """Build a tarball from the current interpreter's installed environment.

    This avoids creating a fresh venv and re-installing dependencies from indexes.
    It's primarily used for builtin registry sync where dependencies are already
    available in the running container image.
    """

    if output_dir is None:
        output_dir = Path(tempfile.mkdtemp(prefix="tracecat_tarball_venv_"))
    output_dir.mkdir(parents=True, exist_ok=True)

    site_packages_paths = get_installed_site_packages_paths()
    package_dir = package_dir.resolve()
    package_in_site_packages = any(
        package_dir.parent == site_packages.resolve()
        for site_packages in site_packages_paths
    )
    existing_package_entries = [
        site_packages / package_name
        for site_packages in site_packages_paths
        if (site_packages / package_name).exists()
    ]
    package_site_entry = (
        existing_package_entries[0] if existing_package_entries else None
    )
    package_site_entry_is_symlink = (
        package_site_entry.is_symlink() if package_site_entry is not None else False
    )
    should_overlay_editable_package = not package_in_site_packages and (
        package_site_entry is None or package_site_entry_is_symlink
    )

    tarball_name = "site-packages.tar.gz"
    tarball_path = output_dir / tarball_name
    squashfs_path = _squashfs_path_for(tarball_path)

    logger.info(
        "Building tarball from installed environment",
        site_packages_paths=[str(path) for path in site_packages_paths],
        package_name=package_name,
        package_dir=str(package_dir),
        package_in_site_packages=package_in_site_packages,
        package_site_entry=str(package_site_entry) if package_site_entry else None,
        package_site_entry_is_symlink=package_site_entry_is_symlink,
    )

    def _filter_link_entries(member: tarfile.TarInfo) -> tarfile.TarInfo | None:
        if member.issym() or member.islnk():
            logger.info(
                "Skipping link entry while building installed environment tarball",
                member_name=member.name,
                link_target=member.linkname,
            )
            return None
        return member

    def _create_tarball() -> None:
        entries: list[tuple[Path, str]] = []
        for site_packages in site_packages_paths:
            entries.extend((item, item.name) for item in site_packages.iterdir())

        if not package_in_site_packages and should_overlay_editable_package:
            entries.append((package_dir, package_name))

        with tarfile.open(tarball_path, "w:gz", compresslevel=6) as tar:
            for path, arcname in entries:
                tar.add(path, arcname=arcname, filter=_filter_link_entries)

        _create_squashfs_sidecar(
            squashfs_path,
            entries,
            preserve_symlinks=False,
        )

        if not package_in_site_packages and not should_overlay_editable_package:
            logger.warning(
                "Skipping editable package overlay because package already exists in site-packages",
                package_name=package_name,
            )

    await asyncio.to_thread(_create_tarball)

    content_hash = await asyncio.to_thread(_compute_file_hash, tarball_path)
    compressed_size = tarball_path.stat().st_size
    squashfs_size = squashfs_path.stat().st_size if squashfs_path.exists() else None

    logger.info(
        "Installed environment tarball built successfully",
        tarball_name=tarball_name,
        content_hash=content_hash[:16],
        compressed_size_bytes=compressed_size,
        squashfs_size_bytes=squashfs_size,
    )

    return TarballVenvBuildResult(
        tarball_path=tarball_path,
        tarball_name=tarball_name,
        content_hash=content_hash,
        compressed_size_bytes=compressed_size,
        squashfs_path=squashfs_path if squashfs_path.exists() else None,
        squashfs_size_bytes=squashfs_size,
    )


async def build_tarball_venv_from_path(
    package_path: Path,
    output_dir: Path | None = None,
    python_version: str = "3.12",
) -> TarballVenvBuildResult:
    """Build a complete venv with all dependencies and compress as tarball.

    This creates a portable venv tarball that can be extracted and used
    directly without running pip install. Faster for deployment since
    extraction is quicker than package installation.

    Args:
        package_path: Path to the package directory (must contain pyproject.toml)
        output_dir: Directory to output the tarball (defaults to temp dir)
        python_version: Python version for the venv (default: 3.12)

    Returns:
        TarballVenvBuildResult with the tarball path and metadata

    Raises:
        TarballBuildError: If the build fails
    """

    if not package_path.exists():
        raise TarballBuildError(f"Package path does not exist: {package_path}")

    pyproject = package_path / "pyproject.toml"
    if not pyproject.exists():
        raise TarballBuildError(f"No pyproject.toml found in {package_path}")

    # Use temp dir if no output dir specified
    if output_dir is None:
        output_dir = Path(tempfile.mkdtemp(prefix="tracecat_tarball_venv_"))

    output_dir.mkdir(parents=True, exist_ok=True)

    # Create venv directory
    venv_dir = output_dir / "venv"
    venv_dir.mkdir(parents=True, exist_ok=True)

    logger.info(
        "Building tarball venv from package",
        package_path=str(package_path),
        output_dir=str(output_dir),
    )

    # Step 1: Create a minimal venv using uv
    venv_cmd = ["uv", "venv", str(venv_dir), "--python", python_version]
    process = await asyncio.create_subprocess_exec(
        *venv_cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        cwd=str(package_path),
    )
    _, stderr = await process.communicate()

    if process.returncode != 0:
        error_msg = stderr.decode().strip()
        logger.error("Failed to create venv", error=error_msg)
        raise TarballBuildError(f"Failed to create venv: {error_msg}")

    # Step 2: Install the package and all dependencies into the venv
    # Use --frozen to respect the lock file if present
    install_cmd = [
        "uv",
        "pip",
        "install",
        "--python",
        str(venv_dir / "bin" / "python"),
        str(package_path),
    ]
    process = await asyncio.create_subprocess_exec(
        *install_cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        cwd=str(package_path),
    )
    _, stderr = await process.communicate()

    if process.returncode != 0:
        error_msg = stderr.decode().strip()
        logger.error("Failed to install package into venv", error=error_msg)
        raise TarballBuildError(f"Failed to install package: {error_msg}")

    # Step 3: Extract just the site-packages directory (what we need for PYTHONPATH)
    site_packages = venv_dir / "lib" / f"python{python_version}" / "site-packages"
    if not site_packages.exists():
        # Try without minor version
        site_packages_candidates = list(
            (venv_dir / "lib").glob("python*/site-packages")
        )
        if site_packages_candidates:
            site_packages = site_packages_candidates[0]
        else:
            raise TarballBuildError(f"Could not find site-packages in {venv_dir}")

    # Step 4: Pre-compile Python bytecode for faster imports in sandbox
    # This avoids runtime compilation overhead when the tarball is extracted
    logger.info("Pre-compiling Python bytecode", site_packages=str(site_packages))
    compile_cmd = [
        str(venv_dir / "bin" / "python"),
        "-m",
        "compileall",
        "-q",  # Quiet mode
        "-f",  # Force recompile
        "-j",
        "0",  # Use all CPU cores
        str(site_packages),
    ]
    process = await asyncio.create_subprocess_exec(
        *compile_cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    await process.communicate()
    # Ignore errors - bytecode compilation is optional optimization

    # Step 5: Create compressed tarball of site-packages
    tarball_name = "site-packages.tar.gz"
    tarball_path = output_dir / tarball_name
    squashfs_path = _squashfs_path_for(tarball_path)

    logger.info(
        "Compressing site-packages to tarball",
        site_packages=str(site_packages),
        tarball_path=str(tarball_path),
    )

    # Run tar compression in a thread to not block async loop
    def _create_tarball() -> None:
        entries = [(item, item.name) for item in site_packages.iterdir()]
        with tarfile.open(tarball_path, "w:gz", compresslevel=6) as tar:
            # Add site-packages contents with relative paths
            for path, arcname in entries:
                tar.add(path, arcname=arcname)
        _create_squashfs_sidecar(squashfs_path, entries)

    await asyncio.to_thread(_create_tarball)

    # Step 6: Compute content hash
    content_hash = await asyncio.to_thread(_compute_file_hash, tarball_path)
    compressed_size = tarball_path.stat().st_size
    squashfs_size = squashfs_path.stat().st_size if squashfs_path.exists() else None

    logger.info(
        "Tarball venv built successfully",
        tarball_name=tarball_name,
        content_hash=content_hash[:16],
        compressed_size_bytes=compressed_size,
        squashfs_size_bytes=squashfs_size,
    )

    return TarballVenvBuildResult(
        tarball_path=tarball_path,
        tarball_name=tarball_name,
        content_hash=content_hash,
        compressed_size_bytes=compressed_size,
        squashfs_path=squashfs_path if squashfs_path.exists() else None,
        squashfs_size_bytes=squashfs_size,
    )


async def build_tarball_venv_from_git(
    git_url: str,
    commit_sha: str,
    env: SshEnv | None = None,
    output_dir: Path | None = None,
    python_version: str = "3.12",
) -> TarballVenvBuildResult:
    """Build a tarball venv from a git repository.

    Args:
        git_url: Git SSH URL (git+ssh://...)
        commit_sha: Commit SHA to checkout
        env: SSH environment for git operations
        output_dir: Directory to output the tarball
        python_version: Python version for the venv (default: 3.12)

    Returns:
        TarballVenvBuildResult with the tarball path and metadata

    Raises:
        TarballBuildError: If the build fails
    """
    if output_dir is None:
        output_dir = Path(tempfile.mkdtemp(prefix="tracecat_tarball_venv_"))

    output_dir.mkdir(parents=True, exist_ok=True)

    # Clone the repository to a temp directory
    async with aiofiles.tempfile.TemporaryDirectory(
        prefix="tracecat_git_"
    ) as clone_dir:
        clone_path = Path(clone_dir)

        # Strip git+ssh:// prefix for git clone
        clone_url = git_url.replace("git+ssh://", "ssh://")

        logger.info(
            "Cloning repository for tarball venv build",
            url=clone_url,
            commit_sha=commit_sha,
        )

        # Set up environment for git operations
        git_env = os.environ.copy()
        if env:
            git_env.update(env.to_dict())

        # Clone the repository
        clone_cmd = ["git", "clone", "--depth", "1", clone_url, str(clone_path)]
        process = await asyncio.create_subprocess_exec(
            *clone_cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=git_env,
        )
        _, stderr = await process.communicate()

        if process.returncode != 0:
            error_msg = stderr.decode().strip()
            raise TarballBuildError(f"Failed to clone repository: {error_msg}")

        # Fetch the specific commit
        fetch_cmd = ["git", "fetch", "origin", commit_sha]
        process = await asyncio.create_subprocess_exec(
            *fetch_cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=str(clone_path),
            env=git_env,
        )
        _, stderr = await process.communicate()

        if process.returncode != 0:
            error_msg = stderr.decode().strip()
            raise TarballBuildError(f"Failed to fetch commit: {error_msg}")

        # Checkout the specific commit
        checkout_cmd = ["git", "checkout", commit_sha]
        process = await asyncio.create_subprocess_exec(
            *checkout_cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=str(clone_path),
            env=git_env,
        )
        _, stderr = await process.communicate()

        if process.returncode != 0:
            error_msg = stderr.decode().strip()
            raise TarballBuildError(f"Failed to checkout commit: {error_msg}")

        # Build the tarball venv from the cloned repository
        return await build_tarball_venv_from_path(
            clone_path, output_dir, python_version
        )


async def build_builtin_registry_tarball_venv(
    output_dir: Path | None = None,
) -> TarballVenvBuildResult:
    """Build a tarball venv from the builtin tracecat_registry package.

    Args:
        output_dir: Directory to output the tarball

    Returns:
        TarballVenvBuildResult with the tarball path and metadata

    Raises:
        TarballBuildError: If the build fails
    """
    if config.TRACECAT__REGISTRY_SYNC_BUILTIN_USE_INSTALLED_SITE_PACKAGES:
        return await build_tarball_venv_from_installed_environment(
            package_name="tracecat_registry",
            package_dir=Path(tracecat_registry.__file__).resolve().parent,
            output_dir=output_dir,
        )

    package_path = get_builtin_registry_source_path()
    return await build_tarball_venv_from_path(package_path, output_dir)


async def build_squashfs_sidecar_from_tarball(
    *,
    tarball_path: Path,
    squashfs_path: Path,
    work_dir: Path,
) -> bool:
    """Build a SquashFS sidecar from an existing site-packages tarball."""

    def _build_sidecar() -> bool:
        if not tarball_path.exists():
            raise FileNotFoundError(f"Tarball file not found: {tarball_path}")

        extracted_dir = work_dir / "site-packages"
        extracted_dir.mkdir(parents=True, exist_ok=True)
        with tarfile.open(tarball_path, "r:gz") as tar:
            tar.extractall(path=extracted_dir, filter="data")

        entries = [(item, item.name) for item in extracted_dir.iterdir()]
        _create_squashfs_sidecar(
            squashfs_path,
            entries,
            preserve_symlinks=True,
        )
        return squashfs_path.exists()

    return await asyncio.to_thread(_build_sidecar)


def get_tarball_venv_s3_key(
    organization_id: str,
    repository_origin: str,
    version: str,
) -> str:
    """Generate the S3 key for a tarball venv file.

    Format: {org_id}/tarball-venvs/{origin_slug}/{version}/site-packages.tar.gz

    Args:
        organization_id: Organization UUID
        repository_origin: Repository origin (e.g., "tracecat_registry", "git+ssh://...")
        version: Version string

    Returns:
        S3 key string
    """
    origin_slug = _slugify_origin(repository_origin)
    return (
        f"{organization_id}/tarball-venvs/{origin_slug}/{version}/site-packages.tar.gz"
    )


async def upload_tarball_venv(
    tarball_path: Path,
    squashfs_path: Path | None,
    key: str,
    bucket: str,
) -> str:
    """Upload a tarball venv file to S3/MinIO.

    Args:
        tarball_path: Local path to the tarball file
        squashfs_path: Optional local path to the SquashFS sidecar
        key: The S3 object key
        bucket: Bucket name

    Returns:
        The S3 URI of the uploaded tarball (s3://{bucket}/{key})

    Raises:
        FileNotFoundError: If the tarball file doesn't exist
    """
    if not tarball_path.exists():
        raise FileNotFoundError(f"Tarball file not found: {tarball_path}")

    squashfs_key = _squashfs_key_for(key)
    if squashfs_path is not None and not squashfs_path.exists():
        raise FileNotFoundError(f"SquashFS file not found: {squashfs_path}")
    await blob.delete_file(key=squashfs_key, bucket=bucket)

    # Stream from disk via multipart upload so we don't hold hundreds of MB
    # of tarball bytes in the API process memory at once.
    await blob.upload_file_from_path(
        path=tarball_path,
        key=key,
        bucket=bucket,
        content_type="application/gzip",
    )

    if squashfs_path is not None:
        await blob.upload_file_from_path(
            path=squashfs_path,
            key=squashfs_key,
            bucket=bucket,
            content_type="application/vnd.squashfs",
        )

    s3_uri = f"s3://{bucket}/{key}"
    logger.info(
        "Tarball venv uploaded successfully",
        key=key,
        bucket=bucket,
        s3_uri=s3_uri,
        size=tarball_path.stat().st_size,
        squashfs_size=squashfs_path.stat().st_size if squashfs_path else None,
    )
    return s3_uri


async def download_tarball_venv(
    key: str,
    bucket: str,
    output_path: Path,
) -> Path:
    """Download a tarball venv file from S3/MinIO.

    Args:
        key: The S3 object key
        bucket: Bucket name
        output_path: Local path to save the tarball file

    Returns:
        The local path to the downloaded tarball
    """
    size_bytes = await blob.download_file_to_path(
        key=key,
        bucket=bucket,
        output_path=output_path,
    )

    logger.info(
        "Tarball venv downloaded successfully",
        key=key,
        bucket=bucket,
        output_path=str(output_path),
        size=size_bytes,
    )
    return output_path
