import threading
from typing import Dict, List, Set

from .run import RunTask, ModelRunner as run_model_runner
from .snapshot import SnapshotRunner as snapshot_model_runner
from .seed import SeedRunner as seed_runner
from .test import TestRunner as test_runner

from dbt.artifacts.schemas.results import NodeStatus, RunStatus
from dbt.artifacts.schemas.run import RunResult
from dbt.cli.flags import Flags
from dbt.config.runtime import RuntimeConfig
from dbt.contracts.graph.manifest import Manifest
from dbt.graph import ResourceTypeSelector, GraphQueue, Graph
from dbt.node_types import NodeType
from dbt.task.test import TestSelector
from dbt.task.base import BaseRunner, resource_types_from_args
from dbt_common.events.functions import fire_event
from dbt.events.types import LogNodeNoOpResult
from dbt.exceptions import DbtInternalError


class SavedQueryRunner(BaseRunner):
    # Stub. No-op Runner for Saved Queries, which require MetricFlow for execution.
    @property
    def description(self):
        return f"saved query {self.node.name}"

    def before_execute(self):
        pass

    def compile(self, manifest):
        return self.node

    def after_execute(self, result):
        fire_event(
            LogNodeNoOpResult(
                description=self.description,
                index=self.node_index,
                total=self.num_nodes,
                node_info=self.node.node_info,
            )
        )

    def execute(self, compiled_node, manifest):
        # no-op
        return RunResult(
            node=compiled_node,
            status=RunStatus.Success,
            timing=[],
            thread_id=threading.current_thread().name,
            execution_time=0,
            message="NO-OP",
            adapter_response={},
            failures=0,
            agate_table=None,
        )


class BuildTask(RunTask):
    """The Build task processes all assets of a given process and attempts to
    'build' them in an opinionated fashion.  Every resource type outlined in
    RUNNER_MAP will be processed by the mapped runner class.

    I.E. a resource of type Model is handled by the ModelRunner which is
    imported as run_model_runner."""

    MARK_DEPENDENT_ERRORS_STATUSES = [NodeStatus.Error, NodeStatus.Fail, NodeStatus.Skipped]

    RUNNER_MAP = {
        NodeType.Model: run_model_runner,
        NodeType.Snapshot: snapshot_model_runner,
        NodeType.Seed: seed_runner,
        NodeType.Test: test_runner,
        NodeType.Unit: test_runner,
        NodeType.SavedQuery: SavedQueryRunner,
    }
    ALL_RESOURCE_VALUES = frozenset({x for x in RUNNER_MAP.keys()})

    def __init__(self, args: Flags, config: RuntimeConfig, manifest: Manifest) -> None:
        super().__init__(args, config, manifest)
        self.selected_unit_tests: Set = set()
        self.model_to_unit_test_map: Dict[str, List] = {}

    def resource_types(self, no_unit_tests=False):
        resource_types = resource_types_from_args(
            self.args, set(self.ALL_RESOURCE_VALUES), set(self.ALL_RESOURCE_VALUES)
        )

        # First we get selected_nodes including unit tests, then without,
        # and do a set difference.
        if no_unit_tests is True and NodeType.Unit in resource_types:
            resource_types.remove(NodeType.Unit)
        return list(resource_types)

    # overrides get_graph_queue in runnable.py
    def get_graph_queue(self) -> GraphQueue:
        # Following uses self.selection_arg and self.exclusion_arg
        spec = self.get_selection_spec()

        # selector including unit tests
        full_selector = self.get_node_selector(no_unit_tests=False)
        # selected node unique_ids with unit_tests
        full_selected_nodes = full_selector.get_selected(spec)

        # This selector removes the unit_tests from the selector
        selector_wo_unit_tests = self.get_node_selector(no_unit_tests=True)
        # selected node unique_ids without unit_tests
        selected_nodes_wo_unit_tests = selector_wo_unit_tests.get_selected(spec)

        # Get the difference in the sets of nodes with and without unit tests and
        # save it
        selected_unit_tests = full_selected_nodes - selected_nodes_wo_unit_tests
        self.selected_unit_tests = selected_unit_tests
        self.build_model_to_unit_test_map(selected_unit_tests)

        # get_graph_queue in the selector will remove NodeTypes not specified
        # in the node_selector (filter_selection).
        return selector_wo_unit_tests.get_graph_queue(spec)

    # overrides handle_job_queue in runnable.py
    def handle_job_queue(self, pool, callback):
        if self.run_count == 0:
            self.num_nodes = self.num_nodes + len(self.selected_unit_tests)
        node = self.job_queue.get()
        if (
            node.resource_type == NodeType.Model
            and self.model_to_unit_test_map
            and node.unique_id in self.model_to_unit_test_map
        ):
            self.handle_model_with_unit_tests_node(node, pool, callback)

        else:
            self.handle_job_queue_node(node, pool, callback)

    def handle_model_with_unit_tests_node(self, node, pool, callback):
        self._raise_set_error()
        args = [node]
        if self.config.args.single_threaded:
            callback(self.call_model_and_unit_tests_runner(*args))
        else:
            pool.apply_async(self.call_model_and_unit_tests_runner, args=args, callback=callback)

    def call_model_and_unit_tests_runner(self, node) -> RunResult:
        assert self.manifest
        for unit_test_unique_id in self.model_to_unit_test_map[node.unique_id]:
            unit_test_node = self.manifest.unit_tests[unit_test_unique_id]
            unit_test_runner = self.get_runner(unit_test_node)
            # If the model is marked skip, also skip the unit tests
            if node.unique_id in self._skipped_children:
                # cause is only for ephemeral nodes
                unit_test_runner.do_skip(cause=None)
            result = self.call_runner(unit_test_runner)
            self._handle_result(result)
            if result.status in self.MARK_DEPENDENT_ERRORS_STATUSES:
                # The _skipped_children dictionary can contain a run_result for ephemeral nodes,
                # but that should never be the case here.
                self._skipped_children[node.unique_id] = None
        runner = self.get_runner(node)
        if runner.node.unique_id in self._skipped_children:
            cause = self._skipped_children.pop(runner.node.unique_id)
            runner.do_skip(cause=cause)
        return self.call_runner(runner)

    # handle non-model-plus-unit-tests nodes
    def handle_job_queue_node(self, node, pool, callback):
        self._raise_set_error()
        runner = self.get_runner(node)
        # we finally know what we're running! Make sure we haven't decided
        # to skip it due to upstream failures
        if runner.node.unique_id in self._skipped_children:
            cause = self._skipped_children.pop(runner.node.unique_id)
            runner.do_skip(cause=cause)
        args = [runner]
        if self.config.args.single_threaded:
            callback(self.call_runner(*args))
        else:
            pool.apply_async(self.call_runner, args=args, callback=callback)

    # Make a map of model unique_ids to selected unit test unique_ids,
    # for processing before the model.
    def build_model_to_unit_test_map(self, selected_unit_tests):
        dct = {}
        for unit_test_unique_id in selected_unit_tests:
            unit_test = self.manifest.unit_tests[unit_test_unique_id]
            model_unique_id = unit_test.depends_on.nodes[0]
            if model_unique_id not in dct:
                dct[model_unique_id] = []
            dct[model_unique_id].append(unit_test.unique_id)
        self.model_to_unit_test_map = dct

    # We return two different kinds of selectors, one with unit tests and one without
    def get_node_selector(self, no_unit_tests=False) -> ResourceTypeSelector:
        if self.manifest is None or self.graph is None:
            raise DbtInternalError("manifest and graph must be set to get node selection")

        resource_types = self.resource_types(no_unit_tests)

        if resource_types == [NodeType.Test]:
            return TestSelector(
                graph=self.graph,
                manifest=self.manifest,
                previous_state=self.previous_state,
            )
        return ResourceTypeSelector(
            graph=self.graph,
            manifest=self.manifest,
            previous_state=self.previous_state,
            resource_types=resource_types,
        )

    def get_runner_type(self, node):
        return self.RUNNER_MAP.get(node.resource_type)

    # Special build compile_manifest method to pass add_test_edges to the compiler
    def compile_manifest(self) -> None:
        if self.manifest is None:
            raise DbtInternalError("compile_manifest called before manifest was loaded")
        self.graph: Graph = self.compiler.compile(self.manifest, add_test_edges=True)
