import logging
import os
import subprocess
import time
from tempfile import NamedTemporaryFile
from typing import Dict, List, Optional, OrderedDict, Sequence, Set, TypeVar

from tango.executor import ExecutionMetadata, Executor, ExecutorOutput
from tango.step import Step
from tango.step_graph import StepGraph
from tango.step_info import StepState
from tango.workspace import Workspace

logger = logging.getLogger(__name__)

T = TypeVar("T")


@Executor.register("multicore")
class MulticoreExecutor(Executor):
    """
    A ``MulticoreExecutor`` runs the steps in parallel and caches their results.
    """

    def __init__(
        self,
        workspace: Workspace,
        include_package: Optional[Sequence[str]] = None,
        parallelism: Optional[int] = 1,
        num_tries_to_sync_states: int = 3,
        wait_seconds_to_sync_states: int = 3,
    ) -> None:
        super().__init__(workspace, include_package=include_package, parallelism=parallelism or 1)
        assert self.parallelism is not None
        if self.parallelism < 0:
            self.parallelism = min(32, os.cpu_count() or 1)

        # Perhaps there's a better way to do this without these being passed as args.
        self._num_tries_to_sync_states = num_tries_to_sync_states
        self._wait_seconds_to_sync_states = wait_seconds_to_sync_states

    def execute_step_graph(
        self, step_graph: StepGraph, run_name: Optional[str] = None
    ) -> ExecutorOutput:
        """
        Execute a :class:`tango.step_graph.StepGraph`. This attempts to execute steps in parallel.
        If a step fails, its dependent steps are not run, but unrelated steps are still executed.
        Step failures will be logged, but no exceptions will be raised.
        """

        _running: OrderedDict[str, subprocess.Popen] = OrderedDict({})
        _successful: Dict[str, ExecutionMetadata] = {}
        _failed: Dict[str, ExecutionMetadata] = {}
        _queued_steps: List[str] = []

        uncacheable_leaf_steps = step_graph.uncacheable_leaf_steps()

        def _sync_step_states() -> Dict[str, StepState]:
            """
            Update the StepState info.
            Although, this is not really elegant. The issue is as follows: The main multicore executor process
            queues a step, and starts step execution in a different process. If we try to read the StepState
            before that process has had time to update the StepState, the Workspace will throw the out of sync
            error (IOError: process should be running but it's considered incomplete...).

            Hence, we try to read a few times, so that the child process has time to update the step's state.
            """

            attempts = 0
            while attempts < self._num_tries_to_sync_states:
                attempts += 1
                try:
                    step_states = {step.name: self._get_state(step) for step in step_graph.values()}
                    break
                except IOError:
                    if attempts == self._num_tries_to_sync_states:
                        raise
                    step_states = {}
                    time.sleep(self._wait_seconds_to_sync_states)
            return step_states

        def _has_incomplete_steps(step_states: Dict[str, StepState]) -> bool:
            """
            Are there any steps in the graph that are currently:
            1) running, or
            2) queued, or
            3) incomplete (with no failed dependencies).

            If there are any failed dependencies for a step, it will never manage to run.
            """

            def _failed_dependencies(step: Step) -> bool:
                for dependency in step.recursive_dependencies:
                    if (
                        step_states[dependency.name] == StepState.FAILED
                        or dependency.name in _failed
                    ):
                        return True
                return False

            uncacheable_leaf_step_names = {step.name for step in uncacheable_leaf_steps}
            for step_name, step_state in step_states.items():
                if (
                    step_name in _running
                    or step_name in _queued_steps
                    or (
                        # If the workspace already has a previous run, we disregard the failure.
                        step_state in [StepState.INCOMPLETE, StepState.FAILED]
                        and not _failed_dependencies(step_graph[step_name])
                        # We check for failures in this run.
                        and step_name not in _failed
                    )
                    or (
                        # Uncacheable leaf steps need to run, but their StepState will always be UNCACHEABLE.
                        step_name in uncacheable_leaf_step_names
                        and step_name not in _successful
                        and step_name not in _failed
                        and not _failed_dependencies(step_graph[step_name])
                    )
                ):
                    return True
            return False

        def _update_running_steps(step_states: Dict[str, StepState]) -> List[str]:
            """
            Check the running processes for status. We use poll_status to check if the process ended,
            but the StepState for checking completion/failure status, because after the process ends,
            the lock release etc. sometimes takes a beat longer.
            """
            done = []
            errors = []
            for step_name, process in _running.items():
                poll_status = process.poll()
                if poll_status is not None:
                    # The step may have finished since we synced step states.
                    if step_states[step_name] == StepState.RUNNING:
                        step_states[step_name] = self._get_state(step_graph[step_name])

                    if step_states[step_name] == StepState.UNCACHEABLE:
                        if poll_status == 0:
                            done.append(step_name)
                        else:
                            errors.append(step_name)
                    elif step_states[step_name] == StepState.COMPLETED:
                        done.append(step_name)
                    elif (
                        step_states[step_name] == StepState.FAILED
                        or step_states[step_name] == StepState.INCOMPLETE
                    ):
                        # TODO: look into why the step status changes from running back to incomplete sometimes.
                        # Possibly it's due to the workspace being aggressive in marking it as incomplete when
                        # it thinks that the process is not running.
                        errors.append(step_name)
                    else:
                        raise RuntimeError(
                            f"Step '{step_name}' has the state {step_states[step_name]}, "
                            "but the corresponding process has ended!"
                        )

            for step_name in done + errors:
                _running.pop(step_name)

            for step_name in done:
                step = step_graph[step_name]
                _successful[step_name] = ExecutionMetadata(
                    result_location=None
                    if not step.cache_results
                    else self.workspace.step_info(step).result_location
                )

            for step_name in errors:
                _failed[step_name] = ExecutionMetadata()

            return errors

        def _get_steps_to_run(step_states: Dict[str, StepState]) -> Set[str]:
            """
            Returns the steps that can be queued to run immediately.
            Criteria:
                1) All dependencies are available.
                2) Step is not already running or queued.
                3) Step has not run in the past and failed.
                4) Step's state is INCOMPLETE (or FAILED from a previous run), or
                   step's state is UNCACHEABLE and it is a leaf step.

            (We only run uncacheable steps if they are needed for another step downstream,
            as part of the downstream step).
            """

            def _are_dependencies_available(step: Step) -> bool:
                for dependency in step.dependencies:
                    if step_states[dependency.name] not in [
                        StepState.COMPLETED,
                        StepState.UNCACHEABLE,
                    ]:
                        return False
                return True

            to_run: Set[str] = set()
            for step in step_graph.values():
                if (
                    _are_dependencies_available(step)
                    and step.name not in _running  # Not already running.
                    and step.name not in _queued_steps  # Not queued to run.
                    and step.name not in _failed  # Not already failed.
                    # See comment in _has_incomplete_steps
                    and (
                        step_states[step.name] in [StepState.INCOMPLETE, StepState.FAILED]
                        or (
                            step_states[step.name] == StepState.UNCACHEABLE
                            and step in uncacheable_leaf_steps
                            and step.name not in _successful
                        )
                    )
                ):
                    to_run.add(step.name)
            return to_run

        def _queue_step(step_name: str) -> None:
            _queued_steps.append(step_name)
            logger.debug(f"Step {step_name} added to the queue for execution.")

        def _try_to_execute_next_step(config_path: str, run_name: Optional[str] = None) -> None:
            """
            If there are queued steps, try to start processes for them (limited by `parallelism`).
            """
            if len(_queued_steps) == 0:
                logger.debug("No steps in queue!")
                return
            if len(_running) < (self.parallelism or 1):
                step_name = _queued_steps.pop(0)
                command: List[str] = [
                    "tango",
                    "--called-by-executor",
                    "run",
                    config_path,
                    "-s",
                    step_name,
                    "-w",
                    self.workspace.url,
                    "--no-server",
                ]
                if self.include_package is not None:
                    for package in self.include_package:
                        command += ["-i", package]
                if run_name is not None:
                    command += ["-n", run_name]
                process = subprocess.Popen(command, shell=False)
                _running[step_name] = process
            else:
                logger.debug(
                    f"{self.parallelism or 1} steps are already running. Will attempt to execute later."
                )

        # Creates a temporary file in which to store the config. This is passed as a command line
        # argument to child step processes.
        with NamedTemporaryFile(prefix="step-graph-to-file-run", suffix=".jsonnet") as file_ref:
            step_graph.to_file(file_ref.name, include_unique_id=True)
            assert os.path.exists(file_ref.name)

            step_states = _sync_step_states()

            while _has_incomplete_steps(step_states):
                # Cleanup previously running steps.
                _update_running_steps(step_states)

                # Get steps that are ready to run.
                to_run = _get_steps_to_run(step_states)
                if to_run:
                    logger.debug(f"Steps ready to run: {to_run}")

                for step_name in to_run:
                    _queue_step(step_name)

                # Begin processes for any queued steps (if not enough processes are already running).
                while len(_queued_steps) > 0 and len(_running) < (self.parallelism or 1):
                    _try_to_execute_next_step(config_path=file_ref.name, run_name=run_name)

                # Re-sync the StepState info.
                step_states = _sync_step_states()

        assert not _running and not _queued_steps
        _not_run: Dict[str, ExecutionMetadata] = {}
        for step_name, step in step_graph.items():
            if step_name in _successful or step_name in _failed:
                # tried to execute directly
                continue
            elif not step.cache_results and step not in uncacheable_leaf_steps:
                # uncacheable interior step; didn't execute directly.
                continue
            elif (
                step.cache_results
                and step_name in step_states
                and step_states[step_name] == StepState.COMPLETED
            ):
                # step result was found in cache.
                # NOTE: since neither `Step.result()` nor `Step.ensure_result()` will have been
                # called, we invoke the CLI logger here to let users know that we didn't run this
                # step because we found it in the cache.
                step.log_cache_hit()
                _successful[step_name] = ExecutionMetadata(
                    result_location=self.workspace.step_info(step_graph[step_name]).result_location
                )
            else:
                # step wasn't executed because parents failed, or
                # step is uncacheable leaf step, so we do care about what happened to it.
                _not_run[step_name] = ExecutionMetadata()

        return ExecutorOutput(successful=_successful, failed=_failed, not_run=_not_run)

    def _get_state(self, step: Step) -> StepState:
        """
        Returns the StepState as determined by the workspace.
        """
        return self.workspace.step_info(step).state
