from .base_agent import BaseAgent
from ..aws_connector import AWSConnector
import json
import time

class LambdaAgent(BaseAgent):
    """
    An agent specialized in handling AWS Lambda tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to Lambda.
        """
        if command == 'list_functions':
            return self._list_functions(**kwargs)
        elif command == 'smart_create_function':
            function_name = kwargs.get('function_name', '')
            region = kwargs.get('region', '')
            print(f"You are about to smart-create a Lambda function '{function_name}' in {region}.")
            confirm = input("This will create a new Lambda function and an IAM role. Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._smart_create_function(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart Create Lambda function command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by LambdaAgent.")

    def _smart_create_function(self, region: str, function_name: str, runtime: str = 'python3.9', handler: str = 'main.handler', memory: int = 128, timeout: int = 30):
        """
        Creates a simple 'hello world' Python Lambda function and an associated IAM role.
        """
        print(f"LambdaAgent: Smart creating function '{function_name}' in region {region}...")
        try:
            lambda_client = AWSConnector.get_client('lambda', region_name=region)
            iam_client = AWSConnector.get_client('iam', region_name=region)
            
            # 1. Create IAM Role for Lambda
            role_name = f'{function_name}-role'
            assume_role_policy_document = {
                "Version": "2012-10-17",
                "Statement": [
                    {
                        "Effect": "Allow",
                        "Principal": {
                            "Service": "lambda.amazonaws.com"
                        },
                        "Action": "sts:AssumeRole"
                    }
                ]
            }
            try:
                create_role_response = iam_client.create_role(
                    RoleName=role_name,
                    AssumeRolePolicyDocument=json.dumps(assume_role_policy_document),
                    Description=f'IAM role for Lambda function {function_name}'
                )
                role_arn = create_role_response['Role']['Arn']
                print(f"Created IAM role: {role_arn}")

                # Attach basic execution policy
                iam_client.attach_role_policy(
                    RoleName=role_name,
                    PolicyArn='arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
                )
                print("Attached AWSLambdaBasicExecutionRole policy.")
                time.sleep(10) # Give IAM time to propagate

            except iam_client.exceptions.EntityAlreadyExistsException:
                print(f"IAM role '{role_name}' already exists. Fetching ARN...")
                role_arn = iam_client.get_role(RoleName=role_name)['Role']['Arn']
            
            # 2. Create function code (simple 'hello world')
            zip_file_content = b"""
import json

def handler(event, context):
    return {
        'statusCode': 200,
        'body': json.dumps('Hello from Lambda!')
    }
""".strip()

            # 3. Create Lambda Function
            create_function_response = lambda_client.create_function(
                FunctionName=function_name,
                Runtime=runtime,
                Role=role_arn,
                Handler=handler,
                Code={'ZipFile': zip_file_content},
                Description=f'A smart-created hello world Lambda function for {function_name}',
                Timeout=timeout,
                MemorySize=memory,
                Publish=True
            )
            function_arn = create_function_response['FunctionArn']

            return {"status": "success", "message": f"Lambda function '{function_name}' created successfully.", "function_arn": function_arn, "role_arn": role_arn}
        except Exception as e:
            print(f"Error during smart Lambda function creation: {e}")
            # In a real scenario, add cleanup logic for created resources
            return {"status": "error", "message": str(e)}


    def _list_functions(self, region: str):
        """Lists all Lambda functions in a specified region."""
        print(f"LambdaAgent: Listing functions in region {region}...")
        try:
            lambda_client = AWSConnector.get_client('lambda', region_name=region)
            response = lambda_client.list_functions()
            functions = response.get('Functions', [])
            return {"status": "success", "functions": functions}
        except Exception as e:
            print(f"Error listing Lambda functions: {e}")
            return {"status": "error", "message": str(e)}