app/package_manager.py

"""Package installation and checking for the configured environment."""

from __future__ import annotations

import asyncio
import logging
from pathlib import Path

from app.config import Settings, get_settings
from app.models import PackageCheckResult

logger = logging.getLogger(__name__)


class PackageManager:
    """Manages package installation and checking."""

    def __init__(self, settings: Settings | None = None):
        self._settings = settings or get_settings()
        self._install_lock = asyncio.Lock()

    async def install_packages(self, packages: list[str]) -> dict:
        """Install packages into the configured environment."""
        async with self._install_lock:
            cmd = self._settings.get_pip_command() + packages
            logger.info("Installing packages: %s", " ".join(packages))

            try:
                process = await asyncio.create_subprocess_exec(
                    *cmd,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                )

                stdout_bytes, stderr_bytes = await asyncio.wait_for(
                    process.communicate(),
                    timeout=300,  # 5 minute timeout for installs
                )

                stdout = stdout_bytes.decode("utf-8", errors="replace")
                stderr = stderr_bytes.decode("utf-8", errors="replace")
                success = process.returncode == 0

                if success:
                    logger.info("Successfully installed: %s", ", ".join(packages))
                else:
                    logger.error("Failed to install packages: %s", stderr)

                return {
                    "success": success,
                    "stdout": stdout,
                    "stderr": stderr,
                    "return_code": process.returncode,
                    "packages": packages,
                }
            except asyncio.TimeoutError:
                return {
                    "success": False,
                    "stderr": "Package installation timed out after 300s",
                    "stdout": "",
                    "return_code": -1,
                    "packages": packages,
                }
            except Exception as e:
                logger.exception("Error installing packages")
                return {
                    "success": False,
                    "stderr": str(e),
                    "stdout": "",
                    "return_code": -1,
                    "packages": packages,
                }

    async def check_packages(self, packages: list[str]) -> list[PackageCheckResult]:
        """Check if packages are installed."""
        python_exec = self._settings.get_python_executable()

        # Build a single script that checks all packages
        check_script = "import importlib, json\nresults = {}\n"
        for pkg in packages:
            # Normalize package name for import
            import_name = pkg.split("==")[0].split(">=")[0].split("<=")[0].split("[")[0]
            import_name = import_name.replace("-", "_").lower()
            check_script += f"""
try:
    mod = importlib.import_module("{import_name}")
    version = getattr(mod, "__version__", "unknown")
    results["{pkg}"] = {{"installed": True, "version": version}}
except ImportError:
    results["{pkg}"] = {{"installed": False, "version": None}}
"""
        check_script += "print(json.dumps(results))\n"

        try:
            if self._settings.env_type.value == "conda":
                cmd = f"{python_exec} -c {repr(check_script)}"
                process = await asyncio.create_subprocess_shell(
                    cmd,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                )
            else:
                process = await asyncio.create_subprocess_exec(
                    python_exec,
                    "-c",
                    check_script,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                )

            stdout_bytes, _ = await asyncio.wait_for(
                process.communicate(), timeout=30
            )

            import json
            results_data = json.loads(stdout_bytes.decode("utf-8").strip())

            return [
                PackageCheckResult(
                    package=pkg,
                    installed=info["installed"],
                    version=info.get("version"),
                )
                for pkg, info in results_data.items()
            ]
        except Exception as e:
            logger.exception("Error checking packages")
            return [
                PackageCheckResult(package=pkg, installed=False) for pkg in packages
            ]


_manager: PackageManager | None = None


def get_package_manager() -> PackageManager:
    global _manager
    if _manager is None:
        _manager = PackageManager()
    return _manager


def reset_package_manager(settings: Settings | None = None) -> PackageManager:
    global _manager
    _manager = PackageManager(settings)
    return _manager