import json
from typing import Any

from src.worker.connectors.aws import stepfunctions


class StaleTaskToken(Exception):
    pass


_STALE_TOKEN_ERRORS = (
    stepfunctions.exceptions.TaskDoesNotExist,
    stepfunctions.exceptions.TaskTimedOut,
)


def send_task_success(*, task_token: str, output: dict[str, Any]) -> None:
    try:
        stepfunctions.send_task_success(taskToken=task_token, output=json.dumps(output))
    except _STALE_TOKEN_ERRORS as exc:
        raise StaleTaskToken(str(exc)) from exc


def send_task_failure(*, task_token: str, error: str, cause: str) -> None:
    try:
        stepfunctions.send_task_failure(taskToken=task_token, error=error, cause=cause)
    except _STALE_TOKEN_ERRORS as exc:
        raise StaleTaskToken(str(exc)) from exc


def send_task_heartbeat(*, task_token: str) -> None:
    try:
        stepfunctions.send_task_heartbeat(taskToken=task_token)
    except _STALE_TOKEN_ERRORS as exc:
        raise StaleTaskToken(str(exc)) from exc