import os
import yaml
import boto3
from urllib.parse import urlparse


def download_file(s3_uri: str, local_path: str):
    o = urlparse(s3_uri)
    bucket = o.netloc
    path = o.path[1:]

    s3 = boto3.client('s3')
    with open(local_path, 'wb') as f:
        s3.download_fileobj(bucket, path, f)

def upload_file(local_path: str, s3_uri: str):
    o = urlparse(s3_uri)
    bucket = o.netloc
    path = o.path[1:]

    s3 = boto3.client('s3')
    s3.upload_file(local_path, bucket, path)

def main():
    print("CONFIG_PATH:", os.environ['CONFIG_PATH'])
    download_file(os.environ['CONFIG_PATH'], "/tmp/config.yaml")
    with open("/tmp/config.yaml") as stream:
        configs = yaml.safe_load(stream)

    source_output_path = configs["jobOutput"]["source"]
    print("SOURCE_PATH:", source_output_path)
    download_file(source_output_path, "/tmp/source.yaml")
    with open("/tmp/source.yaml") as stream:
        source = yaml.safe_load(stream)

    output_uri = configs["outputPath"]
    print("OUTPUT_PATH:", output_uri)
    with open("/tmp/output.yaml", "w") as file:
        yaml.dump({
            "payload": "this is the output from sink job",
            "source_output": source["payload"],
        }, file)
    upload_file("/tmp/output.yaml", output_uri)

if __name__ == "__main__":
    main()