"""
API Deployment

Advanced deployment tools for ML model APIs including containerization,
cloud deployment, and monitoring setup.
"""

import os
import json
import subprocess
from typing import Dict, List, Any, Optional
from pathlib import Path
import yaml
import shutil
from datetime import datetime
import requests
from dataclasses import dataclass, asdict


@dataclass
class DeploymentConfig:
    """Configuration for API deployment"""
    name: str
    environment: str  # 'development', 'staging', 'production'
    api_key: Optional[str] = None
    enable_auth: bool = False
    enable_rate_limiting: bool = True
    rate_limit_per_minute: int = 60
    enable_monitoring: bool = True
    replicas: int = 1
    cpu_limit: str = "500m"
    memory_limit: str = "512Mi"
    port: int = 8000
    domain: Optional[str] = None
    ssl_enabled: bool = False


class APIDeployment:
    """
    Advanced API deployment system supporting multiple deployment targets
    including Docker, Kubernetes, and cloud platforms.
    """
    
    def __init__(self, api_generator, workspace_dir: str = "deployment_workspace"):
        """
        Initialize deployment manager
        
        Args:
            api_generator: APIGenerator instance
            workspace_dir: Directory for deployment files
        """
        self.api_generator = api_generator
        self.workspace_dir = Path(workspace_dir)
        self.workspace_dir.mkdir(exist_ok=True)
        
        # Create deployment subdirectories
        self.docker_dir = self.workspace_dir / "docker"
        self.k8s_dir = self.workspace_dir / "kubernetes"
        self.cloud_dir = self.workspace_dir / "cloud"
        
        for dir_path in [self.docker_dir, self.k8s_dir, self.cloud_dir]:
            dir_path.mkdir(exist_ok=True)
    
    def deploy_docker(self, config: DeploymentConfig) -> Dict[str, Any]:
        """
        Deploy API using Docker
        
        Args:
            config: Deployment configuration
            
        Returns:
            Deployment status and information
        """
        deployment_info = {
            "method": "docker",
            "status": "preparing",
            "config": asdict(config),
            "timestamp": datetime.now().isoformat()
        }
        
        try:
            # Generate deployment files
            self._prepare_docker_deployment(config)
            
            # Build Docker image
            image_tag = f"{config.name}:{config.environment}"
            build_result = self._build_docker_image(image_tag, self.docker_dir)
            
            if not build_result['success']:
                deployment_info['status'] = 'failed'
                deployment_info['error'] = build_result['error']
                return deployment_info
            
            # Run container
            run_result = self._run_docker_container(image_tag, config)
            
            if run_result['success']:
                deployment_info['status'] = 'running'
                deployment_info['container_id'] = run_result['container_id']
                deployment_info['port'] = config.port
                deployment_info['url'] = f"http://localhost:{config.port}"
            else:
                deployment_info['status'] = 'failed'
                deployment_info['error'] = run_result['error']
            
        except Exception as e:
            deployment_info['status'] = 'failed'
            deployment_info['error'] = str(e)
        
        return deployment_info
    
    def _prepare_docker_deployment(self, config: DeploymentConfig):
        """Prepare Docker deployment files"""
        # Generate Dockerfile
        dockerfile_content = f"""FROM python:3.9-slim

WORKDIR /app

# Install system dependencies
RUN apt-get update && apt-get install -y \\
    gcc \\
    && rm -rf /var/lib/apt/lists/*

# Copy requirements first for better caching
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy application code
COPY . .

# Create non-root user
RUN useradd --create-home --shell /bin/bash app \\
    && chown -R app:app /app
USER app

# Expose port
EXPOSE {config.port}

# Health check
HEALTHCHECK --interval=30s --timeout=30s --start-period=5s --retries=3 \\
    CMD curl -f http://localhost:{config.port}/health || exit 1

# Run the application
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "{config.port}"]
"""
        
        with open(self.docker_dir / "Dockerfile", 'w') as f:
            f.write(dockerfile_content)
        
        # Generate requirements.txt
        self.api_generator.generate_requirements(self.docker_dir / "requirements.txt")
        
        # Generate main.py
        self._generate_main_script(config, self.docker_dir / "main.py")
        
        # Generate docker-compose.yml
        docker_compose_content = f"""version: '3.8'

services:
  api:
    build: .
    ports:
      - "{config.port}:{config.port}"
    environment:
      - ENVIRONMENT={config.environment}
      - API_KEY={config.api_key or ''}
    deploy:
      replicas: {config.replicas}
      resources:
        limits:
          cpus: '{config.cpu_limit}'
          memory: {config.memory_limit}
    restart: unless-stopped
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:{config.port}/health"]
      interval: 30s
      timeout: 10s
      retries: 3
      start_period: 40s
"""
        
        if config.enable_monitoring:
            docker_compose_content += """
  prometheus:
    image: prom/prometheus:latest
    ports:
      - "9090:9090"
    volumes:
      - ./prometheus.yml:/etc/prometheus/prometheus.yml
  
  grafana:
    image: grafana/grafana:latest
    ports:
      - "3000:3000"
    environment:
      - GF_SECURITY_ADMIN_PASSWORD=admin
"""
        
        with open(self.docker_dir / "docker-compose.yml", 'w') as f:
            f.write(docker_compose_content)
    
    def _build_docker_image(self, tag: str, build_dir: Path) -> Dict[str, Any]:
        """Build Docker image"""
        try:
            result = subprocess.run(
                ["docker", "build", "-t", tag, str(build_dir)],
                capture_output=True,
                text=True,
                timeout=300
            )
            
            if result.returncode == 0:
                return {"success": True, "output": result.stdout}
            else:
                return {"success": False, "error": result.stderr}
                
        except subprocess.TimeoutExpired:
            return {"success": False, "error": "Build timeout"}
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    def _run_docker_container(self, image_tag: str, config: DeploymentConfig) -> Dict[str, Any]:
        """Run Docker container"""
        try:
            # Stop existing container if running
            subprocess.run(
                ["docker", "stop", config.name],
                capture_output=True
            )
            subprocess.run(
                ["docker", "rm", config.name],
                capture_output=True
            )
            
            # Run new container
            cmd = [
                "docker", "run", "-d",
                "--name", config.name,
                "-p", f"{config.port}:{config.port}",
                "--memory", config.memory_limit,
                "--cpus", config.cpu_limit.replace('m', '') + 'm'
            ]
            
            if config.enable_auth and config.api_key:
                cmd.extend(["-e", f"API_KEY={config.api_key}"])
            
            cmd.extend([image_tag])
            
            result = subprocess.run(cmd, capture_output=True, text=True)
            
            if result.returncode == 0:
                container_id = result.stdout.strip()
                return {"success": True, "container_id": container_id}
            else:
                return {"success": False, "error": result.stderr}
                
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    def deploy_kubernetes(self, config: DeploymentConfig) -> Dict[str, Any]:
        """
        Deploy API to Kubernetes
        
        Args:
            config: Deployment configuration
            
        Returns:
            Deployment status and information
        """
        deployment_info = {
            "method": "kubernetes",
            "status": "preparing",
            "config": asdict(config),
            "timestamp": datetime.now().isoformat()
        }
        
        try:
            # Generate Kubernetes manifests
            self._generate_kubernetes_manifests(config)
            
            # Apply manifests
            apply_result = self._apply_kubernetes_manifests()
            
            if apply_result['success']:
                deployment_info['status'] = 'deploying'
                deployment_info['namespace'] = config.environment
                
                # Wait for deployment to be ready
                rollout_result = self._wait_for_deployment(config.name, config.environment)
                
                if rollout_result['success']:
                    deployment_info['status'] = 'running'
                    deployment_info['service_url'] = rollout_result['service_url']
                else:
                    deployment_info['status'] = 'failed'
                    deployment_info['error'] = rollout_result['error']
            else:
                deployment_info['status'] = 'failed'
                deployment_info['error'] = apply_result['error']
            
        except Exception as e:
            deployment_info['status'] = 'failed'
            deployment_info['error'] = str(e)
        
        return deployment_info
    
    def _generate_kubernetes_manifests(self, config: DeploymentConfig):
        """Generate Kubernetes deployment manifests"""
        # Namespace
        namespace_manifest = {
            "apiVersion": "v1",
            "kind": "Namespace",
            "metadata": {
                "name": config.environment
            }
        }
        
        # Deployment
        deployment_manifest = {
            "apiVersion": "apps/v1",
            "kind": "Deployment",
            "metadata": {
                "name": config.name,
                "namespace": config.environment,
                "labels": {
                    "app": config.name
                }
            },
            "spec": {
                "replicas": config.replicas,
                "selector": {
                    "matchLabels": {
                        "app": config.name
                    }
                },
                "template": {
                    "metadata": {
                        "labels": {
                            "app": config.name
                        }
                    },
                    "spec": {
                        "containers": [{
                            "name": config.name,
                            "image": f"{config.name}:latest",
                            "ports": [{
                                "containerPort": config.port
                            }],
                            "resources": {
                                "limits": {
                                    "cpu": config.cpu_limit,
                                    "memory": config.memory_limit
                                }
                            },
                            "env": [
                                {
                                    "name": "ENVIRONMENT",
                                    "value": config.environment
                                }
                            ],
                            "livenessProbe": {
                                "httpGet": {
                                    "path": "/health",
                                    "port": config.port
                                },
                                "initialDelaySeconds": 30,
                                "periodSeconds": 10
                            },
                            "readinessProbe": {
                                "httpGet": {
                                    "path": "/health",
                                    "port": config.port
                                },
                                "initialDelaySeconds": 5,
                                "periodSeconds": 5
                            }
                        }]
                    }
                }
            }
        }
        
        # Service
        service_manifest = {
            "apiVersion": "v1",
            "kind": "Service",
            "metadata": {
                "name": f"{config.name}-service",
                "namespace": config.environment,
                "labels": {
                    "app": config.name
                }
            },
            "spec": {
                "selector": {
                    "app": config.name
                },
                "ports": [{
                    "protocol": "TCP",
                    "port": 80,
                    "targetPort": config.port
                }],
                "type": "ClusterIP"
            }
        }
        
        # Ingress (if domain provided)
        manifests = [namespace_manifest, deployment_manifest, service_manifest]
        
        if config.domain:
            ingress_manifest = {
                "apiVersion": "networking.k8s.io/v1",
                "kind": "Ingress",
                "metadata": {
                    "name": f"{config.name}-ingress",
                    "namespace": config.environment,
                    "annotations": {
                        "nginx.ingress.kubernetes.io/rewrite-target": "/"
                    }
                },
                "spec": {
                    "rules": [{
                        "host": config.domain,
                        "http": {
                            "paths": [{
                                "path": "/",
                                "pathType": "Prefix",
                                "backend": {
                                    "service": {
                                        "name": f"{config.name}-service",
                                        "port": {
                                            "number": 80
                                        }
                                    }
                                }
                            }]
                        }
                    }]
                }
            }
            
            if config.ssl_enabled:
                ingress_manifest["spec"]["tls"] = [{
                    "hosts": [config.domain],
                    "secretName": f"{config.name}-tls"
                }]
            
            manifests.append(ingress_manifest)
        
        # Write manifests to files
        for i, manifest in enumerate(manifests):
            kind = manifest["kind"].lower()
            filename = f"{kind}.yaml"
            with open(self.k8s_dir / filename, 'w') as f:
                yaml.dump(manifest, f, default_flow_style=False)
    
    def _apply_kubernetes_manifests(self) -> Dict[str, Any]:
        """Apply Kubernetes manifests"""
        try:
            result = subprocess.run(
                ["kubectl", "apply", "-f", str(self.k8s_dir)],
                capture_output=True,
                text=True,
                timeout=60
            )
            
            if result.returncode == 0:
                return {"success": True, "output": result.stdout}
            else:
                return {"success": False, "error": result.stderr}
                
        except subprocess.TimeoutExpired:
            return {"success": False, "error": "Apply timeout"}
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    def _wait_for_deployment(self, name: str, namespace: str, timeout: int = 300) -> Dict[str, Any]:
        """Wait for Kubernetes deployment to be ready"""
        try:
            result = subprocess.run(
                ["kubectl", "rollout", "status", f"deployment/{name}", 
                 "-n", namespace, "--timeout", f"{timeout}s"],
                capture_output=True,
                text=True
            )
            
            if result.returncode == 0:
                # Get service URL
                service_result = subprocess.run(
                    ["kubectl", "get", "service", f"{name}-service", 
                     "-n", namespace, "-o", "jsonpath='{.status.loadBalancer.ingress[0].ip}'"],
                    capture_output=True,
                    text=True
                )
                
                service_url = None
                if service_result.returncode == 0 and service_result.stdout.strip():
                    service_url = f"http://{service_result.stdout.strip()}"
                else:
                    # Try minikube or local cluster
                    service_result = subprocess.run(
                        ["kubectl", "get", "service", f"{name}-service", 
                         "-n", namespace, "-o", "jsonpath='{.spec.clusterIP}'"],
                        capture_output=True,
                        text=True
                    )
                    if service_result.returncode == 0 and service_result.stdout.strip():
                        service_url = f"http://{service_result.stdout.strip()}"
                
                return {"success": True, "service_url": service_url}
            else:
                return {"success": False, "error": result.stderr}
                
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    def deploy_cloud(self, config: DeploymentConfig, platform: str = "aws") -> Dict[str, Any]:
        """
        Deploy API to cloud platform
        
        Args:
            config: Deployment configuration
            platform: Cloud platform ('aws', 'gcp', 'azure')
            
        Returns:
            Deployment status and information
        """
        deployment_info = {
            "method": f"cloud_{platform}",
            "status": "preparing",
            "config": asdict(config),
            "timestamp": datetime.now().isoformat()
        }
        
        try:
            if platform.lower() == "aws":
                result = self._deploy_aws(config)
            elif platform.lower() == "gcp":
                result = self._deploy_gcp(config)
            elif platform.lower() == "azure":
                result = self._deploy_azure(config)
            else:
                raise ValueError(f"Unsupported cloud platform: {platform}")
            
            deployment_info.update(result)
            
        except Exception as e:
            deployment_info['status'] = 'failed'
            deployment_info['error'] = str(e)
        
        return deployment_info
    
    def _deploy_aws(self, config: DeploymentConfig) -> Dict[str, Any]:
        """Deploy to AWS using ECS"""
        # Generate ECS task definition
        task_definition = {
            "family": config.name,
            "networkMode": "awsvpc",
            "requiresCompatibilities": ["FARGATE"],
            "cpu": config.cpu_limit.replace('m', ''),
            "memory": config.memory_limit.replace('Mi', ''),
            "executionRoleArn": f"arn:aws:iam::ACCOUNT:role/ecsTaskExecutionRole",
            "containerDefinitions": [{
                "name": config.name,
                "image": f"{config.name}:latest",
                "portMappings": [{
                    "containerPort": config.port,
                    "protocol": "tcp"
                }],
                "environment": [
                    {"name": "ENVIRONMENT", "value": config.environment}
                ],
                "logConfiguration": {
                    "logDriver": "awslogs",
                    "options": {
                        "awslogs-group": f"/ecs/{config.name}",
                        "awslogs-region": "us-west-2",
                        "awslogs-stream-prefix": "ecs"
                    }
                }
            }]
        }
        
        # Save task definition
        with open(self.cloud_dir / "ecs-task-def.json", 'w') as f:
            json.dump(task_definition, f, indent=2)
        
        # Generate CloudFormation template
        template = {
            "AWSTemplateFormatVersion": "2010-09-09",
            "Description": f"ECS deployment for {config.name}",
            "Resources": {
                "ECSCluster": {
                    "Type": "AWS::ECS::Cluster",
                    "Properties": {
                        "ClusterName": f"{config.name}-cluster"
                    }
                },
                "TaskDefinition": {
                    "Type": "AWS::ECS::TaskDefinition",
                    "Properties": task_definition
                },
                "Service": {
                    "Type": "AWS::ECS::Service",
                    "Properties": {
                        "Cluster": {"Ref": "ECSCluster"},
                        "TaskDefinition": {"Ref": "TaskDefinition"},
                        "DesiredCount": config.replicas,
                        "LaunchType": "FARGATE",
                        "NetworkConfiguration": {
                            "AwsvpcConfiguration": {
                                "AssignPublicIp": "ENABLED",
                                "SecurityGroupIds": ["sg-12345678"],
                                "Subnets": ["subnet-12345678"]
                            }
                        }
                    }
                }
            }
        }
        
        with open(self.cloud_dir / "cloudformation.yaml", 'w') as f:
            yaml.dump(template, f)
        
        return {
            "status": "ready_for_deployment",
            "instructions": [
                "1. Push Docker image to AWS ECR",
                "2. Run: aws cloudformation deploy --template-file cloudformation.yaml --stack-name ml-api-stack",
                "3. Update task definition with correct image URI"
            ]
        }
    
    def _deploy_gcp(self, config: DeploymentConfig) -> Dict[str, Any]:
        """Deploy to Google Cloud Platform using Cloud Run"""
        # Generate cloudbuild.yaml
        cloudbuild_content = f"""steps:
- name: 'gcr.io/cloud-builders/docker'
  args: ['build', '-t', 'gcr.io/$PROJECT_ID/{config.name}', '.']
- name: 'gcr.io/cloud-builders/docker'
  args: ['push', 'gcr.io/$PROJECT_ID/{config.name}']
- name: 'gcr.io/google.com/cloudsdktool/cloud-sdk'
  entrypoint: 'gcloud'
  args:
  - 'run'
  - 'deploy'
  - '{config.name}'
  - '--image=gcr.io/$PROJECT_ID/{config.name}'
  - '--region=us-central1'
  - '--platform=managed'
  - '--allow-unauthenticated'
  - '--memory={config.memory_limit}'
  - '--cpu={config.cpu_limit}'
  - '--max-instances={config.replicas}'
"""
        
        with open(self.cloud_dir / "cloudbuild.yaml", 'w') as f:
            f.write(cloudbuild_content)
        
        return {
            "status": "ready_for_deployment",
            "instructions": [
                "1. Enable Cloud Run and Cloud Build APIs",
                "2. Run: gcloud builds submit --config cloudbuild.yaml",
                "3. Service will be deployed automatically"
            ]
        }
    
    def _deploy_azure(self, config: DeploymentConfig) -> Dict[str, Any]:
        """Deploy to Azure using Container Instances"""
        # Generate ARM template
        template = {
            "$schema": "https://schema.management.azure.com/schemas/2019-04-01/deploymentTemplate.json#",
            "contentVersion": "1.0.0.0",
            "parameters": {
                "containerName": {
                    "type": "string",
                    "defaultValue": config.name
                }
            },
            "resources": [{
                "type": "Microsoft.ContainerInstance/containerGroups",
                "apiVersion": "2021-09-01",
                "name": "[parameters('containerName')]",
                "location": "[resourceGroup().location]",
                "properties": {
                    "containers": [{
                        "name": config.name,
                        "properties": {
                            "image": f"{config.name}:latest",
                            "resources": {
                                "requests": {
                                    "cpu": config.cpu_limit.replace('m', '') + 'm',
                                    "memoryInGB": float(config.memory_limit.replace('Mi', '')) / 1024
                                }
                            },
                            "ports": [{
                                "port": config.port,
                                "protocol": "TCP"
                            }]
                        }
                    }],
                    "osType": "Linux",
                    "ipAddress": {
                        "type": "Public",
                        "ports": [{
                            "port": config.port,
                            "protocol": "TCP"
                        }]
                    }
                }
            }]
        }
        
        with open(self.cloud_dir / "azure-template.json", 'w') as f:
            json.dump(template, f, indent=2)
        
        return {
            "status": "ready_for_deployment",
            "instructions": [
                "1. Push Docker image to Azure Container Registry",
                "2. Run: az deployment group create --resource-group myResourceGroup --template-file azure-template.json"
            ]
        }
    
    def _generate_main_script(self, config: DeploymentConfig, output_path: Path):
        """Generate main.py script with configuration"""
        main_script_content = f'''from fastapi import FastAPI
import uvicorn
import os
from api_generator import APIGenerator

# Configuration
ENVIRONMENT = os.getenv("ENVIRONMENT", "development")
API_KEY = os.getenv("API_KEY", "")

# Initialize API generator
api_generator = APIGenerator(
    model_path="model.pkl",
    model_id="{config.name}"
)

# Generate FastAPI app
app = api_generator.generate_api(
    title="ML Model API",
    description="Auto-generated API for ML model",
    version="1.0.0"
)

# Add authentication if enabled
if config.enable_auth and API_KEY:
    api_generator.add_authentication(API_KEY)

# Add rate limiting if enabled
if config.enable_rate_limiting:
    api_generator.enable_rate_limiting({config.rate_limit_per_minute})

# Add monitoring
if config.enable_monitoring:
    api_generator.add_monitoring_middleware()

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port={config.port})
'''
        
        with open(output_path, 'w') as f:
            f.write(main_script_content)
    
    def get_deployment_status(self, deployment_id: str) -> Dict[str, Any]:
        """Get status of a deployment"""
        # This would integrate with deployment monitoring
        # For now, return a placeholder
        return {
            "deployment_id": deployment_id,
            "status": "unknown",
            "timestamp": datetime.now().isoformat(),
            "message": "Status tracking not implemented in this version"
        }
    
    def cleanup_deployment(self, config: DeploymentConfig, method: str = "docker"):
        """Clean up deployment resources"""
        cleanup_info = {
            "method": method,
            "status": "cleaning",
            "timestamp": datetime.now().isoformat()
        }
        
        try:
            if method == "docker":
                # Stop and remove container
                subprocess.run(["docker", "stop", config.name], capture_output=True)
                subprocess.run(["docker", "rm", config.name], capture_output=True)
                
                # Remove image
                subprocess.run(["docker", "rmi", f"{config.name}:{config.environment}"], 
                             capture_output=True)
                
                cleanup_info['status'] = 'completed'
                
            elif method == "kubernetes":
                # Delete Kubernetes resources
                subprocess.run(["kubectl", "delete", "-f", str(self.k8s_dir)], 
                             capture_output=True)
                
                cleanup_info['status'] = 'completed'
            
        except Exception as e:
            cleanup_info['status'] = 'failed'
            cleanup_info['error'] = str(e)
        
        return cleanup_info
