import json
from unittest import SkipTest
from uuid import uuid4

import boto3

from moto import mock_aws, settings

from . import verify_execution_result


@mock_aws(config={"stepfunctions": {"execute_state_machine": True}})
def test_state_machine_calling_sqs_with_heartbeat():
    if not settings.TEST_DECORATOR_MODE:
        raise SkipTest("No point in testing this in ServerMode")
    sqs = boto3.client("sqs", "us-east-1")
    sfn = boto3.client("stepfunctions", "us-east-1")
    queue_name = f"queue{uuid4()}"
    queue_url = sqs.create_queue(QueueName=queue_name)["QueueUrl"]

    message_txt = "test_message_txt"
    exec_input = {"QueueUrl": queue_url, "Message": message_txt}

    expected_status = "RUNNING"
    tmpl_name = "services/sqs_heartbeat"

    def _verify_result(client, execution, execution_arn):
        resp = sqs.receive_message(QueueUrl=queue_url)
        if "Messages" in resp:
            task_token = json.loads(resp["Messages"][0]["Body"])["TaskToken"]

            sfn.send_task_heartbeat(taskToken=task_token)
            sfn.send_task_failure(taskToken=task_token)
            return True

        return False

    verify_execution_result(
        _verify_result, expected_status, tmpl_name, exec_input=json.dumps(exec_input)
    )


@mock_aws(config={"stepfunctions": {"execute_state_machine": True}})
def test_state_machine_calling_sqs_with_task_success():
    if not settings.TEST_DECORATOR_MODE:
        raise SkipTest("No point in testing this in ServerMode")
    sqs = boto3.client("sqs", "us-east-1")
    sfn = boto3.client("stepfunctions", "us-east-1")
    queue_name = f"queue{uuid4()}"
    queue_url = sqs.create_queue(QueueName=queue_name)["QueueUrl"]

    message_txt = "test_message_txt"
    exec_input = {"QueueUrl": queue_url, "Message": message_txt}

    expected_status = "RUNNING"
    tmpl_name = "services/sqs_heartbeat"

    def _verify_result(client, execution, execution_arn):
        resp = sqs.receive_message(QueueUrl=queue_url)
        if "Messages" in resp:
            task_token = json.loads(resp["Messages"][0]["Body"])["TaskToken"]

            sfn.send_task_success(taskToken=task_token, output="it's done")
            return True

        return False

    verify_execution_result(
        _verify_result, expected_status, tmpl_name, exec_input=json.dumps(exec_input)
    )