# Copyright (c) 2020 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

import time
from datetime import datetime, timedelta, timezone
from typing import Any, Callable, List

import boto3
from mypy_boto3_stepfunctions.type_defs import ExecutionListItemTypeDef
from toolz import valfilter

from czc import json
from czc.locking import Lease
from czc.logging import get_logger

step_functions_client = boto3.client('stepfunctions')
logger = get_logger()


def start_stm(table: Any, stm_arn: str, lock: dict, base_name: str, **params):
    with Lease(table, lock, timedelta(minutes=15)) as lease_id:
        if lease_id is None:
            logger.info(f'STM {stm_arn} for {base_name} is locked for startup by another caller')
            return

        stm = StateMachine(stm_arn, base_name)
        if stm.is_running():
            logger.info(f'STM {stm_arn} for {base_name} is already running')
            return

        execution_arn = stm.start(**params)
        stm.wait_for_running()

    logger.info(f'Successfully started STM {stm_arn} for {base_name}: {execution_arn}')
    return execution_arn


class StateMachine:
    def __init__(self, arn: str, base_name: str):
        self._arn = arn
        self._base_name = base_name
        self._epoch = datetime.now(timezone.utc)

    @property
    def name(self):
        return f'{self._base_name}.{int(time.mktime(self._epoch.timetuple()))}'

    def _find_execution(self,
                        match: Callable[[str], bool],
                        status_filter: str = None,
                        stop_paging: Callable[[List[ExecutionListItemTypeDef]], bool] = lambda _: False,
                        pagination_token: str = None):
        response = step_functions_client.list_executions(**valfilter(lambda x: x is not None, {
            'stateMachineArn': self._arn,
            'statusFilter': status_filter,
            'nextToken': pagination_token
        }))
        if any(match(x['name']) for x in response['executions']):
            return True
        elif not response.get('nextToken') or stop_paging(response['executions']):
            return False
        return self._find_execution(match, status_filter, stop_paging, response['nextToken'])

    def is_running(self):
        return self._find_execution(lambda x: x.startswith(self._base_name), status_filter='RUNNING')

    def start(self, **kwargs):
        response = step_functions_client.start_execution(
            stateMachineArn=self._arn,
            name=self.name,
            input=json.dumps(kwargs)
        )
        return response['executionArn']

    def wait_for_running(self):
        end_time = datetime.now(timezone.utc) + timedelta(seconds=30)

        def stop_paging(executions):
            return executions[-1]['startDate'] < self._epoch - timedelta(minutes=1)

        while not self._find_execution(lambda x: x == self.name, stop_paging=stop_paging):
            time.sleep(1)
            if datetime.now(timezone.utc) > end_time:
                raise TimeoutError()
