import prefect
from prefect.storage import S3, Docker
from prefect.run_configs import ECSRun
from prefect import task, Flow


TASK_ARN = "arn:aws:iam::***:role/ECSTaskS3ECRRole"
RUN_CONFIG = ECSRun(
    run_task_kwargs={
        "cluster": "prefect-agent-cluster",
        "networkConfiguration":{'awsvpcConfiguration': {'assignPublicIp': 'ENABLED', 'subnets': ['subnet-00637e8e1b0c3b2e8', 'subnet-02847af14d62b73c1'], 'securityGroups': []}}
    },
    labels=['s3-flow-storage'],
    task_role_arn=TASK_ARN,
    execution_role_arn='arn:aws:iam::***:role/ECSFullyLoaded',
    # image='prefecthq/prefect:latest-python3.8',
    # memory=512,
    # cpu=256
    )

STORAGE = Docker(
    registry_url='***.<http://dkr.ecr.ap-southeast-2.amazonaws.com/|dkr.ecr.ap-southeast-2.amazonaws.com/>', 
    image_name='prefect-tutorial', 
    image_tag='latest',
    dockerfile='./Dockerfile'
    )


@task
def say_hello():
    logger = prefect.context.get("logger")
    <http://logger.info|logger.info>('This Worked!')



with Flow(
    "s3_pandas", 
    storage=STORAGE,
    run_config=RUN_CONFIG
    ) as flow:

    say_hello()

flow.register(project_name="aws-flow")
# flow.run()