# Copyright (c) CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

from typing import Optional

from .types import AsyncQueryOperation
from czc import json
from czc.logging import get_logger

import boto3

logger = get_logger()


step_functions_client = boto3.client("stepfunctions")

MAX_NAME_SIZE = 80


def start(stm_arn: str, operation: AsyncQueryOperation) -> Optional[str]:
    query_id = operation.query_id.query_id
    name_parts = []
    if operation.context.organization_id is not None:
        name_parts.append(f"org_{operation.context.organization_id}")
    name_parts.append(query_id)
    name = ".".join(name_parts)

    logger.info(f"Starting state machine.  Name: {name}")
    response = step_functions_client.start_execution(
        stateMachineArn=stm_arn, name=name, input=json.dumps(operation)
    )
    return response["executionArn"]


def send_success(token: str, output: Optional[dict]):
    output_data = json.dumps(output or {})
    step_functions_client.send_task_success(taskToken=token, output=output_data)
