import json

import boto3
from prefect import Flow, flatten, task
from prefect.engine import signals
from prefect.engine.results.s3_result import S3Result
from prefect.executors import LocalExecutor
from prefect.run_configs import LocalRun
from prefect.storage import S3


@task(
    name="get_names",
    target="result/{flow_name}/{flow_run_name}/{task_full_name}",
    checkpoint=True,
    result=S3Result("prefect-flow"),
)
def get_names():
    return ["Peter", "James", "Gee"]


@task(
    name="get_skills",
    target="result/{flow_name}/{flow_run_name}/{task_full_name}",
    checkpoint=True,
    result=S3Result("prefect-flow"),
)
def get_skills(name):
    skills_mapping = {
        "Peter": ["IT", "Accounting"],
        "James": ["Mathematics", "English"],
        "Gee": ["Python", "SQL"],
    }
    return skills_mapping.get(name, [])


@task(
    name="invoke_lambda",
    target="result/{flow_name}/{flow_run_name}/{task_full_name}",
    checkpoint=True,
    result=S3Result("prefect-flow"),
)
def invoke_lambda(skill):
    print(f"invoking lambda with skill: {skill}")
    lambda_client = boto3.client("lambda", region_name="ap-southeast-1")
    response = lambda_client.invoke(
        FunctionName="prefect-test", Payload=json.dumps(skill)
    )
    payload = json.loads(response["Payload"].read())
    if payload.get("errorMessage"):
        raise signals.FAIL()
    return payload


with Flow(
    "pikachu",
    executor=LocalExecutor(),
    run_config=LocalRun(),
    storage=S3(
        bucket="prefect-flow",
    ),
) as flow:
    names = get_names()
    skills = get_skills.map(names)
    response = invoke_lambda.map(flatten(skills))


flow.register(
    project_name="jeremy",
    labels=["dev"],
)