utils/dependency_manager.py

"""
utils/dependency_manager.py
----------------------------
Dependency inference and installation for AutoDevAgent.

Before executing generated Python code, this module scans the code
for import statements, checks whether each package is already
available, and attempts to install missing ones via pip install --user.

Design:
  - Uses Python's ast module to parse imports — more reliable than
    regex since it handles multi-line and aliased imports correctly.
  - Only attempts installation for packages that are not already
    importable — avoids unnecessary pip calls.
  - Maps common import names to their pip package names where they
    differ (e.g. "cv2" → "opencv-python", "PIL" → "Pillow").
  - Fails gracefully — if installation fails, returns a clear message
    rather than crashing. The executor will then surface the error.
  - C extensions and system-level packages will fail — documented in
    README as a known limitation.

Known limitation:
  - pip install --user works for pure Python libraries only.
    Packages with C extensions (numpy, pandas, etc.) may fail on
    HuggingFace Spaces free tier. Surface a clear message if so.

Usage:
    from utils.dependency_manager import DependencyManager

    manager = DependencyManager()
    result = manager.resolve(code)
    print(result.missing)         # packages that were not installed
    print(result.installed)       # packages successfully installed
    print(result.failed)          # packages that failed to install
"""

import ast
import importlib
import logging
import subprocess
import sys
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)


# ------------------------------------------------------------------ #
#  Import name → pip package name mapping                              #
#                                                                      #
#  Many packages have different import names vs pip install names.     #
#  This map covers the most common cases seen in generated code.       #
# ------------------------------------------------------------------ #

IMPORT_TO_PIP: dict[str, str] = {
    "cv2":          "opencv-python",
    "PIL":          "Pillow",
    "sklearn":      "scikit-learn",
    "bs4":          "beautifulsoup4",
    "yaml":         "pyyaml",
    "dotenv":       "python-dotenv",
    "dateutil":     "python-dateutil",
    "attr":         "attrs",
    "gi":           "PyGObject",
    "wx":           "wxPython",
    "Crypto":       "pycryptodome",
    "serial":       "pyserial",
    "usb":          "pyusb",
    "magic":        "python-magic",
    "Image":        "Pillow",
}

# Standard library modules — never attempt to pip install these
STDLIB_MODULES: frozenset[str] = frozenset({
    "abc", "ast", "asyncio", "base64", "collections", "contextlib",
    "copy", "csv", "datetime", "decimal", "difflib", "email", "enum",
    "fnmatch", "functools", "glob", "hashlib", "heapq", "html",
    "http", "inspect", "io", "itertools", "json", "logging", "math",
    "multiprocessing", "operator", "os", "pathlib", "pickle",
    "platform", "pprint", "queue", "random", "re", "shutil", "signal",
    "socket", "sqlite3", "ssl", "stat", "string", "struct",
    "subprocess", "sys", "tempfile", "textwrap", "threading", "time",
    "timeit", "traceback", "typing", "unicodedata", "unittest",
    "urllib", "uuid", "warnings", "weakref", "xml", "zipfile",
    "zlib", "builtins", "gc", "getpass", "getopt", "gzip",
    "importlib", "inspect", "keyword", "linecache", "locale",
    "numbers", "os.path", "posixpath", "pdb", "profile", "pstats",
    "pty", "pwd", "resource", "select", "shelve", "sysconfig",
    "syslog", "termios", "tty", "types", "dataclasses", "fractions",
    "statistics", "secrets", "token", "tokenize", "trace",
})


# ------------------------------------------------------------------ #
#  Result dataclass                                                    #
# ------------------------------------------------------------------ #

@dataclass
class DependencyResult:
    """
    Result of a dependency resolution attempt.

    Attributes:
        imports:   All top-level package names found in the code.
        available: Packages already importable — no action needed.
        installed: Packages successfully installed by this call.
        failed:    Packages that could not be installed.
        messages:  Human-readable log of what happened per package.
    """
    imports:   list[str] = field(default_factory=list)
    available: list[str] = field(default_factory=list)
    installed: list[str] = field(default_factory=list)
    failed:    list[str] = field(default_factory=list)
    messages:  list[str] = field(default_factory=list)

    @property
    def all_resolved(self) -> bool:
        """True if every import is either available or successfully installed."""
        return len(self.failed) == 0

    def summary(self) -> str:
        """Return a one-line summary suitable for the UI status bar."""
        if not self.imports:
            return "No external imports detected."
        parts = []
        if self.available:
            parts.append(f"{len(self.available)} already available")
        if self.installed:
            parts.append(f"{len(self.installed)} installed")
        if self.failed:
            parts.append(f"{len(self.failed)} failed: {', '.join(self.failed)}")
        return " · ".join(parts) if parts else "All dependencies resolved."


# ------------------------------------------------------------------ #
#  Manager                                                             #
# ------------------------------------------------------------------ #

class DependencyManager:
    """
    Scans generated code for imports and installs missing packages.

    Workflow:
      1. Parse the code with ast to extract all import names.
      2. Filter out stdlib modules.
      3. For each remaining import: check if importable, skip if so.
      4. For missing imports: attempt pip install --user.
      5. Return a DependencyResult with full accounting.
    """

    def resolve(self, code: str) -> DependencyResult:
        """
        Resolve all dependencies for the given code string.

        Args:
            code: Raw Python source code to analyse.

        Returns:
            DependencyResult with installed/failed/available breakdown.
        """
        result = DependencyResult()

        # ── Step 1: Extract all top-level import names ────────────── #
        imports = _extract_imports(code)
        result.imports = imports

        if not imports:
            logger.debug("DependencyManager: no external imports found")
            return result

        logger.info("DependencyManager: found imports: %s", imports)

        # ── Step 2: Check and install each package ────────────────── #
        for import_name in imports:
            pip_name = IMPORT_TO_PIP.get(import_name, import_name)

            if _is_importable(import_name):
                result.available.append(import_name)
                result.messages.append(f"✓ {import_name} — already available")
                logger.debug("DependencyManager: %s already available", import_name)
                continue

            # Attempt installation
            logger.info("DependencyManager: installing %s (pip: %s)", import_name, pip_name)
            success, msg = _pip_install(pip_name)

            if success:
                result.installed.append(import_name)
                result.messages.append(f"✓ {import_name} — installed successfully")
                logger.info("DependencyManager: installed %s", pip_name)
            else:
                result.failed.append(import_name)
                result.messages.append(
                    f"✗ {import_name} — install failed: {msg[:120]}"
                )
                logger.warning(
                    "DependencyManager: failed to install %s: %s", pip_name, msg[:120]
                )

        return result


# ------------------------------------------------------------------ #
#  Module-level helpers                                                #
# ------------------------------------------------------------------ #

def _extract_imports(code: str) -> list[str]:
    """
    Parse Python source code and extract top-level package names.

    Uses ast.parse so it handles all valid import forms:
      - import os
      - import os.path
      - from os import path
      - from os.path import join
      - import numpy as np

    Filters out stdlib modules automatically.

    Args:
        code: Raw Python source code.

    Returns:
        List of unique top-level package names that are not stdlib.
    """
    imports: set[str] = set()

    try:
        tree = ast.parse(code)
    except SyntaxError:
        # Code has a syntax error — executor will catch this later
        logger.debug("DependencyManager: syntax error while parsing imports")
        return []

    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                # Take only the top-level name: "os.path" → "os"
                top = alias.name.split(".")[0]
                if top not in STDLIB_MODULES:
                    imports.add(top)

        elif isinstance(node, ast.ImportFrom):
            if node.module:
                top = node.module.split(".")[0]
                if top not in STDLIB_MODULES:
                    imports.add(top)

    return sorted(imports)


def _is_importable(package_name: str) -> bool:
    """
    Check whether a package can be imported in the current environment.

    Args:
        package_name: Top-level package name (e.g. "numpy", "requests").

    Returns:
        True if importable, False otherwise.
    """
    try:
        importlib.import_module(package_name)
        return True
    except ImportError:
        return False


def _pip_install(pip_name: str) -> tuple[bool, str]:
    """
    Attempt to install a package via pip install --user.

    Uses --user to avoid permission errors on HuggingFace Spaces and
    other environments where the site-packages directory is read-only.
    Uses --quiet to suppress verbose pip output in logs.

    Args:
        pip_name: The pip package name to install.

    Returns:
        Tuple of (success: bool, message: str).
    """
    try:
        proc = subprocess.run(
            [
                sys.executable, "-m", "pip", "install",
                "--user", "--quiet", "--no-warn-script-location",
                pip_name,
            ],
            capture_output=True,
            text=True,
            timeout=60,   # pip installs can be slow on cold starts
        )

        if proc.returncode == 0:
            return True, "installed successfully"
        else:
            # Surface the most useful part of pip's error output
            error_output = (proc.stderr or proc.stdout).strip()
            lines = error_output.splitlines()
            # Keep last 5 lines — pip errors tend to be most specific there
            excerpt = "\n".join(lines[-5:]) if len(lines) > 5 else error_output
            return False, excerpt

    except subprocess.TimeoutExpired:
        return False, "pip install timed out after 60 seconds"
    except Exception as e:
        return False, str(e)