#!/usr/bin/env python3
"""
Shutdown script for terminating crawler processes.
This script can be called from GitHub workflows to gracefully stop crawler operations.
"""

import os
import sys
import json
import logging
import boto3
from datetime import datetime
from typing import Dict, Any

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


def create_shutdown_signal(bucket_name: str, resource_tag: str) -> str:
    """Create a shutdown signal file in S3 to indicate crawler should stop"""
    try:
        s3 = boto3.client('s3')
        
        # Create shutdown signal
        shutdown_signal = {
            "shutdown": True,
            "timestamp": datetime.now().isoformat(),
            "reason": "Manual shutdown via GitHub workflow",
            "resource_tag": resource_tag
        }
        
        # Upload shutdown signal to S3
        key = "shutdown_signal.json"
        s3.put_object(
            Bucket=bucket_name,
            Key=key,
            Body=json.dumps(shutdown_signal, indent=2),
            ContentType="application/json",
            Tagging=f"resource-tag={resource_tag}"
        )
        
        logger.info(f"Shutdown signal created in s3://{bucket_name}/{key}")
        return f"s3://{bucket_name}/{key}"
        
    except Exception as e:
        logger.error(f"Failed to create shutdown signal: {e}")
        raise


def check_running_processes(region: str, resource_tag: str) -> Dict[str, Any]:
    """Check for running EC2 instances or Lambda functions"""
    try:
        ec2 = boto3.client('ec2', region_name=region)
        lambda_client = boto3.client('lambda', region_name=region)
        
        # Check for running EC2 instances with our tag
        ec2_response = ec2.describe_instances(
            Filters=[
                {'Name': 'tag:resource-tag', 'Values': [resource_tag]},
                {'Name': 'instance-state-name', 'Values': ['running', 'pending']}
            ]
        )
        
        running_instances = []
        for reservation in ec2_response['Reservations']:
            for instance in reservation['Instances']:
                running_instances.append({
                    'instance_id': instance['InstanceId'],
                    'state': instance['State']['Name'],
                    'launch_time': instance['LaunchTime'].isoformat()
                })
        
        # Check for Lambda functions (if any are running)
        lambda_functions = []
        try:
            lambda_response = lambda_client.list_functions()
            for func in lambda_response['Functions']:
                if resource_tag in func.get('Description', ''):
                    lambda_functions.append({
                        'function_name': func['FunctionName'],
                        'runtime': func['Runtime'],
                        'last_modified': func['LastModified']
                    })
        except Exception as e:
            logger.warning(f"Could not list Lambda functions: {e}")
        
        return {
            'running_ec2_instances': running_instances,
            'lambda_functions': lambda_functions
        }
        
    except Exception as e:
        logger.error(f"Failed to check running processes: {e}")
        return {'error': str(e)}


def terminate_ec2_instances(region: str, resource_tag: str) -> Dict[str, Any]:
    """Terminate EC2 instances with the specified tag"""
    try:
        ec2 = boto3.client('ec2', region_name=region)
        
        # Get instances to terminate
        response = ec2.describe_instances(
            Filters=[
                {'Name': 'tag:resource-tag', 'Values': [resource_tag]},
                {'Name': 'instance-state-name', 'Values': ['running', 'pending', 'stopping']}
            ]
        )
        
        instance_ids = []
        for reservation in response['Reservations']:
            for instance in reservation['Instances']:
                instance_ids.append(instance['InstanceId'])
        
        if not instance_ids:
            logger.info("No EC2 instances found to terminate")
            return {'terminated_instances': []}
        
        # Terminate instances
        terminate_response = ec2.terminate_instances(InstanceIds=instance_ids)
        
        terminated_instances = []
        for instance in terminate_response['TerminatingInstances']:
            terminated_instances.append({
                'instance_id': instance['InstanceId'],
                'previous_state': instance['PreviousState']['Name'],
                'current_state': instance['CurrentState']['Name']
            })
        
        logger.info(f"Terminated {len(terminated_instances)} EC2 instances")
        return {'terminated_instances': terminated_instances}
        
    except Exception as e:
        logger.error(f"Failed to terminate EC2 instances: {e}")
        return {'error': str(e)}


def main():
    """Main shutdown function"""
    try:
        # Get configuration from environment
        region = os.getenv("AWS_REGION", "us-east-1")
        bucket_name = os.getenv("S3_BUCKET_PROCESSED", "ransom-analysis-processed")
        resource_tag = os.getenv("RESOURCE_TAG", "ransom-analysis:local")
        
        # Handle empty environment variables (common in test mode)
        if not bucket_name or bucket_name.strip() == "":
            bucket_name = "ransom-analysis-processed"
        if not resource_tag or resource_tag.strip() == "":
            resource_tag = "ransom-analysis:local"
        
        logger.info(f"Starting crawler shutdown process...")
        logger.info(f"Region: {region}")
        logger.info(f"Bucket: {bucket_name}")
        logger.info(f"Resource Tag: {resource_tag}")
        
        # Check current running processes
        logger.info("Checking for running processes...")
        running_processes = check_running_processes(region, resource_tag)
        logger.info(f"Running processes: {json.dumps(running_processes, indent=2)}")
        
        # Create shutdown signal
        logger.info("Creating shutdown signal...")
        shutdown_signal_path = create_shutdown_signal(bucket_name, resource_tag)
        
        # Terminate EC2 instances if any are running
        if running_processes.get('running_ec2_instances'):
            logger.info("Terminating EC2 instances...")
            termination_result = terminate_ec2_instances(region, resource_tag)
            logger.info(f"Termination result: {json.dumps(termination_result, indent=2)}")
        else:
            logger.info("No EC2 instances to terminate")
        
        # Summary
        result = {
            'status': 'success',
            'shutdown_signal_created': shutdown_signal_path,
            'running_processes_before_shutdown': running_processes,
            'timestamp': datetime.now().isoformat()
        }
        
        logger.info("Crawler shutdown process completed successfully")
        print(json.dumps(result, indent=2))
        
    except Exception as e:
        logger.error(f"Shutdown process failed: {e}")
        result = {
            'status': 'error',
            'error': str(e),
            'timestamp': datetime.now().isoformat()
        }
        print(json.dumps(result, indent=2))
        sys.exit(1)


if __name__ == "__main__":
    main()
