import logging
from typing import Dict, Any
import json
from zenyatta.common import get_airflow_connection
from zenyatta.aws import boto_client
from zenyatta.aws.s3 import get_base_s3_path


def generate_etl_finished_message(s3_path: str, job_id: str) -> str:
    return json.dumps({'s3_path': s3_path, 'job_id': job_id})


def generate_s3_path_for_finished_message(s3_meta: Dict[str, str], ts_nodash: str, conn_id: str) -> str:
    return "s3://" + s3_meta['bucket'] + '/' + get_base_s3_path(s3_meta['s3_key'], ts_nodash, conn_id)


def send_sns_message(message: str, sns_arn: str, role_arn: str) -> Dict:
    sns, _ = boto_client('sns', role_arn)  # don't care about the creds
    logging.info("sending {message} to {sns_arn}".format(**locals()))
    resp = sns.publish(TopicArn=sns_arn, Message=message)
    logging.info("sent {message} and response is: {resp}".format(**locals()))
    return resp


def send_sns_message_wrapper(conn_id: str=None,  # string representing airflow.models.Connection.conn_id
                             s3_meta: Dict[str, Any]=None,
                             sns_arn: str=None,
                             role_arn: str=None,
                             **kwargs) -> Dict[str, Any]:
    connection = get_airflow_connection(conn_id)
    message = generate_etl_finished_message(generate_s3_path_for_finished_message(
        s3_meta, kwargs.get('ts_nodash'),
        connection.conn_id), connection.conn_id)
    sns_resp = send_sns_message(message, sns_arn, role_arn)
    return sns_resp
