# SPDX-License-Identifier: Apache-2.0
from collections import deque
from typing import Any, Callable, Dict, List, Optional, Tuple, Sequence, Union

import numpy as np

import torch
import torch.fx
import onnx
from onnx import numpy_helper

from .exceptions import UnsupportedDTypeError, UnsupportedOpError, ValueNotFoundError
from .op_registry import get_handler
from .utils.dtype import DTYPE_MAP
from .utils.external_data import resolve_external_data
from .utils.names import sanitize_name

# Import ops module to register all operators
from . import ops  # noqa: F401


def _collect_all_inputs(node: onnx.NodeProto) -> set:
    """Recursively collect all inputs from a node including subgraph inputs.

    For control flow nodes like If and Loop, this also collects inputs that
    are referenced by the subgraphs (then_branch, else_branch, body).

    Parameters
    ----------
    node : onnx.NodeProto
        The ONNX node to collect inputs from.

    Returns
    -------
    set
        Set of all input names referenced by this node and its subgraphs.
    """
    inputs = set(node.input)

    # Collect inputs from subgraphs (for If, Loop, etc.)
    for attr in node.attribute:
        if attr.type == onnx.AttributeProto.GRAPH:
            subgraph = attr.g
            # Collect subgraph's own initializers and inputs as local values
            local_values = set()
            for init in subgraph.initializer:
                local_values.add(init.name)
            for inp in subgraph.input:
                local_values.add(inp.name)

            # Recursively collect inputs from subgraph nodes
            for sub_node in subgraph.node:
                sub_inputs = _collect_all_inputs(sub_node)
                # Add outputs of this subgraph node to local values
                for out in sub_node.output:
                    if out:
                        local_values.add(out)
                # Inputs not satisfied locally are outer references
                for sub_inp in sub_inputs:
                    if sub_inp and sub_inp not in local_values:
                        inputs.add(sub_inp)

    return inputs


def _topological_sort(
    nodes: List[onnx.NodeProto],
    graph_inputs: set,
    initializers: set,
) -> List[onnx.NodeProto]:
    """Topologically sort ONNX graph nodes using Kahn's algorithm.

    Some ONNX models have nodes in non-topological order (e.g., Cast nodes
    at the end of the graph but their outputs used earlier). This function
    reorders nodes so dependencies are processed before their consumers.

    This function also considers inputs referenced by subgraphs (for nodes
    like If and Loop) to ensure proper ordering.

    Parameters
    ----------
    nodes : List[onnx.NodeProto]
        The list of ONNX nodes to sort.
    graph_inputs : set
        Set of graph input names.
    initializers : set
        Set of initializer names.

    Returns
    -------
    List[onnx.NodeProto]
        Topologically sorted list of nodes.
    """
    if not nodes:
        return []

    # Pre-compute all inputs for each node (including subgraph inputs)
    node_all_inputs: Dict[int, set] = {}
    for node in nodes:
        node_all_inputs[id(node)] = _collect_all_inputs(node)

    # Build output->node mapping (which node produces each output)
    output_to_node: Dict[str, onnx.NodeProto] = {}
    for node in nodes:
        for output in node.output:
            if output:  # Skip empty outputs
                output_to_node[output] = node

    # Available values: graph inputs + initializers
    available = graph_inputs | initializers

    # Compute in-degree for each node (number of unsatisfied dependencies)
    in_degree: Dict[int, int] = {}
    node_id: Dict[int, onnx.NodeProto] = {}
    for i, node in enumerate(nodes):
        node_id[id(node)] = node
        # Count inputs that are neither available nor empty
        deps = 0
        all_inputs = node_all_inputs[id(node)]
        for inp in all_inputs:
            if inp and inp not in available:
                deps += 1
        in_degree[id(node)] = deps

    # Initialize queue with nodes that have no dependencies
    queue = deque()
    for node in nodes:
        if in_degree[id(node)] == 0:
            queue.append(node)

    sorted_nodes: List[onnx.NodeProto] = []
    while queue:
        node = queue.popleft()
        sorted_nodes.append(node)

        # Mark this node's outputs as available
        for output in node.output:
            if output:
                available.add(output)

        # Reduce in-degree for nodes that depend on this node's outputs
        for candidate in nodes:
            if in_degree[id(candidate)] > 0:
                # Check if any of candidate's inputs are now satisfied
                all_inputs = node_all_inputs[id(candidate)]
                for inp in all_inputs:
                    if inp in node.output and inp:
                        in_degree[id(candidate)] -= 1
                        if in_degree[id(candidate)] == 0:
                            queue.append(candidate)

    # If we couldn't sort all nodes, there's a cycle or missing dependency
    # Fall back to original order
    if len(sorted_nodes) != len(nodes):
        return list(nodes)

    return sorted_nodes


class GraphBuilder:
    """Builds a PyTorch FX GraphModule from an ONNX model.

    This class handles the conversion of ONNX graph structure to PyTorch FX,
    including initializer loading, placeholder creation, node conversion,
    and output creation.

    Parameters
    ----------
    model : onnx.ModelProto
        The ONNX model to convert.

    Attributes
    ----------
    model : onnx.ModelProto
        The input ONNX model (with shape inference if successful).
    graph : torch.fx.Graph
        The FX graph being constructed.
    env : Dict[str, torch.fx.Node]
        Mapping from ONNX tensor names to FX nodes.
    opset_version : int
        The opset version for the default ONNX domain.
    """

    def __init__(
        self,
        model: onnx.ModelProto,
        *,
        base_dir: Optional[str] = None,
        memmap_external_data: bool = False,
    ) -> None:
        # Try shape inference but preserve original model if it fails
        # (shape_inference may drop graph contents for large models with external data)
        try:
            inferred_model = onnx.shape_inference.infer_shapes(model)
            # Check if shape inference preserved the model structure
            if len(inferred_model.graph.node) > 0:
                model = inferred_model
            # If nodes were lost, keep original model
        except Exception:
            pass
        self.model: onnx.ModelProto = model
        self.graph: torch.fx.Graph = torch.fx.Graph()
        self._base_dir = base_dir
        self._memmap_external_data = memmap_external_data
        self.value_info_map = self._create_value_info_map()
        self.initializer_map = self._create_initializer_map()
        self.input_names: List[str] = []
        self.env: Dict[str, torch.fx.Node] = {}
        self._constants: Dict[str, torch.Tensor] = {}
        self._submodules: Dict[str, torch.nn.Module] = {}
        self._opset_versions: Dict[str, int] = self._extract_opset_versions()

    def _extract_opset_versions(self) -> Dict[str, int]:
        """Extract opset versions for all domains from the model.

        Returns
        -------
        Dict[str, int]
            Dictionary mapping domain names to their opset versions.
            Empty string "" represents the default ONNX domain.
        """
        versions: Dict[str, int] = {}
        for opset in self.model.opset_import:
            domain = opset.domain if opset.domain else ""
            versions[domain] = opset.version
        return versions

    def _resolve_handler(
        self, node: onnx.NodeProto
    ) -> tuple[Callable[["GraphBuilder", onnx.NodeProto], Any], str, int]:
        """Resolve the handler for an ONNX node and return handler, domain, opset."""
        domain = node.domain if node.domain else ""
        opset = self.get_opset_version(domain)
        handler = get_handler(node.op_type, domain, opset)
        if handler is None:
            raise UnsupportedOpError(node.op_type, domain=domain, opset_version=opset)
        return handler, domain, opset

    def _tag_operator_node(
        self, node: onnx.NodeProto, fx_node: torch.fx.Node, domain: str
    ) -> None:
        """Attach ONNX metadata to an operator node."""
        if fx_node is not None and hasattr(fx_node, "meta"):
            self._set_onnx_metadata(
                fx_node,
                op_type=node.op_type,
                name=node.name,
                domain=domain,
            )

    def _register_outputs(
        self, node: onnx.NodeProto, fx_node: torch.fx.Node, domain: str
    ) -> None:
        """Register node outputs in the environment."""
        if len(node.output) == 1:
            self.env[node.output[0]] = fx_node
            return

        for i, output_name in enumerate(node.output):
            if output_name:  # Skip empty output names
                getitem_node = self.graph.call_function(
                    lambda x, idx=i: x[idx] if isinstance(x, (tuple, list)) else x,
                    args=(fx_node, i),
                )
                self._set_onnx_metadata(
                    getitem_node,
                    op_type=node.op_type,
                    name=node.name,
                    domain=domain,
                    output_index=i,
                )
                self.env[output_name] = getitem_node

    @property
    def opset_version(self) -> int:
        """Get the opset version for the default ONNX domain.

        Returns
        -------
        int
            The opset version number. Defaults to 1 if not specified.
        """
        return self._opset_versions.get("", 1)

    def get_opset_version(self, domain: str = "") -> int:
        """Get the opset version for a specific domain.

        Parameters
        ----------
        domain : str, optional
            The ONNX domain. Default is "" (standard ONNX domain).

        Returns
        -------
        int
            The opset version number. Defaults to 1 if not specified.
        """
        return self._opset_versions.get(domain, 1)

    def build(self) -> torch.fx.GraphModule:
        self._load_initializers()
        self._create_placeholders()
        self._convert_nodes()
        self._create_outputs()
        root_module = torch.nn.Module()
        # Register constants as buffers
        for name, tensor in self._constants.items():
            root_module.register_buffer(sanitize_name(name), tensor)
        # Register submodules
        for name, submod in self._submodules.items():
            root_module.add_module(name, submod)
        module = torch.fx.GraphModule(root_module, self.graph)
        if self._memmap_external_data:
            module._onnx2fx_inference_only = True
        module.graph.lint()
        return module

    @staticmethod
    def _set_onnx_metadata(
        node: torch.fx.Node,
        *,
        op_type: str,
        name: Optional[str] = None,
        domain: Optional[str] = None,
        shape: Optional[List[Optional[int]]] = None,
        dtype: Optional[torch.dtype] = None,
        output_index: Optional[int] = None,
    ) -> None:
        """Populate standard ONNX metadata on an FX node."""
        node.meta["onnx_op_type"] = op_type
        if name is not None:
            node.meta["onnx_name"] = name
        if domain is not None:
            node.meta["onnx_domain"] = domain
        if shape is not None:
            node.meta["onnx_shape"] = shape
        if dtype is not None:
            node.meta["onnx_dtype"] = dtype
        if output_index is not None:
            node.meta["onnx_output_index"] = output_index

    def get_value(self, name: str) -> torch.fx.Node:
        """Get a value (node) by name from the environment.

        Parameters
        ----------
        name : str
            The name of the value.

        Returns
        -------
        torch.fx.Node
            The corresponding FX node.

        Raises
        ------
        KeyError
            If the name is not found in the environment.
        """
        if name not in self.env:
            raise ValueNotFoundError(name, available=list(self.env.keys()))
        return self.env[name]

    def has_value(self, name: str) -> bool:
        """Check if a value exists in the environment."""
        return name in self.env

    def call_function(
        self,
        func: Callable[..., Any],
        args: Sequence[Union[torch.fx.Node, Any]] = (),
        kwargs: Optional[Dict[str, Any]] = None,
    ) -> torch.fx.Node:
        """Create a function call node in the FX graph.

        Parameters
        ----------
        func : Callable[..., Any]
            The function to call. Can be a PyTorch function, lambda, or any callable.
        args : Sequence[Union[torch.fx.Node, Any]], optional
            Positional arguments to the function. Can include FX nodes or constants.
        kwargs : Optional[Dict[str, Any]], optional
            Keyword arguments to the function.

        Returns
        -------
        torch.fx.Node
            The FX node representing this function call.
        """
        fx_node = self.graph.call_function(func, args=tuple(args), kwargs=kwargs or {})
        return fx_node

    def register_submodule(self, name: str, module: torch.nn.Module) -> str:
        """Register a submodule for use in the graph.

        Parameters
        ----------
        name : str
            Base name for the submodule.
        module : torch.nn.Module
            The submodule to register.

        Returns
        -------
        str
            The actual name used (may be modified to avoid conflicts).
        """
        # Sanitize name
        safe_name = sanitize_name(name)
        # Ensure unique name
        if safe_name in self._submodules:
            counter = 0
            while f"{safe_name}_{counter}" in self._submodules:
                counter += 1
            safe_name = f"{safe_name}_{counter}"
        self._submodules[safe_name] = module
        return safe_name

    def call_module(
        self,
        module_name: str,
        args: Sequence[Union[torch.fx.Node, Any]] = (),
        kwargs: Optional[Dict[str, Any]] = None,
    ) -> torch.fx.Node:
        """Create a module call node in the FX graph.

        Parameters
        ----------
        module_name : str
            The name of a registered submodule.
        args : Sequence[Union[torch.fx.Node, Any]], optional
            Positional arguments to the module.
        kwargs : Optional[Dict[str, Any]], optional
            Keyword arguments to the module.

        Returns
        -------
        torch.fx.Node
            The FX node representing this module call.
        """
        fx_node = self.graph.call_module(
            module_name, args=tuple(args), kwargs=kwargs or {}
        )
        return fx_node

    def _create_value_info_map(
        self,
    ) -> Dict[str, Tuple[Optional[List[Optional[int]]], Optional[torch.dtype]]]:
        """Build a mapping from value names to their shape and dtype info."""

        def extract_tensor_shape(
            value: onnx.ValueInfoProto,
        ) -> Optional[List[Optional[int]]]:
            """Extract a list-based representation of a tensor shape from a value info."""

            tensor_type = value.type.tensor_type
            if not tensor_type.HasField("shape"):
                return None
            dims: List[Optional[int]] = []
            for dim in tensor_type.shape.dim:
                if dim.HasField("dim_value"):
                    dims.append(int(dim.dim_value))
                elif dim.HasField("dim_param"):
                    dims.append(None)
                else:
                    dims.append(None)
            return dims

        def extract_tensor_dtype(value: onnx.ValueInfoProto) -> Optional[torch.dtype]:
            """Extract the Torch dtype that corresponds to a value info."""

            onnx_dtype = value.type.tensor_type.elem_type
            if onnx_dtype == 0:
                return None
            torch_dtype = DTYPE_MAP.get(onnx_dtype)
            if torch_dtype is None:
                if onnx_dtype == onnx.TensorProto.STRING:
                    return None
                raise UnsupportedDTypeError(
                    onnx_dtype=onnx_dtype,
                    tensor_name=value.name,
                    details="value_info dtype not supported",
                )
            return torch_dtype

        info_map = {}
        for value_info in (
            list(self.model.graph.input)
            + list(self.model.graph.value_info)
            + list(self.model.graph.output)
        ):
            info_map[value_info.name] = (
                extract_tensor_shape(value_info),
                extract_tensor_dtype(value_info),
            )
        return info_map

    def is_optional_type(self, name: str) -> bool:
        """Check if a value has optional type in the ONNX model.

        Parameters
        ----------
        name : str
            The name of the value to check.

        Returns
        -------
        bool
            True if the value has optional type, False otherwise.
        """
        # Search in graph inputs
        for value_info in self.model.graph.input:
            if value_info.name == name:
                return value_info.type.HasField("optional_type")
        # Search in value_info
        for value_info in self.model.graph.value_info:
            if value_info.name == name:
                return value_info.type.HasField("optional_type")
        # Search in outputs
        for value_info in self.model.graph.output:
            if value_info.name == name:
                return value_info.type.HasField("optional_type")
        return False

    def _create_initializer_map(self) -> Dict[str, torch.Tensor]:
        """Build a mapping from initializer names to PyTorch tensors."""
        init_map = {}
        for initializer in self.model.graph.initializer:
            init_map[initializer.name] = self.load_tensor(initializer)
        return init_map

    def load_tensor(self, tensor: onnx.TensorProto) -> torch.Tensor:
        """Load an ONNX TensorProto into a Torch tensor."""
        onnx_dtype = tensor.data_type
        if DTYPE_MAP.get(onnx_dtype) is None:
            raise UnsupportedDTypeError(
                onnx_dtype=onnx_dtype,
                tensor_name=tensor.name or "<unnamed>",
                details="initializer dtype not supported",
            )

        if self._memmap_external_data and (
            tensor.data_location == onnx.TensorProto.EXTERNAL or tensor.external_data
        ):
            info = resolve_external_data(
                tensor,
                base_dir=self._base_dir,
                strict=True,
            )
            memmap_array = np.memmap(
                info.path,
                dtype=info.numpy_dtype,
                mode="r",
                offset=info.offset,
                shape=info.shape,
            )
            return torch.from_numpy(memmap_array)

        np_array = numpy_helper.to_array(tensor)
        return torch.from_numpy(np_array.copy())

    def _load_initializers(self) -> None:
        """Load ONNX initializers as constant nodes in the FX graph."""
        for name, tensor in self.initializer_map.items():
            # Store in constants dict for later registration as buffers
            safe_name = sanitize_name(name)
            self._constants[safe_name] = tensor

            # Create a get_attr node to access the buffer
            fx_node = self.graph.get_attr(safe_name)
            self._set_onnx_metadata(
                fx_node,
                op_type="Initializer",
                name=name,
                shape=list(tensor.shape),
                dtype=tensor.dtype,
            )
            self.env[name] = fx_node

    def _create_placeholders(self) -> None:
        """Create FX placeholder nodes for graph inputs.

        Note: Inputs that are already loaded as initializers are skipped.
        """
        for value in self.model.graph.input:
            # Skip if already loaded as initializer
            if value.name in self.env:
                continue

            # Sanitize name for valid Python identifier
            safe_name = sanitize_name(value.name)
            placeholder = self.graph.placeholder(safe_name)
            info = self.value_info_map.get(value.name)
            self._set_onnx_metadata(
                placeholder,
                op_type="Input",
                name=value.name,
                shape=info[0] if info else None,
                dtype=info[1] if info else None,
            )
            self.env[value.name] = placeholder
            self.input_names.append(value.name)

    def _convert_nodes(self) -> None:
        # Get graph inputs and initializers for topological sort
        graph_inputs = {inp.name for inp in self.model.graph.input}
        initializers = set(self.initializer_map.keys())

        # Topologically sort nodes to handle out-of-order dependencies
        sorted_nodes = _topological_sort(
            list(self.model.graph.node),
            graph_inputs,
            initializers,
        )

        for node in sorted_nodes:
            # Get handler with domain and opset version support
            handler, domain, _opset = self._resolve_handler(node)
            fx_node = handler(self, node)

            # Add ONNX metadata to the operator node
            # Some handlers return a list of nodes (e.g., gradient ops)
            self._tag_operator_node(node, fx_node, domain)
            self._register_outputs(node, fx_node, domain)

    def _create_outputs(self) -> None:
        output_nodes = [self.get_value(value.name) for value in self.model.graph.output]
        if len(output_nodes) == 1:
            self.graph.output(output_nodes[0])
        else:
            self.graph.output(tuple(output_nodes))


__all__ = ["GraphBuilder"]
