from typing import (
    Any,
    Callable,
    Dict,
    List,
    Optional,
    Tuple,
    TYPE_CHECKING,
    Union,
)
import collections
import dataclasses
import enum
import itertools

import numpy as np

from realhf.api.core.config import ModelBackend, ModelFamily, ModelName
import realhf.base.logging as logging
import realhf.base.namedarray as namedarray

logger = logging.getLogger("DataFlowGraph", "benchmark")


@dataclasses.dataclass
class OffloadHook:
    pass


@dataclasses.dataclass
class SyncParamHook:
    source: Optional[ModelName] = None
    target: Optional[ModelName] = None
    interval: int = 1


RPCHook = Union[OffloadHook, SyncParamHook]


@dataclasses.dataclass
class ModelInterface:
    type_: str
    args: Dict[str, Any] = dataclasses.field(default_factory=dict)

    def __eq__(self, other: "ModelInterface"):
        return self.type_ == other.type_ and self.args == other.args


class ModelInterfaceType(enum.Enum):
    GENERATE = "generate"
    TRAIN_STEP = "train_step"
    EVALUATE = "evaluate"
    INFERENCE = "inference"


@dataclasses.dataclass
class MFCDef:
    """This class defines a model function call (MFC) node in the dataflow graph.

    Dependencies are automatically resolved by the `build_graph` function
    according to the input/output data keys and the model name.

    :param model_name: The model identifier to be used by the node.
    :type model_name: ModelName
    :param model_type: The specification of the LLM, e.g., LLaMA-7B.
        Used by the profiler and search engine to produce an optimal execution plan.
    :type model_type: ModelFamily
    :param model_path: The path to the model file. Used to get the config.
    :type model_path: str
    :param interface_type: The interface type of this node, e.g., "generate".
    :type interface_type: ModelInterfaceType
    :param interface_impl: The actual interface implementation
        when running the this node.
    :type interface_impl: ModelInterface
    :param min_n_seqs: The minimum number of sequences to be processed in a batch.
    :type min_n_seqs: int
    :param max_n_seqs: The maximum number of sequences to be processed in a batch.
    :type max_n_seqs: int
    :param input_data: The data keys required by this node.
    :type input_data: List[str]
    :param input_key_remap: Remap input keys to let the interface implementation
        recognize them. Keys are ``input_data`` and values are the identifiers known to the
        interface.
    :type input_key_remap: Dict[str, str]
    :param output_data: The data keys generated by this node.
    :type output_data: List[str]
    :param output_key_remap: Remap output keys to let MFC recognize them.
        Keys are identifiers known to the interface and values are ``output_data``.
    :type output_key_remap: Dict[str, str]
    :param log_return_value: Whether to log the return value of the interface implementation.
    :type log_return_value: bool
    :param min_n_seqs_per_dp: The minimum number of sequences per data parallelism.
    :type min_n_seqs_per_dp: int
    :param balanced_dp: Whether to balance the data parallelism.
    :type balanced_dp: bool
    :param max_concurrent_calls: The maximum number of concurrent calls.
    :type max_concurrent_calls: int
    :param pre_hooks: The hooks to be executed before the interface implementation.
        Typically parameter reallocation. Should not be manually changed in
        using quickstart experiments.
    :type pre_hooks: List[RPCHook]
    :param post_hooks: The hooks to be executed after the interface implementation.
        Typically offload or parameter reallocation. Should not be manually changed in
        using quickstart experiments.
    :type post_hooks: List[RPCHook]
    """

    model_name: ModelName
    model_type: ModelFamily
    model_path: str
    interface_type: ModelInterfaceType
    interface_impl: ModelInterface

    # batch sizes
    min_n_seqs: int
    max_n_seqs: int

    input_data: List[str] = dataclasses.field(default_factory=lambda: [])
    input_key_remap: Dict[str, str] = dataclasses.field(
        default_factory=lambda: {}
    )
    output_data: List[str] = dataclasses.field(default_factory=lambda: [])
    output_key_remap: Dict[str, str] = dataclasses.field(
        default_factory=lambda: {}
    )
    log_return_value: bool = False

    min_n_seqs_per_dp: int = 1
    balanced_dp: bool = False

    max_concurrent_calls: int = 1

    # hooks
    pre_hooks: List[RPCHook] = dataclasses.field(default_factory=lambda: [])
    post_hooks: List[RPCHook] = dataclasses.field(default_factory=lambda: [])

    # The followings will be automatically filled.
    max_min_flow_seqs: int = 1

    parents: List[str] = dataclasses.field(default_factory=lambda: [])
    children: List[str] = dataclasses.field(default_factory=lambda: [])

    parent_rpcs: List["MFCDef"] = dataclasses.field(default_factory=lambda: [])
    children_rpcs: List["MFCDef"] = dataclasses.field(
        default_factory=lambda: []
    )

    # data key -> model name
    data_producers: Dict[str, ModelName] = dataclasses.field(
        default_factory=lambda: {}
    )
    # data key -> rpc names
    data2required_rpc_names: Dict[str, List[str]] = dataclasses.field(
        default_factory=lambda: {}
    )

    def __post_init__(self):
        if isinstance(self.model_name, str):
            self.model_name = ModelName(self.model_name, 0)
        assert isinstance(self.model_name, ModelName)
        if self.min_n_seqs > self.max_n_seqs:
            raise RuntimeError("Invalid min/max n_seqs.")
        if self.is_src and self.max_n_seqs > 1e4:
            raise RuntimeError(
                "The maximum batch size of the source node in the dataflow graph is too large. "
                f"The maximum number of sequences is {self.max_n_seqs} > budget {int(1e4)}. "
                "Please set a smaller value."
            )
        if "@" in self.model_name.role or "@" in self.interface_type.value:
            raise ValueError(
                f"Invalid model name or interface type: {self.model_name}, {self.interface_type}."
            )

    def __repr__(self):
        return f"MFCDef({self.model_name}, {self.interface_type})"

    @property
    def name(self):
        return f"{self.model_name}@{self.interface_type.value}"

    @property
    def is_src(self):
        return len(self.parents) == 0

    @property
    def is_dst(self):
        return len(self.children) == 0

    @property
    def is_dst_of_model_role(self):

        def _has_children_of_model_name(rpc: "MFCDef", model_name: ModelName):
            if rpc.is_dst:
                return False
            return any(
                [
                    r.model_name.role == model_name.role
                    or _has_children_of_model_name(r, model_name)
                    for r in rpc.children_rpcs
                ]
            )

        return not _has_children_of_model_name(self, self.model_name)

    def remap_input_keys(self, input_batch: Dict) -> namedarray.NamedArray:
        data = {}
        for k in self.input_data:
            if k not in self.input_key_remap:
                data[k] = input_batch[k]
            else:
                data[self.input_key_remap[k]] = input_batch[k]
        return namedarray.from_dict(data)

    def remap_output_keys(self, output_batch: Dict) -> namedarray.NamedArray:
        res_data = {}
        for k, v in output_batch.items():
            if k not in self.output_data:
                continue
            if k in self.output_key_remap:
                res_data[self.output_key_remap[k]] = v
            else:
                res_data[k] = v
        return namedarray.from_dict(res_data)


def build_graph(
    rpcs: List[MFCDef], verbose: bool = False
) -> Tuple[List[MFCDef], List[List[Tuple[str]]]]:
    # Resolve dependencies between model interfaces.
    children: List[List[str]] = [[] for _ in rpcs]
    parents: List[List[str]] = [[] for _ in rpcs]
    parent_rpcs: List[List[MFCDef]] = [[] for _ in rpcs]
    children_rpcs: List[List[MFCDef]] = [[] for _ in rpcs]
    edges: List[List[Tuple[str]]] = [[() for _ in rpcs] for _ in rpcs]

    required_data_entries: List[Tuple[str]] = [() for _ in rpcs]
    generated_data_entries: List[Tuple[str]] = [() for _ in rpcs]
    for i, rpc in enumerate(rpcs):
        required_data_entries[i] = (*required_data_entries[i], *rpc.input_data)
        generated_data_entries[i] = (
            *generated_data_entries[i],
            *[
                k if k not in rpc.output_key_remap else rpc.output_key_remap[k]
                for k in rpc.output_data
            ],
        )
    data_producers = {}
    for rpc, gd in zip(rpcs, generated_data_entries):
        for k in gd:
            data_producers[k] = rpc.model_name
    data2required_rpc_names = collections.defaultdict(list)
    for rpc, data_keys in zip(rpcs, required_data_entries):
        for k in data_keys:
            data2required_rpc_names[k].append(rpc.name)

    for i, rpc in enumerate(rpcs):
        for k in required_data_entries[i]:
            for j, parent_rpc in enumerate(rpcs):
                if parent_rpc.name == rpc.name:
                    continue
                if k in generated_data_entries[j]:
                    if parent_rpc.name not in parents[i]:
                        parents[i].append(parent_rpc.name)
                        parent_rpcs[i].append(parent_rpc)
                    if rpc.name not in children[j]:
                        children[j].append(rpc.name)
                        children_rpcs[j].append(rpc)
                    edges[i][j] = (*edges[i][j], k)
                    if verbose:
                        logger.info(
                            f"Dependency added: {rpc.name}"
                            f" <- {parent_rpc.name} "
                            f"because of data entry `{k}`."
                        )
    if verbose:
        for i, rpc in enumerate(rpcs):
            logger.info(
                f"Dependency: {rpc.name} <- { {x.name: deps for x, deps in zip(rpcs, edges[i]) if deps} }."
            )
    for rpc, p, c, pr, cr in zip(
        rpcs, parents, children, parent_rpcs, children_rpcs
    ):
        rpc.parents = p
        rpc.children = c
        rpc.parent_rpcs = pr
        rpc.children_rpcs = cr

    for rpc in rpcs:
        rpc.max_min_flow_seqs = max(
            [
                r.min_n_seqs
                for r in rpcs
                if r.model_name.role == rpc.model_name.role
            ]
        )
        rpc.data_producers = data_producers
        rpc.data2required_rpc_names = data2required_rpc_names

    # sanity check of hooks
    for rpc in rpcs:
        for h in itertools.chain(rpc.pre_hooks, rpc.post_hooks):
            assert isinstance(h, RPCHook), type(h)
            if isinstance(h, SyncParamHook):
                assert any(h.target == r.model_name for r in rpcs) or (
                    h.source == r.model_name for r in rpcs
                )
        for h in rpc.pre_hooks:
            if isinstance(h, OffloadHook):
                raise ValueError("Offload can only be post hooks!")
    return rpcs, edges
