import logging
from collections.abc import Callable, Hashable, Sequence
from itertools import chain
from types import ModuleType
from typing import Literal, TypeVar, overload

from ordeq.framework.graph import NodeGraph
from ordeq.framework.hook import NodeHook, RunHook
from ordeq.framework.io import Input, Output, _InputCache
from ordeq.framework.nodes import NODE_REGISTRY, Node, get_node

logger = logging.getLogger(__name__)

T = TypeVar("T")

DataStoreType = dict[Input[T] | Output[T], T]
# The save mode determines which outputs are saved. When set to:
# - 'all', all outputs are saved, including those of intermediate nodes.
# - 'sinks', only outputs of sink nodes are saved, i.e. those w/o successors.
# - 'none', to dry-run and save no outputs
# Future extension:
# - 'last', which saves the output of the last node for which no error
# occurred. This can be useful for debugging.
SaveMode = Literal["all", "sinks", "none"]


def _save_outputs(
    node: Node, values: Sequence[T], save: bool = True
) -> dict[Input[T] | Output[T], T]:
    computed: dict[Input[T] | Output[T], T] = {}
    for output_dataset, data in zip(node.outputs, values, strict=False):
        computed[output_dataset] = data

        # TODO: this can be handled in the `save_wrapper`
        if save:
            output_dataset.save(data)

    return computed


def _run_node(
    node: Node, *, hooks: Sequence[NodeHook] = (), save: bool = True
) -> DataStoreType:
    node.validate()

    for node_hook in hooks:
        node_hook.before_node_run(node)

    args = [input_dataset.load() for input_dataset in node.inputs]

    # persisting loaded data
    for node_input, data in zip(node.inputs, args, strict=True):
        if isinstance(node_input, _InputCache):
            node_input.persist(data)

    logger.info("Running node %s", node)

    try:
        values = node.func(*args)
    except Exception as exc:
        for node_hook in hooks:
            node_hook.on_node_call_error(node, exc)
        raise exc

    if len(node.outputs) == 0:
        values = ()
    elif len(node.outputs) == 1:
        values = (values,)
    else:
        values = tuple(values)

    computed = _save_outputs(node, values, save=save)

    # persisting computed data only if outputs are loaded again later
    for node_output in node.outputs:
        if isinstance(node_output, _InputCache):
            node_output.persist(computed[node_output])  # ty: ignore[call-non-callable]

    for node_hook in hooks:
        node_hook.after_node_run(node)

    return computed


def _run_graph(
    graph: NodeGraph,
    *,
    hooks: Sequence[NodeHook] = (),
    save: SaveMode = "all",
    io: dict[Input[T] | Output[T], Input[T] | Output[T]] | None = None,
) -> DataStoreType:
    """Runs nodes in a graph topologically, ensuring IOs are loaded only once.

    Args:
        graph: node graph to run
        hooks: hooks to apply
        hooks: hooks to apply
        save: 'all' | 'sinks' | 'none'.
            If 'sinks', only saves the outputs of sink nodes in the graph.
        io: mapping of IO objects to their replacements

    Returns:
        a dict mapping each IO to the computed data

    """

    patched_nodes: dict[Node, Node] = {}
    for node in graph.nodes:
        patched_nodes[node] = _patch_io(node, io or {})

    data_store: dict = {}  # For each IO, the loaded data

    # TODO: Create _Patch wrapper for IO?
    for node in graph.topological_ordering:
        if (save == "sinks" and node in graph.sink_nodes) or save == "all":
            save_node = True
        else:
            save_node = False

        computed = _run_node(patched_nodes[node], hooks=hooks, save=save_node)
        data_store.update(computed)

    reverse_io = {v: k for k, v in (io or {}).items()}
    patched_data_store = {}
    for k, v in data_store.items():
        patched_data_store[reverse_io.get(k, k)] = v

    # unpersist IO objects
    for gnode in graph.nodes:
        io_objs = chain(gnode.inputs, gnode.outputs)
        for io_obj in io_objs:
            if isinstance(io_obj, _InputCache):
                io_obj.unpersist()

    return patched_data_store


def _patch_io(
    node: Node, io: dict[Input[T] | Output[T], Input[T] | Output[T]]
) -> Node:
    """Patches the inputs and outputs of a node with the provided IO mapping.

    Args:
        node: the original node
        io: mapping of Input/Output objects to their replacements

    Returns:
        the node with patched inputs and outputs

    """

    return node._replace(
        inputs=[io.get(ip, ip) for ip in node.inputs],  # type: ignore[misc]
        outputs=[io.get(op, op) for op in node.outputs],  # type: ignore[misc]
    )


def _gather_nodes_from_module(module: ModuleType) -> list[Node]:
    """Gathers all nodes defined in a module.

    Args:
        module: the module to gather nodes from

    Returns:
        a list of nodes defined in the module

    """

    nodes = []
    for attr in dir(module):
        obj = getattr(module, attr)
        if isinstance(obj, Hashable) and obj in NODE_REGISTRY:
            nodes.append(NODE_REGISTRY.get(obj))
    return nodes


@overload
def run(
    *runnables: ModuleType,
    hooks: Sequence[NodeHook | RunHook] = (),
    save: SaveMode = "all",
    verbose: bool = False,
    io: dict[Input[T] | Output[T], Input[T] | Output[T]] | None = None,
) -> DataStoreType: ...


@overload
def run(
    *runnables: Callable,
    hooks: Sequence[NodeHook | RunHook] = (),
    save: SaveMode = "all",
    verbose: bool = False,
    io: dict[Input[T] | Output[T], Input[T] | Output[T]] | None = None,
) -> DataStoreType: ...


def run(
    *runnables: ModuleType | Callable,
    hooks: Sequence[NodeHook | RunHook] = (),
    save: SaveMode = "all",
    verbose: bool = False,
    io: dict[Input[T] | Output[T], Input[T] | Output[T]] | None = None,
) -> DataStoreType:
    """Runs nodes in topological order.

    Args:
        runnables: the nodes to run, or a modules containing nodes
        hooks: hooks to apply
        save: 'all' | 'sinks'. If 'sinks', only saves the sink outputs
        verbose: whether to print the node graph
        io: mapping of IO objects to their replacements

    Returns:
        a dict mapping each IO to the computed data

    Raises:
        TypeError: if runnables are not all modules or all callables
    """

    if all(isinstance(r, ModuleType) for r in runnables):
        modules: tuple[ModuleType, ...] = runnables  # type: ignore[assignment]
        nodes_ = []
        for module in modules:
            nodes_.extend(_gather_nodes_from_module(module))
    elif all(callable(r) for r in runnables):
        funcs: tuple[Callable, ...] = runnables  # type: ignore[assignment]
        nodes_ = [get_node(func) for func in funcs]
    else:
        raise TypeError("All runnables must be either modules or nodes.")

    graph = NodeGraph.from_nodes(nodes_)

    if verbose:
        print(graph)

    node_hooks: list[NodeHook] = [h for h in hooks if isinstance(h, NodeHook)]
    run_hooks: list[RunHook] = [h for h in hooks if isinstance(h, RunHook)]

    for run_hook in run_hooks:
        run_hook.before_run(graph)

    result = _run_graph(graph, hooks=node_hooks, save=save, io=io)

    for run_hook in run_hooks:
        run_hook.after_run(graph, result)

    return result
