#!/usr/bin/python3
# Copyright (C) 2025 itachi_re <xanbenson99@gmail.com>
# ZypperX - Safe Parallel Wrapper for Zypper
# v1.1.2 - Fixed argparse conflicts, /var deletion bug, mount cleanup, and more

import os
import sys
import asyncio
import shutil
import tempfile
import subprocess
import argparse
import logging
import fcntl
import re
import signal
import xml.etree.ElementTree as ET
from typing import List, Tuple, Optional
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn, MofNCompleteColumn
from rich.logging import RichHandler
from rich.prompt import Confirm
from rich.panel import Panel

# ─── Constants ────────────────────────────────────────────────────────────────
ZYPPER_PID_FILE = "/run/zypp.pid"
ZYPPER_ENV = {"ZYPP_CURL2": "1", "ZYPP_PCK_PRELOAD": "1", "ZYPP_SINGLE_RPMTRANS": "1"}
VERSION = "1.1.1"

# CA cert directories that need to be bind-mounted for SSL to work in chroot
CERT_DIRS = [
    "/etc/pki",
    "/etc/ca-certificates",
    "/usr/share/ca-certificates",
]


class ExitCode:
    SUCCESS = 0
    ERROR = 1
    LOCKED = 5
    PERMISSION_DENIED = 126


console = Console()
logging.basicConfig(
    level="INFO",
    format="%(message)s",
    datefmt="[%X]",
    handlers=[RichHandler(console=console, markup=True, show_path=False)],
)
log = logging.getLogger("zypperx")


# ─── Mount Safety Helpers ─────────────────────────────────────────────────────

def get_active_mounts_under(path: str) -> List[str]:
    """
    Return all active mount points that are under `path`, sorted deepest-first
    so they can be safely unmounted in order.
    """
    path = os.path.realpath(path)
    active: List[str] = []
    try:
        with open("/proc/mounts", "r") as f:
            for line in f:
                parts = line.split()
                if len(parts) >= 2:
                    mp = parts[1]
                    if mp == path or mp.startswith(path + "/"):
                        active.append(mp)
    except IOError:
        pass
    # Deepest paths first — required for nested unmounts
    return sorted(active, key=lambda x: x.count("/"), reverse=True)


def safe_rmtree(path: str) -> bool:
    """
    Remove a directory tree ONLY after confirming no bind/tmpfs mounts remain
    under it. Refuses to delete if any are still active.

    BUG FIX: The original code called shutil.rmtree() without checking whether
    bind mounts were still present. If a bind mount of /var/cache/zypp or
    /var/lib/zypp happened to be attached (e.g. unshare namespace cleanup
    lagged, or an older code path used host-namespace mounts), rmtree would
    follow the mount and delete the REAL /var/cache/zypp or /var/lib/zypp.
    This is the most likely cause of the /var deletion incident.
    """
    if not os.path.exists(path):
        return True
    active = get_active_mounts_under(path)
    if active:
        log.error(f"[bold red]REFUSING to delete {path}: active mounts remain![/bold red]")
        log.error("Manual cleanup required:")
        for mp in active:
            log.error(f"  sudo umount -l '{mp}'")
        log.error(f"  sudo rm -rf '{path}'")
        return False
    shutil.rmtree(path, ignore_errors=True)
    return True


# ─── Safety Checker ──────────────────────────────────────────────────────────

class SafetyChecker:

    @staticmethod
    def verify_no_ghost_mounts(temp_prefix: str = "zypperx_") -> bool:
        """Abort if leftover mounts from a previous crashed run are present."""
        try:
            with open("/proc/mounts", "r") as f:
                content = f.read()
            if temp_prefix in content:
                console.print("[bold red]DANGER: Leftover mounts from a previous run detected![/bold red]")
                console.print("[yellow]Clean them up first:[/yellow]")
                console.print(
                    f"  grep '{temp_prefix}' /proc/mounts | awk '{{print $2}}' "
                    f"| sort -r | xargs -r sudo umount -l"
                )
                console.print(f"  sudo rm -rf /tmp/{temp_prefix}*")
                return False
        except IOError:
            pass
        return True

    @staticmethod
    def verify_filesystem_health() -> bool:
        for path in ["/var", "/var/cache", "/etc", "/tmp"]:
            if not os.path.exists(path) or not os.access(path, os.R_OK | os.W_OK):
                console.print(f"[bold red]CRITICAL: {path} is not accessible![/bold red]")
                return False
        return True

    @staticmethod
    def check_unshare_support() -> bool:
        """
        BUG FIX: The original check only ran 'unshare --version' which always
        succeeds even if mount namespaces are disabled by the kernel. Now we
        actually try to create a mount namespace to confirm it works.
        """
        try:
            result = subprocess.run(
                ["unshare", "--mount", "--propagation", "private", "true"],
                capture_output=True,
                timeout=5,
            )
            if result.returncode != 0:
                raise RuntimeError(result.stderr.decode().strip())
            return True
        except FileNotFoundError:
            console.print("[bold red]ERROR: 'unshare' not found.[/bold red]")
            console.print("  Install: sudo zypper install util-linux")
            return False
        except Exception as e:
            console.print(f"[bold red]ERROR: mount namespaces not functional: {e}[/bold red]")
            console.print(
                "  This is required for safe isolation. "
                "Check /proc/sys/kernel/unprivileged_userns_clone"
            )
            return False


# ─── Namespaced Worker ───────────────────────────────────────────────────────

class NamespacedWorker:
    def __init__(self, worker_id: str, task_type: str, task_item: str, force: bool = False):
        self.worker_id = worker_id
        self.task_type = task_type
        self.task_item = task_item
        self.force = force
        self.workspace: Optional[str] = None

    async def execute(self) -> Tuple[int, str]:
        self.workspace = tempfile.mkdtemp(prefix=f"zypperx_worker_{self.worker_id}_")
        try:
            # BUG FIX: Build zypper_args without empty strings.
            # Original: ["refresh", "--force" if self.force else "", item]
            # The empty string was filtered out later, but it's cleaner to never add it.
            if self.task_type == "refresh":
                zypper_args = ["refresh"]
                if self.force:
                    zypper_args.append("--force")
                zypper_args.append(self.task_item)
            elif self.task_type == "download":
                zypper_args = ["download", self.task_item]
            else:
                return 1, f"Unknown task type: {self.task_type!r}"

            cmd = [
                "unshare", "--mount", "--propagation", "private",
                "sh", "-c", self._build_isolated_script(zypper_args),
            ]

            proc = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=self._build_env(),
            )
            stdout, stderr = await proc.communicate()

            if proc.returncode != 0:
                # Combine stdout + stderr — zypper writes its human-readable
                # error (including the curl error code) to stdout, not stderr.
                out = stdout.decode("utf-8", errors="ignore").strip()
                err = stderr.decode("utf-8", errors="ignore").strip()
                combined = "\n".join(filter(None, [out, err]))
                return proc.returncode, combined
            return 0, "Success"

        finally:
            await self._cleanup()

    async def _cleanup(self):
        """
        Safely clean up the workspace after the worker finishes.

        BUG FIX: The original cleanup used a hardcoded list of mount points
        that was both INCOMPLETE (missing /etc/zypp, /etc/ca-certificates,
        /usr/share/ca-certificates) and in the WRONG ORDER (parent mounts like
        /etc were listed after their children /etc/ssl, /etc/pki, but that's
        only coincidentally correct — the list was not systematically ordered).

        New approach: read /proc/mounts to find what's actually mounted and
        sort deepest-first. Then use safe_rmtree() which re-checks /proc/mounts
        before deleting — so if unmounting fails for any reason, we refuse to
        rmtree rather than risk deleting real data.
        """
        if not self.workspace or not os.path.exists(self.workspace):
            return

        # Give the kernel a moment to finalize private-namespace cleanup
        await asyncio.sleep(0.05)

        active = get_active_mounts_under(self.workspace)
        if active:
            log.warning(
                f"Worker {self.worker_id}: "
                f"{len(active)} lingering mount(s) found (namespace cleanup lag?) — unmounting..."
            )
            for mp in active:  # already sorted deepest-first
                result = subprocess.run(
                    ["umount", "-l", mp],
                    capture_output=True,
                    timeout=10,
                    check=False,
                )
                if result.returncode != 0:
                    log.error(
                        f"  Failed to unmount {mp}: "
                        f"{result.stderr.decode(errors='ignore').strip()}"
                    )

        # safe_rmtree re-checks /proc/mounts before deleting
        if not safe_rmtree(self.workspace):
            log.error(
                f"Workspace {self.workspace} was NOT deleted. "
                "Please clean up manually (see above)."
            )

    def _build_isolated_script(self, zypper_args: List[str]) -> str:
        """
        Build the shell script that runs inside the private mount namespace.

        Key safety improvements vs original:
        - 'set -euo pipefail' instead of just 'set -e' (catches unset vars too)
        - Individual mkdir calls instead of brace expansion (safer if CHROOT is weird)
        - /var/lib/zypp is now mounted READ-ONLY — workers only need to read
          the package DB for dependency resolution, not write to it
        - /etc/zypp is explicitly mounted and listed for cleanup
        - All cert dirs handled uniformly
        - Uses 'realpath' for symlinked dirs (/bin, /sbin, /lib, /lib64 on Tumbleweed)
        """
        cert_lines = "\n".join(
            f'if [ -d "{d}" ]; then\n'
            f'    mkdir -p "$CHROOT{d}"\n'
            f'    mount --bind -o ro "{d}" "$CHROOT{d}"\n'
            f'fi'
            for d in CERT_DIRS
        )

        script = f"""
set -euo pipefail
CHROOT="{self.workspace}/chroot"

# Create directory skeleton individually (safer than brace expansion)
mkdir -p "$CHROOT/dev"
mkdir -p "$CHROOT/proc"
mkdir -p "$CHROOT/sys"
mkdir -p "$CHROOT/tmp"
mkdir -p "$CHROOT/run"
mkdir -p "$CHROOT/var/cache/zypp"
mkdir -p "$CHROOT/var/lib/zypp"
mkdir -p "$CHROOT/etc"
mkdir -p "$CHROOT/usr"
mkdir -p "$CHROOT/bin"
mkdir -p "$CHROOT/sbin"
mkdir -p "$CHROOT/lib"
mkdir -p "$CHROOT/lib64"

# ── CA bundle copy — must happen BEFORE any bind-mounts ─────────────────────
# /etc/ssl and /var/lib/ca-certificates are plain empty directories at this
# point (created by mkdir above). We resolve the real CA bundle file on the
# host (following all symlinks) and copy the actual bytes in. This avoids
# the read-only-filesystem error that occurs if we try to write after
# bind-mounting /etc/ssl, and avoids the dangling-symlink error that occurs
# because ca-bundle.pem -> /var/lib/ca-certificates/ca-bundle.pem is an
# absolute symlink that points outside the chroot.
mkdir -p "$CHROOT/etc/ssl/certs"
mkdir -p "$CHROOT/var/lib/ca-certificates"
_ca_copied=0
for _candidate in \\
    /var/lib/ca-certificates/ca-bundle.pem \\
    /etc/ssl/ca-bundle.pem \\
    /etc/ssl/certs/ca-certificates.crt \\
    /etc/pki/tls/certs/ca-bundle.crt \\
    /etc/ssl/cert.pem; do
    if [ -e "$_candidate" ]; then
        _real=$(realpath "$_candidate" 2>/dev/null || echo "$_candidate")
        if [ -f "$_real" ]; then
            # Copy the real bundle to every path curl/openssl might look for
            cp "$_real" "$CHROOT/etc/ssl/ca-bundle.pem"
            cp "$_real" "$CHROOT/etc/ssl/certs/ca-certificates.crt"
            cp "$_real" "$CHROOT/var/lib/ca-certificates/ca-bundle.pem"
            _ca_copied=1
            break
        fi
    fi
done
[ "$_ca_copied" = "0" ] && echo "WARNING: no CA bundle found on host" >&2

# /etc/ssl/certs/ — bind-mount the host hash directory on top of what we
# just created. OpenSSL uses this directory for certificate chain lookup
# (hashed symlinks like ab12cd34.0 -> individual CA certs). Without it,
# intermediate certificate verification fails with error 20 even when the
# root bundle is present. We mount only this subdirectory, not /etc/ssl
# itself, to avoid overwriting the ca-bundle.pem we just copied above.
if [ -d /etc/ssl/certs ]; then
    mount --bind -o ro /etc/ssl/certs "$CHROOT/etc/ssl/certs"
fi

# Ephemeral tmpfs for /tmp and /run (no host data leaks)
mount -t tmpfs tmpfs "$CHROOT/tmp"
mount -t tmpfs tmpfs "$CHROOT/run"

# Read-only system directories
mount --bind -o ro /dev   "$CHROOT/dev"
mount --bind -o ro /proc  "$CHROOT/proc"
mount --bind -o ro /sys   "$CHROOT/sys"

# Bind OS dirs; use realpath to handle /bin -> /usr/bin symlinks on Tumbleweed
for _d in /usr /bin /sbin /lib /lib64; do
    _real=$(realpath "$_d" 2>/dev/null || echo "$_d")
    if [ -d "$_real" ]; then
        mount --bind -o ro "$_real" "$CHROOT$_d"
    fi
done

# /var/cache/zypp — read-WRITE so zypper can write downloaded packages here
if [ -d /var/cache/zypp ]; then
    mount --bind /var/cache/zypp "$CHROOT/var/cache/zypp"
fi

# /var/lib/zypp — read-ONLY: workers only need to READ the DB, never write
# BUG FIX: was previously mounted rw, creating unnecessary write risk
if [ -d /var/lib/zypp ]; then
    mount --bind -o ro /var/lib/zypp "$CHROOT/var/lib/zypp"
fi

# Copy essential network/OS config files (not bind-mounted, so no risk)
for _f in resolv.conf hosts nsswitch.conf os-release; do
    [ -e "/etc/$_f" ] && cp -L "/etc/$_f" "$CHROOT/etc/$_f" 2>/dev/null || true
done

# zypper config (read-only)
if [ -d /etc/zypp ]; then
    mkdir -p "$CHROOT/etc/zypp"
    mount --bind -o ro /etc/zypp "$CHROOT/etc/zypp"
fi

# Bind-mount remaining cert dirs (/etc/pki, /etc/ca-certificates,
# /usr/share/ca-certificates). /etc/ssl and /var/lib/ca-certificates
# were already populated above by direct file copy, so they are excluded.
{cert_lines}

# Execute zypper inside the isolated chroot.
# SSL_CERT_FILE and CURL_CA_BUNDLE tell libcurl/openssl exactly where the
# bundle is so they never fall back to compiled-in defaults that may not
# exist inside the chroot.
chroot "$CHROOT" env -i \\
    ZYPP_CURL2=1 ZYPP_PCK_PRELOAD=1 ZYPP_SINGLE_RPMTRANS=1 \\
    PATH=/usr/bin:/bin:/sbin HOME=/tmp LANG=C.UTF-8 \\
    SSL_CERT_FILE=/etc/ssl/ca-bundle.pem \\
    CURL_CA_BUNDLE=/etc/ssl/ca-bundle.pem \\
    zypper --non-interactive {" ".join(zypper_args)}
"""
        return script

    @staticmethod
    def _build_env() -> dict:
        return {
            "PATH": "/usr/bin:/bin:/usr/sbin:/sbin",
            "HOME": "/tmp",
            "LANG": "C.UTF-8",
        }


# ─── Engine ──────────────────────────────────────────────────────────────────

class ZypperXEngine:
    def __init__(self, jobs: int = 10, force: bool = False):
        self.jobs = jobs
        self.force = force
        self.lock_fd: Optional[int] = None
        self.sem: Optional[asyncio.Semaphore] = None  # created in async context

    def check_prerequisites(self):
        if os.geteuid() != 0:
            console.print("[bold red]Error: Must run as root (sudo).[/bold red]")
            sys.exit(ExitCode.ERROR)
        if not SafetyChecker.verify_no_ghost_mounts():
            sys.exit(ExitCode.ERROR)
        if not SafetyChecker.verify_filesystem_health():
            sys.exit(ExitCode.ERROR)
        if not SafetyChecker.check_unshare_support():
            sys.exit(ExitCode.ERROR)
        missing = [p for p in ["zypper", "mount", "umount", "chroot", "unshare"]
                   if shutil.which(p) is None]
        if missing:
            console.print(f"[bold red]Missing tools: {', '.join(missing)}[/bold red]")
            sys.exit(ExitCode.ERROR)

    def get_lock(self):
        try:
            self.lock_fd = os.open(ZYPPER_PID_FILE, os.O_CREAT | os.O_WRONLY, 0o644)
            fcntl.flock(self.lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
            os.ftruncate(self.lock_fd, 0)
            os.write(self.lock_fd, str(os.getpid()).encode())
        except BlockingIOError:
            console.print("[bold red]Zypper is already locked by another process.[/bold red]")
            sys.exit(ExitCode.LOCKED)
        except OSError as e:
            console.print(f"[bold red]Lock error: {e}[/bold red]")
            sys.exit(ExitCode.ERROR)

    def release_lock(self):
        if self.lock_fd is not None:
            try:
                fcntl.flock(self.lock_fd, fcntl.LOCK_UN)
                os.close(self.lock_fd)
            except Exception:
                pass
            finally:
                self.lock_fd = None

    @staticmethod
    def validate_input(item: str) -> bool:
        # Allow package names, repo aliases, and edition strings
        return bool(re.match(r"^[\w.\-+:@/]+$", item))

    async def execute_tasks(
        self, task_list: List[str], task_type: str, description: str
    ) -> int:
        if self.sem is None:
            self.sem = asyncio.Semaphore(self.jobs)

        valid = [t for t in task_list if self.validate_input(t)]
        invalid = [t for t in task_list if not self.validate_input(t)]
        if invalid:
            log.warning(f"Skipping {len(invalid)} item(s) with invalid characters: {invalid[:5]}")

        successes = 0
        failures: List[Tuple[str, str]] = []

        with Progress(
            SpinnerColumn(),
            TextColumn("[bold blue]{task.description}"),
            BarColumn(),
            MofNCompleteColumn(),
            console=console,
            transient=False,
        ) as progress:
            task_id = progress.add_task(description, total=len(valid))

            async def run_worker(idx: int, item: str):
                nonlocal successes
                async with self.sem:
                    w = NamespacedWorker(str(idx), task_type, item, self.force)
                    code, msg = await w.execute()
                    if code == 0:
                        successes += 1
                    else:
                        failures.append((item, msg))
                        log.error(f"Failed: {item} - {msg.splitlines()[0][:120]}")
                    progress.advance(task_id)

            await asyncio.gather(
                *[run_worker(i, item) for i, item in enumerate(valid)]
            )

        if failures:
            console.print(f"\n[yellow]⚠  {len(failures)} task(s) failed (see errors above)[/yellow]")

        console.print(
            f"[green]✓ {successes}/{len(valid)} tasks completed successfully[/green]"
        )
        return successes

    def get_enabled_repos(self) -> List[str]:
        console.print("[cyan]Scanning enabled repositories...[/cyan]")
        res = subprocess.run(
            ["zypper", "--non-interactive", "--no-cd", "--xmlout", "repos"],
            capture_output=True,
            text=True,
        )
        if res.returncode != 0:
            console.print("[bold red]Failed to list repositories.[/bold red]")
            sys.exit(ExitCode.ERROR)
        try:
            root = ET.fromstring(res.stdout)
            return [
                r.attrib["alias"]
                for r in root.iter("repo")
                if r.attrib.get("enabled") == "1"
            ]
        except ET.ParseError as e:
            console.print(f"[bold red]Failed to parse repo list: {e}[/bold red]")
            sys.exit(ExitCode.ERROR)

    def get_transaction_packages(
        self, mode: str, packages: Optional[List[str]] = None
    ) -> Tuple[List[str], int]:
        console.print(f"[cyan]Calculating transaction for {mode}...[/cyan]")
        base = ["zypper", "--non-interactive", "--no-cd", "--xmlout"]
        if mode == "dist-upgrade":
            cmd = base + ["dist-upgrade", "--dry-run"]
        elif mode == "install":
            cmd = base + ["install", "--dry-run"] + (packages or [])
        elif mode == "install-new-recommends":
            cmd = base + ["install-new-recommends", "--dry-run"]
        else:
            return [], 0

        res = subprocess.run(cmd, capture_output=True, text=True)
        if "Nothing to do" in res.stdout or res.returncode not in (0, 100, 101):
            if res.returncode not in (0, 100, 101):
                console.print(
                    f"[bold red]zypper dry-run failed (exit {res.returncode}):[/bold red]"
                )
                console.print(res.stderr[:500])
                sys.exit(ExitCode.ERROR)
            return [], 0

        try:
            root = ET.fromstring(res.stdout)
            summary = root.find(".//install-summary")
            download_size = (
                int(summary.attrib.get("download-size", 0))
                if summary is not None
                else 0
            )

            # BUG FIX: Original iterated ALL solvable nodes in the document,
            # which included nodes inside <to-remove> as well. Now we only
            # collect from <to-install> explicitly.
            to_install: set = set()
            for node in root.iter("to-install"):
                for solv in node.findall("solvable"):
                    if solv.attrib.get("type") == "package":
                        name = solv.attrib.get("name", "")
                        edition = solv.attrib.get("edition", "")
                        arch = solv.attrib.get("arch", "")
                        if name and edition and arch:
                            to_install.add(f"{name}-{edition}.{arch}")

            to_remove: set = set()
            for node in root.iter("to-remove"):
                for solv in node.findall("solvable"):
                    name = solv.attrib.get("name", "")
                    edition = solv.attrib.get("edition", "")
                    arch = solv.attrib.get("arch", "")
                    if name and edition and arch:
                        to_remove.add(f"{name}-{edition}.{arch}")

            return sorted(to_install - to_remove), download_size

        except Exception as e:
            console.print(f"[bold red]Failed to parse transaction XML: {e}[/bold red]")
            sys.exit(ExitCode.ERROR)


# ─── Signal / Emergency Cleanup ──────────────────────────────────────────────

def emergency_cleanup(engine: ZypperXEngine):
    """Best-effort cleanup of all zypperx workspaces and the zypp lock."""
    engine.release_lock()
    try:
        tmp = "/tmp"
        for entry in os.listdir(tmp):
            if entry.startswith("zypperx_"):
                path = os.path.join(tmp, entry)
                active = get_active_mounts_under(path)
                for mp in active:
                    subprocess.run(["umount", "-l", mp], capture_output=True, check=False)
                safe_rmtree(path)
    except Exception as exc:
        log.warning(f"Emergency cleanup warning: {exc}")


# ─── Main ─────────────────────────────────────────────────────────────────────

async def main():
    # ── Argument Parser ──────────────────────────────────────────────────────
    #
    # BUG FIX: The original code defined -f/--force and -d/--download-only on
    # BOTH the root parser and the subparsers. argparse detects this as a
    # conflict and raises:
    #   "Fatal error: argument -f/--force: conflicting option strings: -f, --force"
    #
    # Fix: Remove -f and -d from the root parser entirely. They only make sense
    # on the specific subcommands that use them anyway.
    #
    parser = argparse.ArgumentParser(
        prog="zypperx",
        description=f"ZypperX v{VERSION} — Safe parallel accelerator for zypper",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Examples:\n"
            "  sudo zypperx refresh             # cleans cache then refreshes all repos\n"
            "  sudo zypperx ref --no-clean      # refresh without clearing cache first\n"
            "  sudo zypperx ref -f              # force-refresh all repos\n"
            "  sudo zypperx dup -j 20           # parallel dist-upgrade, 20 workers\n"
            "  sudo zypperx in firefox -y       # install firefox, no confirmation\n"
        ),
    )
    parser.add_argument(
        "-j", "--jobs", type=int, default=10, metavar="N",
        help="Number of parallel workers (default: 10)",
    )
    parser.add_argument(
        "-y", "--no-confirm", action="store_true",
        help="Auto-confirm all prompts (non-interactive)",
    )
    parser.add_argument(
        "--version", action="store_true",
        help="Show version and exit",
    )

    sub = parser.add_subparsers(dest="command", metavar="COMMAND")

    # refresh / ref
    p_ref = sub.add_parser("refresh", aliases=["ref"], help="Refresh repositories in parallel (always cleans cache first)")
    p_ref.add_argument(
        "-f", "--force", action="store_true",
        help="Force refresh even if repos appear up to date",
    )
    p_ref.add_argument(
        "--no-clean", action="store_true",
        help="Skip the automatic 'zypper clean --all' that normally runs before refresh",
    )

    # dist-upgrade / dup
    p_dup = sub.add_parser("dist-upgrade", aliases=["dup"], help="Distribution upgrade in parallel")
    p_dup.add_argument("-d", "--download-only", action="store_true", help="Download only, don't install")

    # install / in
    p_in = sub.add_parser("install", aliases=["in"], help="Install packages in parallel")
    p_in.add_argument("packages", nargs="+", help="Package name(s) to install")
    p_in.add_argument("-d", "--download-only", action="store_true", help="Download only, don't install")

    # install-new-recommends / inr
    p_inr = sub.add_parser("install-new-recommends", aliases=["inr"],
                            help="Install newly recommended packages")
    p_inr.add_argument("-d", "--download-only", action="store_true", help="Download only, don't install")

    args = parser.parse_args()

    if args.version:
        console.print(f"[bold cyan]ZypperX v{VERSION}[/bold cyan]")
        return

    if not args.command:
        parser.print_help()
        return

    # Pull per-subcommand flags (absent on subcommands that don't define them)
    force = getattr(args, "force", False)
    download_only = getattr(args, "download_only", False)
    no_clean = getattr(args, "no_clean", False)
    no_confirm = args.no_confirm

    engine = ZypperXEngine(jobs=args.jobs, force=force)
    engine.check_prerequisites()

    # Signal handlers — ensure lock and workspaces are cleaned on Ctrl-C / kill
    def _on_signal(sig, _frame):
        console.print(f"\n[yellow]Signal {sig} received — cleaning up...[/yellow]")
        emergency_cleanup(engine)
        sys.exit(130)

    signal.signal(signal.SIGINT, _on_signal)
    signal.signal(signal.SIGTERM, _on_signal)

    # ── Always clean repo metadata cache before refresh ───────────────────────
    # Stale/corrupt metadata in /var/cache/zypp is the #1 cause of SSL and
    # "invalid repository" errors. Cleaning first ensures workers always
    # fetch a fresh copy. Pass --no-clean to skip this if you're in a hurry.
    if args.command in ("refresh", "ref"):
        if not no_clean:
            console.print("[cyan]Cleaning repo metadata cache (zypper clean --all)...[/cyan]")
            result = subprocess.run(
                ["zypper", "--non-interactive", "clean", "--all"],
                check=False,
            )
            if result.returncode == 0:
                console.print("[green]✓ Cache cleared.[/green]")
            else:
                console.print(
                    "[yellow]Warning: zypper clean returned a non-zero exit code "
                    "(continuing anyway).[/yellow]"
                )
        else:
            console.print("[yellow]Skipping cache clean (--no-clean passed).[/yellow]")

    # ── Build task list ───────────────────────────────────────────────────────
    task_items: List[str] = []
    mode = ""

    if args.command in ("refresh", "ref"):
        task_items = engine.get_enabled_repos()
        mode = "refresh"
    else:
        if args.command in ("dist-upgrade", "dup"):
            mode = "dist-upgrade"
        elif args.command in ("install", "in"):
            mode = "install"
        else:
            mode = "install-new-recommends"

        target_pkgs: List[str] = getattr(args, "packages", [])
        task_items, size = engine.get_transaction_packages(mode, target_pkgs)

        if not task_items:
            console.print("[green]Nothing to do.[/green]")
            return

        if size > 0:
            size_mb = size / (1024 * 1024)
            console.print(
                Panel(
                    f"Packages to download: [bold]{len(task_items)}[/bold]\n"
                    f"Download size:        [bold]{size_mb:.2f} MB[/bold]",
                    title="Transaction Summary",
                    expand=False,
                )
            )
            if not no_confirm and not Confirm.ask("Proceed with parallel download?"):
                sys.exit(0)
        else:
            console.print("[green]All packages already in cache.[/green]")
            task_items = []

    # ── Execute parallel tasks ────────────────────────────────────────────────
    if task_items:
        try:
            engine.get_lock()
            label = "Refreshing repos..." if mode == "refresh" else "Downloading packages..."
            task_type = "refresh" if mode == "refresh" else "download"
            await engine.execute_tasks(task_items, task_type, label)
        finally:
            engine.release_lock()

    # ── Hand off to native zypper for installation ────────────────────────────
    if mode != "refresh" and not download_only:
        console.print("\n[bold yellow]>> Handing off to native zypper for installation...[/bold yellow]")
        final_cmd = ["zypper"]
        if no_confirm:
            final_cmd.append("--non-interactive")
        final_cmd.append("--no-cd")
        if mode == "dist-upgrade":
            final_cmd.append("dist-upgrade")
        elif mode == "install-new-recommends":
            final_cmd.append("install-new-recommends")
        elif mode == "install":
            final_cmd += ["install"] + getattr(args, "packages", [])

        env = os.environ.copy()
        env.update(ZYPPER_ENV)
        os.execvpe("zypper", final_cmd, env)

    elif mode != "refresh":
        console.print("[bold green]✓ Download complete. Run without -d to install.[/bold green]")


# ─── Entry Point ──────────────────────────────────────────────────────────────

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        console.print("\n[yellow]Interrupted.[/yellow]")
        sys.exit(130)
    except Exception as exc:
        console.print(f"[bold red]Fatal error:[/bold red] {exc}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
