import datetime
import time
from typing import Any, Callable, Dict, Iterable, Optional, Tuple

import pendulum

import prefect
from prefect.client import Client
from prefect.core import Edge, Task
from prefect.engine.result import Result
from prefect.engine.runner import ENDRUN, call_state_handlers
from prefect.engine.state import (
    Cached,
    ClientFailed,
    Failed,
    Mapped,
    Queued,
    Retrying,
    State,
)
from prefect.engine.task_runner import TaskRunner, TaskRunnerInitializeResult
from prefect.utilities.executors import tail_recursive
from prefect.utilities.graphql import with_args


class CloudTaskRunner(TaskRunner):
    """
    TaskRunners handle the execution of Tasks and determine the State of a Task
    before, during and after the Task is run.

    In particular, through the TaskRunner you can specify the states of any upstream dependencies,
    and what state the Task should be initialized with.

    Args:
        - task (Task): the Task to be run / executed
        - state_handlers (Iterable[Callable], optional): A list of state change handlers
            that will be called whenever the task changes state, providing an
            opportunity to inspect or modify the new state. The handler
            will be passed the task runner instance, the old (prior) state, and the new
            (current) state, with the following signature: `state_handler(TaskRunner, old_state, new_state) -> State`;
            If multiple functions are passed, then the `new_state` argument will be the
            result of the previous handler.
        - result (Result, optional): the result instance used to retrieve and store task results during execution;
            if not provided, will default to the one on the provided Task
        - default_result (Result, optional): the fallback result type to use for retrieving and storing state results
            during execution (to be used on upstream inputs if they don't provide their own results)
    """

    def __init__(
        self,
        task: Task,
        state_handlers: Iterable[Callable] = None,
        result: Result = None,
        default_result: Result = None,
    ) -> None:
        self.client = Client()
        super().__init__(
            task=task,
            state_handlers=state_handlers,
            result=result,
            default_result=default_result,
        )

    def _heartbeat(self) -> bool:
        try:
            task_run_id = self.task_run_id  # type: str
            self.heartbeat_cmd = ["prefect", "heartbeat", "task-run", "-i", task_run_id]
            self.client.update_task_run_heartbeat(task_run_id)

            # use empty string for testing purposes
            flow_run_id = prefect.context.get("flow_run_id", "")  # type: str
            query = {
                "query": {
                    with_args("flow_run_by_pk", {"id": flow_run_id}): {
                        "flow": {"settings": True},
                    }
                }
            }
            flow_run = self.client.graphql(query).data.flow_run_by_pk
            if not flow_run.flow.settings.get("heartbeat_enabled", True):
                return False
            return True
        except Exception as exc:
            self.logger.exception(
                "Heartbeat failed for Task '{}'".format(self.task.name)
            )
            return False

    def call_runner_target_handlers(self, old_state: State, new_state: State) -> State:
        """
        A special state handler that the TaskRunner uses to call its task's state handlers.
        This method is called as part of the base Runner's `handle_state_change()` method.

        Args:
            - old_state (State): the old (previous) state
            - new_state (State): the new (current) state

        Returns:
            - State: the new state
        """
        raise_on_exception = prefect.context.get("raise_on_exception", False)

        try:
            new_state = super().call_runner_target_handlers(
                old_state=old_state, new_state=new_state
            )
        except Exception as exc:
            msg = "Exception raised while calling state handlers: {}".format(repr(exc))
            self.logger.exception(msg)
            if raise_on_exception:
                raise exc
            new_state = Failed(msg, result=exc)

        task_run_id = prefect.context.get("task_run_id")
        version = prefect.context.get("task_run_version")

        try:
            cloud_state = new_state
            state = self.client.set_task_run_state(
                task_run_id=task_run_id,
                version=version,
                state=cloud_state,
                cache_for=self.task.cache_for,
            )
        except Exception as exc:
            self.logger.exception(
                "Failed to set task state with error: {}".format(repr(exc))
            )
            raise ENDRUN(state=ClientFailed(state=new_state))

        if state.is_queued():
            state.state = old_state  # type: ignore
            raise ENDRUN(state=state)

        if version is not None:
            prefect.context.update(task_run_version=version + 1)  # type: ignore

        return new_state

    def initialize_run(  # type: ignore
        self, state: Optional[State], context: Dict[str, Any]
    ) -> TaskRunnerInitializeResult:
        """
        Initializes the Task run by initializing state and context appropriately.

        Args:
            - state (Optional[State]): the initial state of the run
            - context (Dict[str, Any]): the context to be updated with relevant information

        Returns:
            - tuple: a tuple of the updated state, context, and upstream_states objects
        """

        # if the map_index is not None, this is a dynamic task and we need to load
        # task run info for it
        map_index = context.get("map_index")
        if map_index not in [-1, None]:
            try:
                task_run_info = self.client.get_task_run_info(
                    flow_run_id=context.get("flow_run_id", ""),
                    task_id=context.get("task_id", ""),
                    map_index=map_index,
                )

                # if state was provided, keep it; otherwise use the one from db
                state = state or task_run_info.state  # type: ignore
                context.update(
                    task_run_id=task_run_info.id,  # type: ignore
                    task_run_version=task_run_info.version,  # type: ignore
                )
            except Exception as exc:
                self.logger.exception(
                    "Failed to retrieve task state with error: {}".format(repr(exc))
                )
                if state is None:
                    state = Failed(
                        message="Could not retrieve state from Prefect Cloud",
                        result=exc,
                    )
                raise ENDRUN(state=state)

        # we assign this so it can be shared with heartbeat thread
        self.task_run_id = context.get("task_run_id", "")  # type: str
        context.update(checkpointing=True)

        return super().initialize_run(state=state, context=context)

    @call_state_handlers
    def check_task_is_cached(self, state: State, inputs: Dict[str, Result]) -> State:
        """
        Checks if task is cached in the DB and whether any of the caches are still valid.

        Args:
            - state (State): the current state of this task
            - inputs (Dict[str, Result]): a dictionary of inputs whose keys correspond
                to the task's `run()` arguments.

        Returns:
            - State: the state of the task after running the check

        Raises:
            - ENDRUN: if the task is not ready to run
        """
        if state.is_cached() is True:
            assert isinstance(state, Cached)  # mypy assert
            sanitized_inputs = {key: res.value for key, res in inputs.items()}
            if self.task.cache_validator(
                state, sanitized_inputs, prefect.context.get("parameters")
            ):
                state = state.load_result(self.result)
                return state

        if self.task.cache_for is not None:
            oldest_valid_cache = datetime.datetime.utcnow() - self.task.cache_for
            cached_states = self.client.get_latest_cached_states(
                task_id=prefect.context.get("task_id", ""),
                cache_key=self.task.cache_key,
                created_after=oldest_valid_cache,
            )

            if not cached_states:
                self.logger.debug(
                    "Task '{name}': can't use cache because no Cached states were found".format(
                        name=prefect.context.get("task_full_name", self.task.name)
                    )
                )
            else:
                self.logger.debug(
                    "Task '{name}': {num} candidate cached states were found".format(
                        name=prefect.context.get("task_full_name", self.task.name),
                        num=len(cached_states),
                    )
                )

            for candidate_state in cached_states:
                assert isinstance(candidate_state, Cached)  # mypy assert
                candidate_state.load_cached_results(inputs)
                sanitized_inputs = {key: res.value for key, res in inputs.items()}
                if self.task.cache_validator(
                    candidate_state, sanitized_inputs, prefect.context.get("parameters")
                ):
                    return candidate_state.load_result(self.result)

                self.logger.debug(
                    "Task '{name}': can't use cache because no candidate Cached states "
                    "were valid".format(
                        name=prefect.context.get("task_full_name", self.task.name)
                    )
                )

        return state

    def load_results(
        self, state: State, upstream_states: Dict[Edge, State]
    ) -> Tuple[State, Dict[Edge, State]]:
        """
        Given the task's current state and upstream states, populates all relevant result objects for this task run.

        Args:
            - state (State): the task's current state.
            - upstream_states (Dict[Edge, State]): the upstream state_handlers

        Returns:
            - Tuple[State, dict]: a tuple of (state, upstream_states)

        """
        upstream_results = {}

        try:
            for edge, upstream_state in upstream_states.items():
                upstream_states[edge] = upstream_state.load_result(
                    edge.upstream_task.result or self.default_result
                )
                if edge.key is not None:
                    upstream_results[edge.key] = (
                        edge.upstream_task.result or self.default_result
                    )

            state.load_cached_results(upstream_results)
            return state, upstream_states
        except Exception as exc:
            new_state = Failed(
                message=f"Failed to retrieve task results: {exc}", result=exc
            )
            final_state = self.handle_state_change(old_state=state, new_state=new_state)
            raise ENDRUN(final_state)

    def get_task_inputs(
        self, state: State, upstream_states: Dict[Edge, State]
    ) -> Dict[str, Result]:
        """
        Given the task's current state and upstream states, generates the inputs for this task.
        Upstream state result values are used. If the current state has `cached_inputs`, they
        will override any upstream values.

        Args:
            - state (State): the task's current state.
            - upstream_states (Dict[Edge, State]): the upstream state_handlers

        Returns:
            - Dict[str, Result]: the task inputs

        """
        task_inputs = super().get_task_inputs(state, upstream_states)

        try:
            ## for mapped tasks, we need to take extra steps to store the cached_inputs;
            ## this is because in the event of a retry we don't want to have to load the
            ## entire upstream array that is being mapped over, instead we need store the
            ## individual pieces of data separately for more efficient retries
            map_index = prefect.context.get("map_index")
            if map_index not in [-1, None]:
                for edge, upstream_state in upstream_states.items():
                    if (
                        edge.key
                        and edge.mapped
                        and edge.upstream_task.checkpoint is not False
                    ):
                        try:
                            task_inputs[edge.key] = task_inputs[edge.key].write(  # type: ignore
                                task_inputs[edge.key].value,
                                filename=f"{edge.key}-{map_index}",
                                **prefect.context,
                            )
                        except NotImplementedError:
                            pass
        except Exception as exc:
            new_state = Failed(
                message=f"Failed to save inputs for mapped task: {exc}", result=exc
            )
            final_state = self.handle_state_change(old_state=state, new_state=new_state)
            raise ENDRUN(final_state)

        return task_inputs

    @tail_recursive
    def run(
        self,
        state: State = None,
        upstream_states: Dict[Edge, State] = None,
        context: Dict[str, Any] = None,
        executor: "prefect.engine.executors.Executor" = None,
    ) -> State:
        """
        The main endpoint for TaskRunners.  Calling this method will conditionally execute
        `self.task.run` with any provided inputs, assuming the upstream dependencies are in a
        state which allow this Task to run.  Additionally, this method will wait and perform Task retries
        which are scheduled for <= 1 minute in the future.

        Args:
            - state (State, optional): initial `State` to begin task run from;
                defaults to `Pending()`
            - upstream_states (Dict[Edge, State]): a dictionary
                representing the states of any tasks upstream of this one. The keys of the
                dictionary should correspond to the edges leading to the task.
            - context (dict, optional): prefect Context to use for execution
            - executor (Executor, optional): executor to use when performing
                computation; defaults to the executor specified in your prefect configuration

        Returns:
            - `State` object representing the final post-run state of the Task
        """
        context = context or {}
        end_state = super().run(
            state=state,
            upstream_states=upstream_states,
            context=context,
            executor=executor,
        )
        while (end_state.is_retrying() or end_state.is_queued()) and (
            end_state.start_time <= pendulum.now("utc").add(minutes=10)  # type: ignore
        ):
            assert isinstance(end_state, (Retrying, Queued))
            naptime = max(
                (end_state.start_time - pendulum.now("utc")).total_seconds(), 0
            )
            time.sleep(naptime)

            # currently required as context has reset to its original state
            task_run_info = self.client.get_task_run_info(
                flow_run_id=context.get("flow_run_id", ""),
                task_id=context.get("task_id", ""),
                map_index=context.get("map_index"),
            )
            context.update(task_run_version=task_run_info.version)  # type: ignore

            end_state = super().run(
                state=end_state,
                upstream_states=upstream_states,
                context=context,
                executor=executor,
            )
        return end_state
