#!/usr/bin/env python3
"""
Phase 1 Production Deployment Script

Deploys production-ready experiments (Exp 7, 9, 11, 13) using gradual rollout strategy.

Usage:
    python infra/deployment/phase1_deploy.py --stage staging_validation
    python infra/deployment/phase1_deploy.py --stage canary_deployment
    python infra/deployment/phase1_deploy.py --stage full_deployment --auto-approve
"""

import argparse
import json
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional


class Phase1Deployment:
    """Manages Phase 1 production deployment of experiments."""
    
    def __init__(self, config_path: Path):
        """Initialize deployment with configuration."""
        self.config_path = config_path
        self.config = self._load_config()
        self.experiments = self.config["experiments"]
        self.deployment_strategy = self.config["deployment_strategy"]
        
    def _load_config(self) -> Dict:
        """Load deployment configuration."""
        with open(self.config_path, 'r') as f:
            return json.load(f)
    
    def validate_environment(self) -> bool:
        """Validate deployment environment."""
        print("=" * 80)
        print("PHASE 1 DEPLOYMENT - Environment Validation")
        print("=" * 80)
        print()
        
        checks = [
            ("Configuration file exists", self.config_path.exists()),
            ("All experiments configured", len(self.experiments) == 4),
            ("Deployment strategy defined", "deployment_strategy" in self.config),
            ("Monitoring configured", "monitoring" in self.config),
        ]
        
        all_passed = True
        for check_name, check_result in checks:
            status = "[PASS]" if check_result else "[FAIL]"
            print(f"{status}: {check_name}")
            if not check_result:
                all_passed = False
        
        print()
        return all_passed
    
    def deploy_stage(self, stage_name: str, auto_approve: bool = False) -> bool:
        """Deploy a specific stage."""
        stage = None
        for s in self.deployment_strategy["stages"]:
            if s["name"] == stage_name:
                stage = s
                break
        
        if not stage:
            print(f"[ERROR] Stage '{stage_name}' not found in deployment strategy")
            return False
        
        print("=" * 80)
        print(f"PHASE 1 DEPLOYMENT - Stage: {stage['name']}")
        print("=" * 80)
        print()
        print(f"Environment: {stage['environment']}")
        print(f"Traffic: {stage['traffic_percent']}%")
        print(f"Duration: {stage['duration_hours']} hours")
        print()
        
        if not auto_approve:
            response = input("Proceed with deployment? (yes/no): ")
            if response.lower() != "yes":
                print("Deployment cancelled.")
                return False
        
        # Deploy each experiment
        print("Deploying experiments...")
        print()
        
        for exp_id, exp_config in self.experiments.items():
            if not exp_config.get("enabled", False):
                print(f"[SKIP] {exp_config['name']} (disabled)")
                continue
            
            print(f"[DEPLOY] {exp_config['name']}")
            print(f"   Status: {exp_config['status']}")
            
            # Simulate deployment (replace with actual deployment logic)
            success = self._deploy_experiment(exp_id, exp_config, stage)
            
            if success:
                print(f"   [PASS] Deployment successful")
            else:
                print(f"   [FAIL] Deployment failed")
                if not auto_approve:
                    response = input("Continue with other experiments? (yes/no): ")
                    if response.lower() != "yes":
                        return False
        
        print()
        print("=" * 80)
        print("VALIDATION")
        print("=" * 80)
        print()
        
        # Validate deployment
        validation_passed = self._validate_stage(stage)
        
        if validation_passed:
            print("[PASS] Stage validation PASSED")
            return True
        else:
            print("[FAIL] Stage validation FAILED")
            print("[WARN] Consider rollback if critical issues detected")
            return False
    
    def _deploy_experiment(self, exp_id: str, exp_config: Dict, stage: Dict) -> bool:
        """Deploy a single experiment."""
        # This is a placeholder - replace with actual deployment logic
        # For example:
        # - Update Kubernetes configs
        # - Apply experiment-specific settings
        # - Enable feature flags
        # - Update vLLM configuration
        
        time.sleep(0.5)  # Simulate deployment time
        return True
    
    def _validate_stage(self, stage: Dict) -> bool:
        """Validate stage deployment against criteria."""
        criteria = stage.get("validation_criteria", {})
        
        # Simulate validation checks
        # In production, this would:
        # - Query metrics from Prometheus
        # - Check error rates
        # - Verify throughput/latency
        # - Check experiment-specific metrics
        
        print("Validating deployment criteria...")
        print()
        
        all_passed = True
        for criterion, expected in criteria.items():
            # Simulate validation (replace with actual checks)
            passed = True  # Placeholder
            status = "[PASS]" if passed else "[FAIL]"
            print(f"{status} {criterion}: {expected}")
            if not passed:
                all_passed = False
        
        return all_passed
    
    def rollback(self, experiment_id: Optional[str] = None) -> bool:
        """Rollback deployment."""
        print("=" * 80)
        print("PHASE 1 DEPLOYMENT - Rollback")
        print("=" * 80)
        print()
        
        if experiment_id:
            print(f"Rolling back experiment: {experiment_id}")
            # Rollback specific experiment
        else:
            print("Rolling back all Phase 1 experiments")
            # Rollback all experiments
        
        # Simulate rollback
        time.sleep(1)
        print("[PASS] Rollback completed")
        return True
    
    def get_status(self) -> Dict:
        """Get current deployment status."""
        return {
            "phase": "Phase 1",
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "experiments": {
                exp_id: {
                    "name": exp_config["name"],
                    "status": exp_config["status"],
                    "enabled": exp_config.get("enabled", False),
                }
                for exp_id, exp_config in self.experiments.items()
            },
            "deployment_strategy": self.deployment_strategy,
        }


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(description="Phase 1 Production Deployment")
    parser.add_argument(
        "--config",
        type=Path,
        default=Path(__file__).parent / "phase1_production_config.json",
        help="Path to deployment configuration file"
    )
    parser.add_argument(
        "--stage",
        type=str,
        choices=["staging_validation", "canary_deployment", "gradual_rollout_10pct", 
                 "gradual_rollout_25pct", "full_deployment"],
        help="Deployment stage to execute"
    )
    parser.add_argument(
        "--validate",
        action="store_true",
        help="Validate environment only"
    )
    parser.add_argument(
        "--status",
        action="store_true",
        help="Show deployment status"
    )
    parser.add_argument(
        "--rollback",
        type=str,
        nargs="?",
        const="all",
        help="Rollback deployment (optionally specify experiment ID)"
    )
    parser.add_argument(
        "--auto-approve",
        action="store_true",
        help="Skip confirmation prompts"
    )
    
    args = parser.parse_args()
    
    # Initialize deployment
    deployment = Phase1Deployment(args.config)
    
    # Validate environment
    if not deployment.validate_environment():
        print("[FAIL] Environment validation failed. Please fix issues before deploying.")
        sys.exit(1)
    
    # Handle different commands
    if args.validate:
        print("[PASS] Environment validation passed")
        sys.exit(0)
    
    if args.status:
        status = deployment.get_status()
        print(json.dumps(status, indent=2))
        sys.exit(0)
    
    if args.rollback:
        experiment_id = None if args.rollback == "all" else args.rollback
        success = deployment.rollback(experiment_id)
        sys.exit(0 if success else 1)
    
    if args.stage:
        success = deployment.deploy_stage(args.stage, args.auto_approve)
        sys.exit(0 if success else 1)
    
    # Default: show help
    parser.print_help()
    sys.exit(1)


if __name__ == "__main__":
    main()
