agent/deployment_engine.py

"""
Deployment orchestration engine.

Orchestrates end-to-end configuration deployment:
- Config generation from templates
- Pre-deployment validation
- Device connectivity and deployment
- Post-deployment verification
- Automatic rollback on failure
- Parallel deployment via Ray integration
"""

import logging
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass, field
from enum import Enum
from datetime import datetime
import time

from agent.device_driver import (
    DeviceDriver, DeviceCredentials, DeviceType,
    ConfigDeploymentResult, ConnectionStatus
)
from agent.config_templates import ConfigTemplateEngine

logger = logging.getLogger(__name__)


class DeploymentStatus(Enum):
    """Deployment status"""
    PENDING = "pending"
    VALIDATING = "validating"
    DEPLOYING = "deploying"
    VERIFYING = "verifying"
    SUCCESS = "success"
    FAILED = "failed"
    ROLLED_BACK = "rolled_back"


@dataclass
class DeploymentTask:
    """Single device deployment task"""
    device_id: str
    device_type: DeviceType
    hostname: str
    username: str
    password: str
    config: str
    dry_run: bool = False
    pre_checks: List[str] = field(default_factory=list)
    post_checks: List[str] = field(default_factory=list)


@dataclass
class DeploymentResult:
    """Result from deployment orchestration"""
    device_id: str
    status: DeploymentStatus
    config_deployed: Optional[str] = None
    config_before: Optional[str] = None
    config_after: Optional[str] = None
    pre_check_results: Dict[str, bool] = field(default_factory=dict)
    post_check_results: Dict[str, bool] = field(default_factory=dict)
    deployment_output: Optional[str] = None
    error: Optional[str] = None
    duration_seconds: float = 0.0
    rolled_back: bool = False
    timestamp: datetime = field(default_factory=datetime.now)


class DeploymentEngine:
    """
    Orchestrates configuration deployment to network devices.
    
    Features:
    - Multi-vendor device support (Cisco, Arista, Juniper)
    - Pre/post deployment validation
    - Automatic rollback on failure
    - Parallel deployment via Ray
    - Dry-run mode for testing
    """
    
    def __init__(self, use_napalm: bool = True, use_ray: bool = False):
        """
        Initialize deployment engine.
        
        Args:
            use_napalm: Prefer NAPALM over Netmiko
            use_ray: Use Ray for parallel deployment
        """
        self.driver = DeviceDriver(use_napalm=use_napalm)
        self.template_engine = ConfigTemplateEngine()
        self.use_ray = use_ray
        
        if use_ray:
            try:
                from agent.ray_executor import RayExecutor
                self.ray_executor = RayExecutor()
                logger.info("DeploymentEngine using Ray for parallel execution")
            except ImportError:
                logger.warning("Ray not available - falling back to serial execution")
                self.use_ray = False
                self.ray_executor = None
        else:
            self.ray_executor = None
        
        self.deployment_history: List[DeploymentResult] = []
    
    def deploy_single_device(self, task: DeploymentTask) -> DeploymentResult:
        """
        Deploy configuration to a single device.
        
        Args:
            task: Deployment task specification
            
        Returns:
            DeploymentResult with outcome
        """
        logger.info(f"Starting deployment to {task.device_id} (dry_run={task.dry_run})")
        
        start_time = time.time()
        result = DeploymentResult(
            device_id=task.device_id,
            status=DeploymentStatus.PENDING
        )
        
        try:
            # 1. Connect to device
            logger.info(f"Connecting to {task.device_id}...")
            credentials = DeviceCredentials(
                hostname=task.hostname,
                username=task.username,
                password=task.password,
                device_type=task.device_type
            )
            
            conn = self.driver.connect(credentials)
            
            if conn.status != ConnectionStatus.CONNECTED:
                raise Exception(f"Failed to connect: {conn.last_error}")
            
            # 2. Pre-deployment checks
            result.status = DeploymentStatus.VALIDATING
            logger.info(f"Running {len(task.pre_checks)} pre-deployment checks...")
            
            for check in task.pre_checks:
                check_result = self._run_validation_check(task.device_id, check)
                result.pre_check_results[check] = check_result
                
                if not check_result:
                    raise Exception(f"Pre-deployment check failed: {check}")
            
            # 3. Get current config (for rollback)
            result.config_before = self.driver.get_config(task.device_id, "running")
            
            # 4. Deploy configuration
            result.status = DeploymentStatus.DEPLOYING
            logger.info(f"Deploying configuration to {task.device_id}...")
            
            deploy_result = self.driver.deploy_config(
                device_id=task.device_id,
                config=task.config,
                dry_run=task.dry_run,
                replace=False  # Merge by default
            )
            
            if not deploy_result.success:
                raise Exception(f"Deployment failed: {deploy_result.error}")
            
            result.config_deployed = task.config
            result.config_after = deploy_result.config_after
            result.deployment_output = deploy_result.output
            
            # 5. Post-deployment verification
            if not task.dry_run:
                result.status = DeploymentStatus.VERIFYING
                logger.info(f"Running {len(task.post_checks)} post-deployment checks...")
                
                # Give device time to apply config
                time.sleep(2)
                
                for check in task.post_checks:
                    check_result = self._run_validation_check(task.device_id, check)
                    result.post_check_results[check] = check_result
                    
                    if not check_result:
                        # Post-check failed - rollback!
                        logger.error(f"Post-deployment check failed: {check}")
                        logger.warning(f"Initiating rollback on {task.device_id}")
                        
                        rollback_result = self.driver.rollback_config(
                            device_id=task.device_id,
                            config=result.config_before
                        )
                        
                        result.rolled_back = rollback_result.success
                        result.status = DeploymentStatus.ROLLED_BACK if result.rolled_back else DeploymentStatus.FAILED
                        raise Exception(f"Post-deployment check failed: {check} (rollback {'succeeded' if result.rolled_back else 'FAILED'})")
            
            # Success!
            result.status = DeploymentStatus.SUCCESS
            result.duration_seconds = time.time() - start_time
            logger.info(f"✓ Deployment to {task.device_id} completed successfully in {result.duration_seconds:.1f}s")
            
        except Exception as e:
            result.status = DeploymentStatus.FAILED
            result.error = str(e)
            result.duration_seconds = time.time() - start_time
            logger.error(f"✗ Deployment to {task.device_id} failed: {e}")
        
        finally:
            # Disconnect
            self.driver.disconnect(task.device_id)
        
        # Store in history
        self.deployment_history.append(result)
        
        return result
    
    def deploy_multiple_devices(self, tasks: List[DeploymentTask],
                                parallel: bool = False,
                                max_failures: Optional[int] = None) -> List[DeploymentResult]:
        """
        Deploy configuration to multiple devices.
        
        Args:
            tasks: List of deployment tasks
            parallel: If True, use parallel execution (Ray if available)
            max_failures: Stop if this many devices fail (None = deploy all)
            
        Returns:
            List of DeploymentResults
        """
        logger.info(f"Deploying to {len(tasks)} devices (parallel={parallel})")
        
        if parallel and self.use_ray and self.ray_executor:
            return self._deploy_parallel_ray(tasks, max_failures)
        elif parallel:
            logger.warning("Parallel mode requested but Ray not available - using serial")
        
        # Serial deployment
        results = []
        failures = 0
        
        for task in tasks:
            result = self.deploy_single_device(task)
            results.append(result)
            
            if result.status == DeploymentStatus.FAILED:
                failures += 1
                
                if max_failures and failures >= max_failures:
                    logger.error(f"Max failures ({max_failures}) reached - stopping deployment")
                    
                    # Mark remaining as pending
                    remaining = len(tasks) - len(results)
                    logger.warning(f"Skipping {remaining} remaining devices")
                    break
        
        return results
    
    def _deploy_parallel_ray(self, tasks: List[DeploymentTask],
                            max_failures: Optional[int]) -> List[DeploymentResult]:
        """Deploy using Ray parallel execution"""
        logger.info(f"Using Ray to deploy to {len(tasks)} devices in parallel")
        
        # This would integrate with Ray executor
        # For now, fall back to serial
        logger.warning("Ray parallel deployment not yet implemented - using serial")
        return self.deploy_multiple_devices(tasks, parallel=False, max_failures=max_failures)
    
    def _run_validation_check(self, device_id: str, check: str) -> bool:
        """
        Run validation check on device.
        
        Args:
            device_id: Device identifier
            check: Check command or test
            
        Returns:
            True if check passes
        """
        try:
            # Parse check type
            if check.startswith("ping:"):
                # Ping test
                target = check.split(":", 1)[1]
                result = self.driver.send_command(device_id, f"ping {target}")
                return result.success and "success" in result.output.lower()
            
            elif check.startswith("interface:"):
                # Interface status check
                interface = check.split(":", 1)[1]
                result = self.driver.send_command(device_id, f"show interface {interface}")
                return result.success and "up" in result.output.lower()
            
            elif check.startswith("command:"):
                # Generic command check (success = command runs without error)
                command = check.split(":", 1)[1]
                result = self.driver.send_command(device_id, command)
                return result.success
            
            else:
                # Default: run as show command
                result = self.driver.send_command(device_id, check)
                return result.success
                
        except Exception as e:
            logger.error(f"Validation check '{check}' failed on {device_id}: {e}")
            return False
    
    def generate_and_deploy(self, device: Any, network_context: Dict[str, Any],
                           credentials: Dict[str, str],
                           dry_run: bool = False,
                           pre_checks: Optional[List[str]] = None,
                           post_checks: Optional[List[str]] = None) -> DeploymentResult:
        """
        Generate configuration from template and deploy to device.
        
        Args:
            device: Device object from NetworkModel
            network_context: Network-level context (VLANs, routing, etc.)
            credentials: Device credentials (username, password)
            dry_run: If True, generate and validate but don't deploy
            pre_checks: Optional pre-deployment validation checks
            post_checks: Optional post-deployment validation checks
            
        Returns:
            DeploymentResult
        """
        device_name = device.get('name') if isinstance(device, dict) else device.name
        logger.info(f"Generating and deploying config for {device_name}")
        
        # 1. Generate configuration from template
        config = self.template_engine.generate_device_config(device, network_context)
        
        # 2. Map device vendor to driver type (handle both dict and object)
        vendor = device.get('vendor') if isinstance(device, dict) else device.vendor
        model = device.get('model', '') if isinstance(device, dict) else getattr(device, 'model', '')
        device_type = self._map_vendor_to_device_type(vendor, model)
        
        # 3. Create deployment task
        device_name = device.get('name') if isinstance(device, dict) else device.name
        mgmt_ip = device.get('mgmt_ip') if isinstance(device, dict) else device.mgmt_ip
        
        task = DeploymentTask(
            device_id=device_name,
            device_type=device_type,
            hostname=mgmt_ip,
            username=credentials.get('username', 'admin'),
            password=credentials.get('password', 'admin'),
            config=config,
            dry_run=dry_run,
            pre_checks=pre_checks or [],
            post_checks=post_checks or []
        )
        
        # 4. Deploy
        return self.deploy_single_device(task)
    
    def _map_vendor_to_device_type(self, vendor: str, model: str) -> DeviceType:
        """Map vendor/model to DeviceType"""
        vendor_lower = vendor.lower()
        model_lower = model.lower()
        
        if 'cisco' in vendor_lower:
            if 'nexus' in model_lower or 'nxos' in model_lower:
                return DeviceType.CISCO_NXOS
            elif 'xe' in model_lower or 'asr' in model_lower or 'isr' in model_lower:
                return DeviceType.CISCO_XE
            else:
                return DeviceType.CISCO_IOS
        
        elif 'arista' in vendor_lower:
            return DeviceType.ARISTA_EOS
        
        elif 'juniper' in vendor_lower:
            return DeviceType.JUNIPER_JUNOS
        
        else:
            logger.warning(f"Unknown vendor {vendor}, defaulting to Cisco IOS")
            return DeviceType.CISCO_IOS
    
    def get_deployment_summary(self) -> Dict[str, Any]:
        """Get summary of deployment history"""
        total = len(self.deployment_history)
        
        if total == 0:
            return {
                'total_deployments': 0,
                'success_count': 0,
                'failed_count': 0,
                'rolled_back_count': 0,
                'success_rate': 0.0
            }
        
        success = sum(1 for r in self.deployment_history if r.status == DeploymentStatus.SUCCESS)
        failed = sum(1 for r in self.deployment_history if r.status == DeploymentStatus.FAILED)
        rolled_back = sum(1 for r in self.deployment_history if r.rolled_back)
        
        return {
            'total_deployments': total,
            'success_count': success,
            'failed_count': failed,
            'rolled_back_count': rolled_back,
            'success_rate': (success / total) * 100 if total > 0 else 0.0,
            'avg_duration': sum(r.duration_seconds for r in self.deployment_history) / total,
            'latest_deployments': [
                {
                    'device_id': r.device_id,
                    'status': r.status.value,
                    'timestamp': r.timestamp.isoformat(),
                    'duration': r.duration_seconds
                }
                for r in self.deployment_history[-10:]  # Last 10
            ]
        }
    
    def verify_deployment(self, device_id: str, expected_config_snippet: Optional[str] = None) -> Dict[str, Any]:
        """
        Verify deployment was successful.
        
        Args:
            device_id: Device identifier
            expected_config_snippet: Optional config snippet to verify is present
            
        Returns:
            Verification results
        """
        results = {
            'device_id': device_id,
            'connectivity': False,
            'config_retrieved': False,
            'snippet_found': None,
            'facts': None,
            'errors': []
        }
        
        try:
            # Check if we have a connection
            if device_id not in self.driver.connections:
                results['errors'].append("Device not connected")
                return results
            
            # Test connectivity
            results['connectivity'] = self.driver.verify_connectivity(device_id)
            
            if not results['connectivity']:
                results['errors'].append("Device not responsive")
                return results
            
            # Get current config
            config = self.driver.get_config(device_id, "running")
            results['config_retrieved'] = config is not None
            
            if not results['config_retrieved']:
                results['errors'].append("Could not retrieve configuration")
                return results
            
            # Check for expected snippet
            if expected_config_snippet:
                results['snippet_found'] = expected_config_snippet in config
                
                if not results['snippet_found']:
                    results['errors'].append(f"Expected config snippet not found: {expected_config_snippet[:50]}...")
            
            # Get device facts
            facts = self.driver.get_device_facts(device_id)
            results['facts'] = facts
            
        except Exception as e:
            results['errors'].append(str(e))
            logger.error(f"Verification failed for {device_id}: {e}")
        
        return results
    
    def cleanup(self):
        """Disconnect all devices and cleanup resources"""
        self.driver.disconnect_all()
        
        if self.ray_executor:
            self.ray_executor.shutdown()
        
        logger.info("DeploymentEngine cleanup complete")