deb-mock/deb_mock/environment_manager.py

"""
Environment Management API for deb-mock

This module provides comprehensive environment management capabilities
for external tools integrating with deb-mock.
"""

import os
import sys
import json
import tempfile
import subprocess
import shutil
from pathlib import Path
from typing import Dict, List, Any, Optional, Union, Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime

from .core import DebMock
from .config import Config
from .exceptions import ConfigurationError, ChrootError, SbuildError


@dataclass
class EnvironmentInfo:
    """Information about a mock environment"""
    name: str
    architecture: str
    suite: str
    status: str
    created: Optional[datetime] = None
    modified: Optional[datetime] = None
    size: int = 0
    packages_installed: List[str] = None
    mounts: List[Dict[str, str]] = None


@dataclass
class BuildResult:
    """Result of a build operation"""
    success: bool
    artifacts: List[str]
    output_dir: str
    log_file: str
    metadata: Dict[str, Any]
    error: Optional[str] = None
    duration: float = 0.0


class EnvironmentManager:
    """
    Comprehensive environment management for deb-mock
    
    This class provides a high-level interface for managing mock environments,
    executing commands, and collecting artifacts.
    """
    
    def __init__(self, config: Optional[Config] = None):
        """Initialize the environment manager"""
        if config is None:
            config = Config.default()
        
        self.config = config
        self.deb_mock = DebMock(config)
        self._active_environments = {}
    
    def create_environment(self, 
                          name: str,
                          arch: str = None,
                          suite: str = None,
                          packages: List[str] = None,
                          force: bool = False) -> EnvironmentInfo:
        """
        Create a new mock environment
        
        Args:
            name: Name for the environment
            arch: Target architecture
            suite: Debian suite
            packages: Initial packages to install
            force: Force creation even if environment exists
        
        Returns:
            EnvironmentInfo object
        """
        if not force and self.environment_exists(name):
            raise ValueError(f"Environment '{name}' already exists")
        
        # Remove existing environment if force is True
        if force and self.environment_exists(name):
            self.remove_environment(name)
        
        try:
            # Create the chroot environment
            self.deb_mock.init_chroot(name, arch, suite)
            
            # Install initial packages if specified
            if packages:
                self.deb_mock.install_packages(packages)
            
            # Get environment info
            info = self.get_environment_info(name, arch, suite)
            self._active_environments[name] = info
            
            return info
            
        except Exception as e:
            raise RuntimeError(f"Failed to create environment '{name}': {e}")
    
    def environment_exists(self, name: str) -> bool:
        """Check if an environment exists"""
        return self.deb_mock.chroot_manager.chroot_exists(name)
    
    def get_environment_info(self, name: str, arch: str = None, suite: str = None) -> EnvironmentInfo:
        """Get detailed information about an environment"""
        if not self.environment_exists(name):
            raise ValueError(f"Environment '{name}' does not exist")
        
        # Get basic chroot info
        chroot_info = self.deb_mock.chroot_manager.get_chroot_info(name)
        
        # Get installed packages
        packages = self._get_installed_packages(name)
        
        # Get mount information
        mounts = self.deb_mock.chroot_manager.list_mounts(name)
        
        return EnvironmentInfo(
            name=name,
            architecture=arch or self.config.architecture,
            suite=suite or self.config.suite,
            status=chroot_info.get('status', 'unknown'),
            created=chroot_info.get('created'),
            modified=chroot_info.get('modified'),
            size=chroot_info.get('size', 0),
            packages_installed=packages,
            mounts=mounts
        )
    
    def list_environments(self) -> List[EnvironmentInfo]:
        """List all available environments"""
        environments = []
        
        for name in self.deb_mock.list_chroots():
            try:
                info = self.get_environment_info(name)
                environments.append(info)
            except Exception as e:
                print(f"Warning: Failed to get info for environment '{name}': {e}")
        
        return environments
    
    def remove_environment(self, name: str, force: bool = False) -> None:
        """Remove an environment"""
        if not self.environment_exists(name):
            if not force:
                raise ValueError(f"Environment '{name}' does not exist")
            return
        
        # Clean up active environment tracking
        if name in self._active_environments:
            del self._active_environments[name]
        
        # Remove the chroot
        self.deb_mock.clean_chroot(name)
    
    def update_environment(self, name: str) -> None:
        """Update packages in an environment"""
        if not self.environment_exists(name):
            raise ValueError(f"Environment '{name}' does not exist")
        
        self.deb_mock.update_chroot(name)
    
    def execute_command(self, 
                       name: str,
                       command: Union[str, List[str]],
                       capture_output: bool = True,
                       check: bool = True,
                       timeout: Optional[int] = None) -> subprocess.CompletedProcess:
        """
        Execute a command in an environment
        
        Args:
            name: Environment name
            command: Command to execute
            capture_output: Whether to capture output
            check: Whether to raise exception on non-zero exit
            timeout: Command timeout in seconds
        
        Returns:
            CompletedProcess object
        """
        if not self.environment_exists(name):
            raise ValueError(f"Environment '{name}' does not exist")
        
        if isinstance(command, str):
            command = command.split()
        
        # Prepare command with timeout if specified
        if timeout:
            command = ['timeout', str(timeout)] + command
        
        try:
            result = self.deb_mock.chroot_manager.execute_in_chroot(
                name, command, capture_output=capture_output
            )
            
            if check and result.returncode != 0:
                raise subprocess.CalledProcessError(
                    result.returncode, command, result.stdout, result.stderr
                )
            
            return result
            
        except subprocess.CalledProcessError as e:
            if check:
                raise
            return e
    
    def install_packages(self, name: str, packages: List[str]) -> Dict[str, Any]:
        """Install packages in an environment"""
        if not self.environment_exists(name):
            raise ValueError(f"Environment '{name}' does not exist")
        
        return self.deb_mock.install_packages(packages)
    
    def copy_files(self, 
                  name: str,
                  source: str,
                  destination: str,
                  direction: str = "in") -> None:
        """
        Copy files to/from an environment
        
        Args:
            name: Environment name
            source: Source path
            destination: Destination path
            direction: "in" to copy into environment, "out" to copy out
        """
        if not self.environment_exists(name):
            raise ValueError(f"Environment '{name}' does not exist")
        
        if direction == "in":
            self.deb_mock.chroot_manager.copy_to_chroot(source, destination, name)
        elif direction == "out":
            self.deb_mock.chroot_manager.copy_from_chroot(source, destination, name)
        else:
            raise ValueError("Direction must be 'in' or 'out'")
    
    def collect_artifacts(self, 
                         name: str,
                         source_patterns: List[str] = None,
                         output_dir: str = None) -> List[str]:
        """
        Collect build artifacts from an environment
        
        Args:
            name: Environment name
            source_patterns: File patterns to search for
            output_dir: Output directory for artifacts
        
        Returns:
            List of collected artifact paths
        """
        if not self.environment_exists(name):
            raise ValueError(f"Environment '{name}' does not exist")
        
        if source_patterns is None:
            source_patterns = [
                '*.deb',
                '*.changes',
                '*.buildinfo',
                '*.dsc',
                '*.tar.*',
                '*.orig.tar.*',
                '*.debian.tar.*'
            ]
        
        if output_dir is None:
            output_dir = tempfile.mkdtemp(prefix='deb-mock-artifacts-')
        
        os.makedirs(output_dir, exist_ok=True)
        artifacts = []
        
        for pattern in source_patterns:
            # Find files matching pattern
            result = self.execute_command(
                name, ['find', '/build', '-name', pattern, '-type', 'f'],
                capture_output=True, check=False
            )
            
            if result.returncode == 0:
                for line in result.stdout.strip().split('\n'):
                    if line.strip():
                        source_path = line.strip()
                        filename = os.path.basename(source_path)
                        dest_path = os.path.join(output_dir, filename)
                        
                        # Copy artifact
                        self.copy_files(name, source_path, dest_path, "out")
                        artifacts.append(dest_path)
        
        return artifacts
    
    def build_package(self, 
                     name: str,
                     source_package: str,
                     output_dir: str = None,
                     **kwargs) -> BuildResult:
        """
        Build a package in an environment
        
        Args:
            name: Environment name
            source_package: Path to source package
            output_dir: Output directory for build artifacts
            **kwargs: Additional build options
        
        Returns:
            BuildResult object
        """
        if not self.environment_exists(name):
            raise ValueError(f"Environment '{name}' does not exist")
        
        start_time = datetime.now()
        
        try:
            # Set chroot name for build
            kwargs['chroot_name'] = name
            if output_dir:
                kwargs['output_dir'] = output_dir
            
            # Build the package
            result = self.deb_mock.build(source_package, **kwargs)
            
            # Calculate duration
            duration = (datetime.now() - start_time).total_seconds()
            
            return BuildResult(
                success=result.get('success', False),
                artifacts=result.get('artifacts', []),
                output_dir=result.get('output_dir', ''),
                log_file=result.get('log_file', ''),
                metadata=result.get('metadata', {}),
                duration=duration
            )
            
        except Exception as e:
            duration = (datetime.now() - start_time).total_seconds()
            return BuildResult(
                success=False,
                artifacts=[],
                output_dir=output_dir or '',
                log_file='',
                metadata={},
                error=str(e),
                duration=duration
            )
    
    @contextmanager
    def environment(self, 
                   name: str,
                   arch: str = None,
                   suite: str = None,
                   packages: List[str] = None,
                   create_if_missing: bool = True) -> Iterator[EnvironmentInfo]:
        """
        Context manager for environment operations
        
        Args:
            name: Environment name
            arch: Target architecture
            suite: Debian suite
            packages: Initial packages to install
            create_if_missing: Create environment if it doesn't exist
        
        Yields:
            EnvironmentInfo object
        """
        env_info = None
        created = False
        
        try:
            # Get or create environment
            if self.environment_exists(name):
                env_info = self.get_environment_info(name)
            elif create_if_missing:
                env_info = self.create_environment(name, arch, suite, packages)
                created = True
            else:
                raise ValueError(f"Environment '{name}' does not exist")
            
            yield env_info
            
        finally:
            # Clean up if we created the environment
            if created and env_info:
                try:
                    self.remove_environment(name)
                except Exception as e:
                    print(f"Warning: Failed to cleanup environment '{name}': {e}")
    
    def _get_installed_packages(self, name: str) -> List[str]:
        """Get list of installed packages in environment"""
        try:
            result = self.execute_command(
                name, ['dpkg', '-l'], capture_output=True, check=False
            )
            
            if result.returncode == 0:
                packages = []
                for line in result.stdout.split('\n'):
                    if line.startswith('ii'):
                        parts = line.split()
                        if len(parts) >= 3:
                            packages.append(parts[1])
                return packages
            
        except Exception:
            pass
        
        return []
    
    def export_environment(self, name: str, output_path: str) -> None:
        """Export environment to a tar archive"""
        if not self.environment_exists(name):
            raise ValueError(f"Environment '{name}' does not exist")
        
        chroot_path = self.deb_mock.config.get_chroot_path()
        
        # Create tar archive
        subprocess.run([
            'tar', '-czf', output_path, '-C', chroot_path, '.'
        ], check=True)
    
    def import_environment(self, name: str, archive_path: str) -> None:
        """Import environment from a tar archive"""
        if self.environment_exists(name):
            raise ValueError(f"Environment '{name}' already exists")
        
        # Create environment directory
        chroot_path = os.path.join(self.config.chroot_dir, name)
        os.makedirs(chroot_path, exist_ok=True)
        
        # Extract archive
        subprocess.run([
            'tar', '-xzf', archive_path, '-C', chroot_path
        ], check=True)
        
        # Create schroot configuration
        self.deb_mock.chroot_manager._create_schroot_config(
            name, chroot_path, self.config.architecture, self.config.suite
        )


# Convenience functions
def create_environment_manager(config: Optional[Config] = None) -> EnvironmentManager:
    """Create a new environment manager"""
    return EnvironmentManager(config)


def quick_environment(name: str = "quick-build",
                     arch: str = "amd64",
                     suite: str = "trixie",
                     packages: List[str] = None) -> EnvironmentManager:
    """Create a quick environment manager with default settings"""
    config = Config(
        chroot_name=name,
        architecture=arch,
        suite=suite,
        chroot_additional_packages=packages or []
    )
    return EnvironmentManager(config)