import boto3
import os
import json
import zipfile
import io
import urllib.request
from datetime import datetime

class LambdaExporter:
    def __init__(self, profile_name=None, region_name='us-east-1'):
        """
        Initialize Lambda Exporter
        
        Args:
            profile_name (str): AWS profile name (optional)
            region_name (str): AWS region name
        """
        try:
            if profile_name:
                session = boto3.Session(profile_name=profile_name)
                self.lambda_client = session.client('lambda', region_name=region_name)
            else:
                self.lambda_client = boto3.client('lambda', region_name=region_name)
        except Exception as e:
            print(f"Error initializing AWS client: {e}")
            raise
    
    def get_all_functions(self):
        """Retrieve all Lambda functions from the account"""
        functions = []
        try:
            paginator = self.lambda_client.get_paginator('list_functions')
            for page in paginator.paginate():
                functions.extend(page['Functions'])
            return functions
        except Exception as e:
            print(f"Error fetching functions: {e}")
            return []
    
    def export_function_code(self, function_name, output_dir):
        """Export function code to local directory"""
        try:
            # Get function details
            response = self.lambda_client.get_function(FunctionName=function_name)
            
            # Download code
            code_url = response['Code']['Location']
            with urllib.request.urlopen(code_url) as response:
                code_data = response.read()
            
            # Extract if it's a zip file
            if zipfile.is_zipfile(io.BytesIO(code_data)):
                with zipfile.ZipFile(io.BytesIO(code_data)) as zip_ref:
                    zip_ref.extractall(output_dir)
                return True
            else:
                # Save as raw file
                with open(os.path.join(output_dir, 'function_code.bin'), 'wb') as f:
                    f.write(code_data)
                return True
                
        except Exception as e:
            print(f"Error exporting code for {function_name}: {e}")
            return False
    
    def export_function_config(self, function_name, output_dir):
        """Export function configuration"""
        try:
            config = self.lambda_client.get_function_configuration(
                FunctionName=function_name
            )
            
            # Convert to serializable format
            config_serializable = json.loads(json.dumps(config, default=str))
            
            with open(os.path.join(output_dir, 'configuration.json'), 'w') as f:
                json.dump(config_serializable, f, indent=2)
            
            return True
        except Exception as e:
            print(f"Error exporting config for {function_name}: {e}")
            return False
    
    def export_all_functions(self, base_dir='lambda_export'):
        """Export all Lambda functions"""
        functions = self.get_all_functions()
        
        if not functions:
            print("No functions found or error retrieving functions")
            return
        
        print(f"Found {len(functions)} Lambda functions")
        
        # Create export directory with timestamp
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        export_dir = os.path.join(base_dir, f"export_{timestamp}")
        
        success_count = 0
        
        for func in functions:
            function_name = func['FunctionName']
            print(f"\nProcessing: {function_name}")
            
            # Create function directory
            func_dir = os.path.join(export_dir, function_name)
            os.makedirs(func_dir, exist_ok=True)
            
            # Export configuration
            config_success = self.export_function_config(function_name, func_dir)
            
            # Export code
            code_success = self.export_function_code(function_name, func_dir)
            
            if config_success and code_success:
                success_count += 1
                print(f"✅ Successfully exported: {function_name}")
            else:
                print(f"⚠ Partially exported: {function_name}")
        
        # Create export summary
        summary = {
            'export_date': datetime.now().isoformat(),
            'total_functions': len(functions),
            'successfully_exported': success_count,
            'export_directory': export_dir
        }
        
        summary_file = os.path.join(export_dir, 'export_summary.json')
        with open(summary_file, 'w') as f:
            json.dump(summary, f, indent=2)
        
        print(f"\n Export completed!")
        print(f" Location: {os.path.abspath(export_dir)}")
        print(f" Summary: {success_count}/{len(functions)} functions exported successfully")

# Usage
if __name__ == "__main__":
    # For specific AWS profile (optional)
    # exporter = LambdaExporter(profile_name='your-profile-name')
    
    # Using default credentials
    exporter = LambdaExporter()
    exporter.export_all_functions()