from .base_agent import BaseAgent
from ..aws_connector import AWSConnector
import time
import zipfile
import os

class ElasticBeanstalkAgent(BaseAgent):
    """
    An agent specialized in handling AWS Elastic Beanstalk tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to Elastic Beanstalk.
        """
        if command == 'smart_create_web_app':
            app_name = kwargs.get('app_name', '')
            region = kwargs.get('region', '')
            print(f"You are about to smart-create an Elastic Beanstalk web application '{app_name}' in {region}.")
            confirm = input("This will create an application, environment, and deploy a sample app. Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._smart_create_web_app(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart Create Elastic Beanstalk web app command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by ElasticBeanstalkAgent.")

    def _smart_create_web_app(self, region: str, app_name: str, environment_name: str = None, solution_stack_name: str = '64bit Amazon Linux 2 v3.4.1 running Python 3.8', instance_type: str = 't3.micro'):
        """
        Creates an Elastic Beanstalk application, environment, and deploys a sample web application.
        """
        print(f"ElasticBeanstalkAgent: Smart creating web app '{app_name}' in region {region}...")
        try:
            eb_client = AWSConnector.get_client('elasticbeanstalk', region_name=region)
            s3_client = AWSConnector.get_client('s3', region_name=region)

            if environment_name is None:
                environment_name = f'{app_name}-env'

            # 1. Create S3 bucket for application versions (if it doesn't exist)
            default_bucket_name = f'elasticbeanstalk-{region}-{AWSConnector.get_account_id()}'
            try:
                s3_client.head_bucket(Bucket=default_bucket_name)
                print(f"Using existing S3 bucket: {default_bucket_name}")
            except s3_client.exceptions.ClientError as e:
                if e.response['Error']['Code'] == '404':
                    print(f"Creating S3 bucket for Elastic Beanstalk: {default_bucket_name}")
                    if region == "us-east-1":
                        s3_client.create_bucket(Bucket=default_bucket_name)
                    else:
                        location = {'LocationConstraint': region}
                        s3_client.create_bucket(Bucket=default_bucket_name, CreateBucketConfiguration=location)
                else:
                    raise

            # 2. Create a sample application ZIP file
            app_zip_name = f'{app_name}-sample-app.zip'
            with zipfile.ZipFile(app_zip_name, 'w') as zf:
                zf.writestr('application.py', """
import platform
from flask import Flask
app = Flask(__name__)

@app.route('/')
def hello_world():
    return f'Hello from Elastic Beanstalk! Running on Python {platform.python_version()} on {platform.system()} {platform.release()}'

if __name__ == '__main__':
    app.run(debug=True)
""")
                zf.writestr('requirements.txt', "Flask")
            print(f"Created sample application zip: {app_zip_name}")

            # 3. Upload application ZIP to S3
            s3_client.upload_file(app_zip_name, default_bucket_name, app_zip_name)
            print(f"Uploaded {app_zip_name} to S3 bucket {default_bucket_name}.")
            os.remove(app_zip_name) # Clean up local zip file

            # 4. Create Application
            eb_client.create_application(ApplicationName=app_name)
            print(f"Elastic Beanstalk application '{app_name}' created.")

            # 5. Create Application Version
            eb_client.create_application_version(
                ApplicationName=app_name,
                VersionLabel='v1',
                SourceBundle={'S3Bucket': default_bucket_name, 'S3Key': app_zip_name},
                AutoCreateApplication=False,
                Process=True
            )
            print("Application version 'v1' created.")

            # 6. Create Environment
            create_env_response = eb_client.create_environment(
                ApplicationName=app_name,
                EnvironmentName=environment_name,
                VersionLabel='v1',
                SolutionStackName=solution_stack_name,
                OptionSettings=[
                    {'Namespace': 'aws:autoscaling:launchconfiguration', 'OptionName': 'InstanceType', 'Value': instance_type},
                    {'Namespace': 'aws:ec2:vpc', 'OptionName': 'AssociatePublicIpAddress', 'Value': 'true'},
                    {'Namespace': 'aws:elasticbeanstalk:environment', 'OptionName': 'LoadBalancerType', 'Value': 'application'},
                ]
            )
            print(f"Elastic Beanstalk environment '{environment_name}' creation initiated.")

            # Wait for environment to be ready (optional, but good for user feedback)
            print("Waiting for environment to become ready... (this may take several minutes)")
            waiter = eb_client.get_waiter('environment_exists')
            waiter.wait(EnvironmentNames=[environment_name])
            
            env_status = eb_client.describe_environments(EnvironmentNames=[environment_name])['Environments'][0]
            cname = env_status['CNAME']

            return {"status": "success", "message": f"Elastic Beanstalk web app '{app_name}' deployed to environment '{environment_name}'. CNAME: {cname}", "cname": cname}
        except Exception as e:
            print(f"Error during smart Elastic Beanstalk web app creation: {e}")
            return {"status": "error", "message": str(e)}