from __future__ import annotations

import ctypes
import datetime
import inspect
import itertools
import logging
import operator
import os
import pathlib
import textwrap
import traceback
import typing
from typing import Any, Literal, Sequence

import numpy as np
import onnx
import onnxscript
import onnxscript.evaluator
import onnxscript.function_libs
import onnxscript.function_libs.torch_lib
import onnxscript.function_libs.torch_lib.registration
import torch
import torch.fx
from onnxscript import ir
from onnxscript.ir import convenience as ir_convenience
from torch.export import graph_signature

from torch_onnx import (
    _building,
    _decomp,
    _dispatching,
    _fx_passes,
    _ir_passes,
    _isolated,
    _onnx_program,
    _registration,
    _reporting,
    _tensors,
    _torchscript_converter,
    errors,
)

# Define utilities to convert PyTorch data types so users do not need to specify manually
_TORCH_DTYPE_TO_ONNX: dict[torch.dtype, ir.DataType] = {
    torch.bfloat16: ir.DataType.BFLOAT16,
    torch.bool: ir.DataType.BOOL,
    torch.complex128: ir.DataType.COMPLEX128,
    torch.complex64: ir.DataType.COMPLEX64,
    torch.float16: ir.DataType.FLOAT16,
    torch.float32: ir.DataType.FLOAT,
    torch.float64: ir.DataType.DOUBLE,
    torch.float8_e4m3fn: ir.DataType.FLOAT8E4M3FN,
    torch.float8_e4m3fnuz: ir.DataType.FLOAT8E4M3FNUZ,
    torch.float8_e5m2: ir.DataType.FLOAT8E5M2,
    torch.float8_e5m2fnuz: ir.DataType.FLOAT8E5M2FNUZ,
    torch.int16: ir.DataType.INT16,
    torch.int32: ir.DataType.INT32,
    torch.int64: ir.DataType.INT64,
    torch.int8: ir.DataType.INT8,
    torch.uint8: ir.DataType.UINT8,
}
_BLUE = "\033[96m"
_END = "\033[0m"

logger = logging.getLogger(__name__)


def _torch_dtype_to_onnx_dtype(dtype: torch.dtype) -> ir.DataType:
    return _TORCH_DTYPE_TO_ONNX[dtype]


class TorchTensor(ir.Tensor):
    def __init__(self, tensor: torch.Tensor, name: str | None = None):
        # Pass the tensor as the raw data to ir.Tensor's constructor
        super().__init__(
            tensor, dtype=_torch_dtype_to_onnx_dtype(tensor.dtype), name=name
        )

    def __array__(self, dtype: Any = None) -> np.ndarray:
        # numpy() calls __array__ in ir.Tensor
        if self.dtype == ir.DataType.BFLOAT16:
            return self.raw.view(torch.uint16).__array__(dtype)
        if self.dtype in {
            ir.DataType.FLOAT8E4M3FN,
            ir.DataType.FLOAT8E4M3FNUZ,
            ir.DataType.FLOAT8E5M2,
            ir.DataType.FLOAT8E5M2FNUZ,
        }:
            # TODO: Use ml_dtypes
            return self.raw.view(torch.uint8).__array__(dtype)
        return self.raw.__array__(dtype)

    def tobytes(self) -> bytes:
        # Implement tobytes to support native PyTorch types so we can use types like bloat16
        # Reading from memory directly is also more efficient because
        # it avoids copying to a NumPy array
        tensor = self.raw.detach().cpu().contiguous()
        return bytes(
            (ctypes.c_ubyte * tensor.element_size() * tensor.numel()).from_address(
                tensor.data_ptr()
            )
        )


# https://github.com/pytorch/pytorch/blob/ee6cb6daa173896f8ea1876266a19775aaa4f610/torch/export/graph_signature.py#L56C1-L62C19
# class InputKind(Enum):
#     USER_INPUT = auto()
#     PARAMETER = auto()
#     BUFFER = auto()
#     CONSTANT_TENSOR = auto()
#     CUSTOM_OBJ = auto()
#     TOKEN = auto()

# https://github.com/pytorch/pytorch/blob/ee6cb6daa173896f8ea1876266a19775aaa4f610/torch/export/graph_signature.py#L89C1-L96C19
# class OutputKind(Enum):
#     USER_OUTPUT = auto()
#     LOSS_OUTPUT = auto()
#     BUFFER_MUTATION = auto()
#     GRADIENT_TO_PARAMETER = auto()
#     GRADIENT_TO_USER_INPUT = auto()
#     USER_INPUT_MUTATION = auto()
#     TOKEN = auto()


def _set_shape_types(
    values: Sequence[ir.Value],
    meta_vals: Sequence[torch.Tensor],
    complex_to_float: bool = True,
):
    for value, meta_val in zip(values, meta_vals):
        _set_shape_type(value, meta_val, complex_to_float=complex_to_float)


def _set_shape_type(
    value: ir.Value,
    meta_val: torch.Tensor | tuple[torch.Tensor],
    complex_to_float: bool,
):
    # TODO: Consider using meta["tensor_meta"] for this? Would it be faster?
    if isinstance(meta_val, tuple):
        logger.warning("Setting shape and type of tensors is not supported yet")
    if isinstance(meta_val, torch.Tensor):
        # FIXME: Consider shape for complex values
        dims = []
        for dim in meta_val.shape:
            if isinstance(dim, int):
                dims.append(dim)
            else:
                dims.append(str(dim.node))
        value.dtype = _torch_dtype_to_onnx_dtype(meta_val.dtype)
        if complex_to_float:
            if meta_val.dtype == torch.complex64:
                value.dtype = ir.DataType.FLOAT
                # Add 2 as the last dimension if the tensor is complex to hold the real/imag parts
                dims.append(2)
            elif meta_val.dtype == torch.complex128:
                value.dtype = ir.DataType.DOUBLE
                # Add 2 as the last dimension if the tensor is complex to hold the real/imag parts
                dims.append(2)

        value.shape = ir.Shape(dims)
    elif isinstance(meta_val, (int, torch.SymInt)):
        # aten::sym_size output is a int, not a tensor, which stands
        # for the size of one dim. We treat it as a scalar.
        value.dtype = ir.DataType.INT64
        value.shape = ir.Shape([])
    elif isinstance(meta_val, (bool, torch.SymBool)):
        value.dtype = ir.DataType.BOOL
        value.shape = ir.Shape([])
    elif isinstance(meta_val, (float, torch.SymFloat)):
        value.dtype = ir.DataType.FLOAT
        value.shape = ir.Shape([])
    else:
        pass


def _get_qualified_module_name(cls: Any) -> str:
    if isinstance(cls, str):
        return cls
    module = cls.__module__
    if module is None or module == str.__class__.__module__:
        return cls.__name__
    return module + "." + cls.__name__


def _get_node_namespace(node: torch.fx.Node) -> tuple[str, list[str], list[str]]:
    """Get the namespace and scope of the node.

    Example::

        {
            'L__self__': ('', <class 'torchvision.models.resnet.ResNet'>),
            'L__self___avgpool': ('avgpool', <class 'torch.nn.modules.pooling.AdaptiveAvgPool2d'>)
        }

    Will yield

    namespace: ": torchvision.models.resnet.ResNet/avgpool: torch.nn.modules.pooling.AdaptiveAvgPool2d/node_name: node_target"
    class_hierarchy: ["torchvision.models.resnet.ResNet", "torch.nn.modules.pooling.AdaptiveAvgPool2d", <node_target>]
    name_scopes: ["", "avgpool", <node_name>]

    Args:
        node: The node to get the namespace and scope of.

    Returns:
        (namespace, class_hierarchy, name_scope)
    """
    nn_module_stack = node.meta.get("nn_module_stack")
    logger.debug("%s", nn_module_stack)
    if nn_module_stack is None:
        logger.warning("nn_module_stack not found for node %s", node.name)
        return f"{node.name}: {node.target}", [str(node.target)], [node.name]
    namespaces = []
    class_hierarchy = []
    name_scopes = []
    for name, nn_module in nn_module_stack.values():
        name_scopes.append(name)
        nn_module_name = _get_qualified_module_name(nn_module)
        class_hierarchy.append(nn_module_name)
        namespaces.append(f"{name}: {_get_qualified_module_name(nn_module)}")
    namespaces.append(f"{node.name}: {node.target}")
    class_hierarchy.append(str(node.target))
    name_scopes.append(node.name)

    return "/".join(namespaces), class_hierarchy, name_scopes


def _set_node_metadata(fx_node: torch.fx.Node, ir_node: ir.Node) -> None:
    """Adds namespace and other node metadata to the ONNX node."""
    namespace, class_hierarchy, name_scopes = _get_node_namespace(fx_node)
    ir_node.metadata_props["namespace"] = namespace
    ir_node.metadata_props["pkg.torch.onnx.class_hierarchy"] = repr(class_hierarchy)
    ir_node.metadata_props["pkg.torch.onnx.name_scopes"] = repr(name_scopes)
    ir_node.metadata_props["pkg.torch.onnx.fx_node"] = str(fx_node.format_node())
    ir_node.metadata_props["pkg.torch.onnx.stack_trace"] = fx_node.meta.get(
        "stack_trace", ""
    )


def _handle_getitem_node(
    node: torch.fx.Node, node_name_to_values: dict[str, ir.Value | Sequence[ir.Value]]
) -> ir.Value:
    """Handle a getitem node.

    Add the input value it is getting to the mapping, then return the value.

    There are two cases for this node:
    1. The output is a Sequence (traced), we can simply get the value from the sequence
    2. The output is produced by a SplitToSequence node, we need to get the value from the sequence value
    This function only handles the first case
    """
    assert len(node.all_input_nodes) == 1
    source = node.all_input_nodes[0]
    source_outputs = node_name_to_values[source.name]
    assert isinstance(
        source_outputs, Sequence
    ), f"Expected {source.name} to output sequence, got {node_name_to_values[source.name]}"
    index = typing.cast(int, node.args[1])
    value = source_outputs[index]
    # Save the getitem value to the values mapping to in case
    # it is one of the graph outputs
    node_name_to_values[node.name] = value
    # Rename the name of value with the getitem name.
    value.name = node.name
    return value


def _handle_call_function_node(
    graph: ir.Graph,
    node: torch.fx.Node,
    node_name_to_values: dict[str, ir.Value | Sequence[ir.Value]],
):
    if node.target == operator.getitem:
        _handle_getitem_node(node, node_name_to_values)
    # Add op to the graph
    op = str(node.target)
    fx_inputs, attributes, input_names, output_names = _get_inputs_and_attributes(node)
    inputs = []
    for i, input_ in enumerate(fx_inputs):
        if input_ is None:
            inputs.append(None)
        elif hasattr(input_, "name"):
            if isinstance(input_, torch.fx.Node) and input_.target == operator.getitem:
                actual_input = _handle_getitem_node(input_, node_name_to_values)
                inputs.append(actual_input)
            else:
                inputs.append(node_name_to_values[input_.name])
        else:
            attributes[f"arg_{i}"] = input_

    outputs = [ir.Value(name=name) for name in output_names]
    if len(outputs) > 1:
        _set_shape_types(outputs, node.meta["val"], complex_to_float=False)
        node_name_to_values[node.name] = outputs
    else:
        _set_shape_type(outputs[0], node.meta["val"], complex_to_float=False)
        node_name_to_values[node.name] = outputs[0]
    ir_node = ir.Node(
        "pkg.torch.ops",
        op,
        inputs,
        attributes=ir_convenience.convert_attributes(attributes),
        outputs=outputs,
        name=node.name,
    )
    ir_node.meta["node"] = node
    ir_node.metadata_props["pkg.torch.onnx.input_names"] = repr(input_names)
    # Record the nn.Module stack for the node
    _set_node_metadata(node, ir_node)

    graph.append(ir_node)


def _convert_fx_arg_to_onnx_arg(
    arg, node_name_to_values: dict[str, ir.Value | Sequence[ir.Value]]
):
    """Convert an FX argument to an ONNX compatible argument.

    This function
    - Converts a torch dtype to an integer
    - Converts a torch device/memory_format/layout to a string
    - Converts a torch.fx.Node to an ir.Value
    - Converts a sequence of torch.fx.Node to a sequence of ir.Value
    """
    if arg is None:
        # None arguments are not modified because when the arg is an ONNX input
        # we need to preserve the None value; when the arg is an ONNX attribute,
        # we want to drop the value.
        # The actual dropping of a None attribute value is done by OpRecorder
        return None
    if hasattr(arg, "name"):
        if isinstance(arg, torch.fx.Node) and arg.target == operator.getitem:
            source = arg.all_input_nodes[0]
            source_outputs = node_name_to_values[source.name]
            if isinstance(source_outputs, Sequence):
                # If the node is getting an input from another node, get the actual value the node is retrieving
                return _handle_getitem_node(arg, node_name_to_values)
            else:
                # `source_outputs` is a sequence(tensor()) value and we need to
                # use SequenceAt to get the value. This is handled by torchlib
                pass
        # If the input is a node, get the value from the mapping
        return node_name_to_values[arg.name]
    if isinstance(arg, (list, tuple)):
        return [_convert_fx_arg_to_onnx_arg(elem, node_name_to_values) for elem in arg]
    if isinstance(arg, (torch.device, torch.memory_format, torch.layout)):
        return str(arg)
    if isinstance(arg, torch.dtype):
        return _torch_dtype_to_onnx_dtype(arg)
    # Maybe a Python value
    return arg


def _get_onnxscript_opset(opset_version: int) -> onnxscript.values.Opset:
    return onnxscript.values.Opset("", opset_version)


def _handle_call_function_node_with_lowering(
    model: ir.Model,
    node: torch.fx.Node,
    node_name_to_values: dict[str, ir.Value | Sequence[ir.Value]],
    constant_farm: dict[Any, ir.Value],
    registry: _registration.OnnxRegistry,
    opset: onnxscript.values.Opset,
):
    if node.target == operator.getitem:
        source = node.all_input_nodes[0]
        source_outputs = node_name_to_values[source.name]
        if isinstance(source_outputs, Sequence):
            _handle_getitem_node(node, node_name_to_values)
            return
        else:
            # `source_outputs` is a sequence(tensor()) value and we need to
            # use SequenceAt to get the value. This is handled by torchlib
            pass

    # Find the matching ONNX overload for the node
    # NOTE: Create different registries for different ONNX opset versions
    # TODO: Log the message here to expose false positives
    onnx_function, message = _dispatching.dispatch(node, registry)

    if onnx_function is None:
        # TODO(justinchuby): Fall back to ATen op or do something else?
        raise errors.DispatchError(
            f"No ONNX function found for {node.target!r}. Failure message: {message}"
        )

    # Map FX inputs to ONNX inputs and fill optional inputs.
    # torch_args and torch_kwargs are for op-level validation
    fx_args = node.args
    fx_kwargs = node.kwargs

    # Replace the input FX nodes with ONNX values
    onnx_args = [
        _convert_fx_arg_to_onnx_arg(input_, node_name_to_values) for input_ in fx_args
    ]

    onnx_kwargs = {}
    for key, value in fx_kwargs.items():
        onnx_kwargs[key] = _convert_fx_arg_to_onnx_arg(value, node_name_to_values)
        if key == "dtype" and onnx_kwargs[key] is None:
            # Set dtype to -1 if it is None
            onnx_kwargs[key] = -1

    with onnxscript.evaluator.default_as(
        tracer := _building.OpRecorder(opset, constant_farm)
    ):
        try:
            outputs = onnx_function(*onnx_args, **onnx_kwargs)
        except Exception as e:
            raise errors.GraphConstructionError(
                f"Error when calling function '{onnx_function}' with args '{onnx_args}' and kwargs '{onnx_kwargs}'"
            ) from e

    output_names = []
    try:
        node_schema = node.target._schema
    except AttributeError:
        logger.debug("No schema found for node %s", node.format_node())
    else:
        for i, output in enumerate(node_schema.returns):
            if not output.name:
                if i == 0:
                    output_names.append(f"val_{node.name}")
                else:
                    # More than one output may not have a name, in which case we
                    # append an index to avoid duplication
                    output_names.append(f"val_{node.name}__{i}")
            else:
                output_names.append(f"val_{node.name}_{output.name}")

    if isinstance(outputs, Sequence):
        _set_shape_types(outputs, node.meta["val"], complex_to_float=True)
        node_name_to_values[node.name] = outputs
        if output_names:
            for output, name in zip(outputs, output_names):
                output.name = name
        else:
            for i, output in enumerate(outputs):
                output.name = f"val_{node.name}_{i}"
    else:
        _set_shape_type(outputs, node.meta["val"], complex_to_float=True)
        node_name_to_values[node.name] = outputs
        if output_names:
            outputs.name = output_names[0]
        else:
            outputs.name = f"val_{node.name}"

    for ir_node in tracer.nodes:
        ir_node.meta["node"] = node

        # Record the nn.Module stack for the node
        _set_node_metadata(node, ir_node)

    # Add the traced nodes to the graph
    model.graph.extend(tracer.nodes)
    # Add the defined functions to the model
    for identifier, onnxscript_function in tracer.functions.items():
        if identifier in model.functions:
            continue
        # TODO: Get IR function directly when onnxscript is updated
        proto = onnxscript_function.to_function_proto()
        ir_function = ir.serde.deserialize_function(proto)
        model.functions[identifier] = ir_function
        if ir_function.domain not in model.opset_imports:
            # FIXME: Record the opset version of the function
            model.opset_imports[ir_function.domain] = 1


def _handle_placeholder_node(
    node: torch.fx.Node,
    node_name_to_values: dict[str, ir.Value | Sequence[ir.Value]],
    *,
    lower: str,
    opset: onnxscript.values.Opset,
) -> None:
    # Placeholder nodes are user inputs
    # We need to create a new tensor for each user input
    # and add it to the graph's inputs
    name = node.name
    input_ = _tensors.SymbolicTensor(opset, name=name)
    input_.meta["node"] = node
    _set_shape_type(input_, node.meta["val"], complex_to_float=lower != "none")
    node_name_to_values[name] = input_
    # The inputs will be added to the graph later


def _add_nodes(
    exported_program: torch.export.ExportedProgram,
    model: ir.Model,
    lower: Literal["at_conversion", "post_conversion", "none"],
    registry: _registration.OnnxRegistry,
) -> dict[str, ir.Value | Sequence[ir.Value]]:
    node_name_to_values: dict[str, ir.Value | Sequence[ir.Value]] = {}
    constant_farm = {}
    opset = _get_onnxscript_opset(registry.opset_version)
    for node in exported_program.graph.nodes:
        logger.debug(
            "%s", (node.name, node.args, node.target, node.op, node.type, node.kwargs)
        )
        try:
            if node.op == "placeholder":
                _handle_placeholder_node(
                    node,
                    node_name_to_values,
                    lower=lower,
                    opset=opset,
                )
            elif node.op == "call_function":
                if lower == "at_conversion":
                    _handle_call_function_node_with_lowering(
                        model,
                        node,
                        node_name_to_values,
                        constant_farm,
                        registry=registry,
                        opset=opset,
                    )
                else:
                    # No lowering
                    _handle_call_function_node(model.graph, node, node_name_to_values)
        except Exception as e:
            raise errors.OnnxConversionError(
                f"Error when translating node {node.format_node()}. See the stack trace for more information."
            ) from e
    return node_name_to_values


def _torch_version_integer() -> int:
    return int(torch.__version__.replace(".", "").split("dev")[0])


def _get_inputs_and_attributes(
    node: torch.fx.Node,
) -> tuple[list[torch.fx.Node | None], dict[str, Any], list[str], list[str]]:
    """Find and Fill in the not provided kwargs with default values.

    Returns:
        (inputs, attributes, input_names, output_names)
    """
    if inspect.isbuiltin(node.target) or isinstance(node.target, str):
        inputs = list(node.args)
        return inputs, {}, [], [node.name]

    # The target should be an ATen operator now
    assert hasattr(
        node.target, "_schema"
    ), f"The target should be an ATen operator now, but node target {node.target} has no schema"
    node_schema: torch.FunctionSchema = node.target._schema

    # This function assumes the order of arguments in FX op is the
    # same as the order of arguments in TorchScript op.
    inputs = []
    input_names = []
    attributes = {}

    if inspect.isbuiltin(node.target):
        inputs = list(node.args)
    else:
        for arg, schema_arg in zip(node.args, node_schema.arguments):
            if arg is None or isinstance(arg, torch.fx.Node):
                inputs.append(arg)
                input_names.append(schema_arg.name)
            elif isinstance(arg, Sequence) and all(
                elem is None or isinstance(elem, torch.fx.Node) for elem in arg
            ):
                inputs.extend(arg)
                input_names.extend([schema_arg.name] * len(arg))
            elif isinstance(arg, torch.device):
                attributes[schema_arg.name] = str(arg)
            elif isinstance(arg, torch.dtype):
                attributes[schema_arg.name] = _torch_dtype_to_onnx_dtype(arg)
            else:
                attributes[schema_arg.name] = arg
        for schema_arg in node_schema.arguments:
            if schema_arg.name not in node.kwargs:
                continue
            kwarg = node.kwargs[schema_arg.name]
            if schema_arg.name in {
                "layout",
                "device",
                "requires_grad",
                "memory_format",
                "implicit",
            } or isinstance(kwarg, torch.device):
                attr = str(kwarg)
            elif isinstance(kwarg, torch.dtype):
                attr = _torch_dtype_to_onnx_dtype(kwarg)
            else:
                attr = kwarg

            attributes[schema_arg.name] = attr

    output_names = [f"{node.name}_{output.name}" for output in node_schema.returns]

    return inputs, attributes, input_names, output_names


def _maybe_start_profiler(should_profile: bool) -> Any:
    if should_profile:
        import pyinstrument

        profiler = pyinstrument.Profiler(async_mode="disabled")
        profiler.start()
        return profiler
    return None


def _maybe_stop_profiler_and_get_result(profiler) -> str | None:
    if profiler is None:
        return None
    profiler.stop()
    return profiler.output_text(unicode=True)


def _format_exception(e: Exception) -> str:
    return "\n".join(traceback.format_exception(type(e), e, e.__traceback__))


def exported_program_to_ir(
    exported_program: torch.export.ExportedProgram,
    *,
    registry: _registration.OnnxRegistry | None = None,
    lower: Literal["at_conversion", "post_conversion", "none"] = "at_conversion",
) -> ir.Model:
    """Convert an exported program to an ONNX IR model.

    Reference:
        - ExportedProgram spec: https://pytorch.org/docs/stable/export.ir_spec.html

    Args:
        exported_program: The exported program to convert.
        lower: Whether to lower the graph to core ONNX operators.
            at_conversion: Lower whe translating the FX graph to ONNX IR.
            post_conversion: Use an IR pass to lower the graph.
            none: Do not lower the graph.
        registry: The registry of all ONNX Script decomposition.
    """
    if registry is None:
        # Trigger op registration
        from onnxscript.function_libs.torch_lib import ops

        del ops
        registry = _registration.OnnxRegistry.from_torchlib(
            onnxscript.function_libs.torch_lib.registration.default_registry
        )
    model = ir.Model(
        graph=ir.Graph(
            [],
            [],
            nodes=[],
            opset_imports={
                "": registry.opset_version,
            },
            name="main_graph",
        ),
        ir_version=9,
        producer_name="torch",
        producer_version=torch.__version__,
    )

    if lower == "none":
        # Add the opset import for the torch ops
        model.opset_imports["pkg.torch.ops"] = _torch_version_integer()
    # NOTE: Function domains are added when translating nodes when lower="at_conversion"

    # 1. Add all nodes to the graph and create a dictionary of values
    if lower != "none":
        # Decompose the graph given the implemented torch ops in ONNX
        decomp_table = _decomp.create_onnx_friendly_decomposition_table(registry)
        exported_program.run_decompositions(decomp_table)

        # Include explicit type promotion nodes
        _fx_passes.insert_type_promotion_nodes(exported_program)
        _fx_passes.remove_assertion_nodes(exported_program)
    values = _add_nodes(exported_program, model, lower=lower, registry=registry)

    # 2. Add user inputs and all parameters/buffers to the graph.
    # Since the node names and the tensor names are different, we need to rename
    # the nodes to match the tensor names later. For now we will just use the node names.
    user_inputs = [
        spec
        for spec in exported_program.graph_signature.input_specs
        if spec.kind == graph_signature.InputKind.USER_INPUT
    ]
    non_user_inputs = [
        spec
        for spec in exported_program.graph_signature.input_specs
        if spec.kind != graph_signature.InputKind.USER_INPUT
    ]

    for spec in itertools.chain(user_inputs, non_user_inputs):
        # Put the user inputs first and then the parameters/buffers
        if isinstance(spec.arg, graph_signature.ConstantArgument):
            logger.debug("Skipping constant argument %s", spec.arg)
            continue
        value_name = spec.arg.name
        input_kind = spec.kind
        persistent = spec.persistent
        value = values[value_name]

        value.metadata_props["pkg.torch.export.graph_signature.InputSpec.kind"] = (
            input_kind.name
        )
        value.metadata_props[
            "pkg.torch.export.graph_signature.InputSpec.persistent"
        ] = str(persistent)

        model.graph.inputs.append(value)  # type: ignore
        if input_kind != graph_signature.InputKind.USER_INPUT:
            model.graph.initializers[value_name] = value

    # 3. Add outputs to the graph. Put the user outputs first.
    user_outputs = [
        spec
        for spec in exported_program.graph_signature.output_specs
        if spec.kind == graph_signature.OutputKind.USER_OUTPUT
    ]
    non_user_outputs = [
        spec
        for spec in exported_program.graph_signature.output_specs
        if spec.kind != graph_signature.OutputKind.USER_OUTPUT
    ]
    for spec in itertools.chain(user_outputs, non_user_outputs):
        if isinstance(spec.arg, graph_signature.ConstantArgument):
            logger.warning("Skipping constant argument %s", spec.arg)
            continue
        value_name = spec.arg.name
        output_kind = spec.kind
        value = values[value_name]

        if not isinstance(value, ir.Value):
            raise TypeError(
                f"Output '{value_name}' should be an ir.Value. Actual type is '{type(value)}': {value!r}. "
                "This may be due to an incorrect implementation of the ONNX function that produced this output."
            )

        value.metadata_props["pkg.torch.export.graph_signature.OutputSpec.kind"] = (
            output_kind.name
        )

        if output_kind == graph_signature.OutputKind.USER_OUTPUT:
            model.graph.outputs.append(value)

    # 4. Rename the initializers to match the tensor names
    for name, param_name in itertools.chain(
        exported_program.graph_signature.inputs_to_parameters.items(),
        exported_program.graph_signature.inputs_to_buffers.items(),
        exported_program.graph_signature.inputs_to_lifted_tensor_constants.items(),
    ):
        initializer = model.graph.initializers.pop(name)
        initializer.name = param_name
        # Record the original name so users can search the metadata and correspond
        # with the FX graph
        initializer.metadata_props["pkg.torch.onnx.original_node_name"] = name
        model.graph.initializers[param_name] = initializer

    # 5. Add initializers to the graph
    # ExportedProgram stores parameters and buffers in state_dict,
    # but non_persistent_buffers and lifted_tensor_constants are not there
    # so we need to get them from the name_* apis.
    for name, torch_tensor in itertools.chain(
        exported_program.named_parameters(),
        exported_program.named_buffers(),
        exported_program.constants.items(),
    ):
        initializer = model.graph.initializers.get(name)
        if initializer is None:
            logger.warning("Tensor '%s' is not one of the initializers", name)
            continue
        if not isinstance(torch_tensor, torch.Tensor):
            raise NotImplementedError(
                f"Tensor '{name}' should be a torch.Tensor. Actual type is '{type(torch_tensor)}': {torch_tensor!r}. "
                "This is unexpected and not yet supported."
            )
        ir_tensor = TorchTensor(torch_tensor, name=name)
        initializer.const_value = ir_tensor
        _set_shape_type(
            initializer,
            torch_tensor,
            complex_to_float=lower != "none",
        )

    # TODO: Decide if we should keep mutated buffers as inputs/outputs

    return model


def _take_first_line(text: str) -> str:
    """Take the first line of a text."""
    lines = text.split("\n", maxsplit=1)
    first_line = lines[0]
    if len(lines) > 1:
        first_line += "[...]"
    return first_line


def export(
    model: torch.nn.Module | torch.export.ExportedProgram,
    args: tuple[Any, ...],
    kwargs: dict[str, Any] | None = None,
    *,
    registry: _registration.OnnxRegistry | None = None,
    dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any] | None = None,
    input_names: Sequence[str] | None = None,
    output_names: Sequence[str] | None = None,
    profile: bool = False,
    error_report: bool = False,
    dump_exported_program: bool = False,
    artifacts_dir: str | os.PathLike = ".",
) -> _onnx_program.ONNXProgram:
    """Export a PyTorch model to ONNXProgram.

    Args:
        model: The model to export. This can be a PyTorch nn.Module or an ExportedProgram.
        args: The arguments to pass to the model.
        kwargs: The keyword arguments to pass to the model.
        registry: The registry of all ONNX decompositions.
        dynamic_shapes: Dynamic shapes in the graph.
        input_names: If provided, rename the inputs.
        output_names: If provided, rename the outputs.
        profile: Whether to profile the export process.
        error_report: Whether to generate an error report if the export fails.
        dump_exported_program: Whether to save the exported program to a file.
        artifacts_dir: The directory to save the exported program and error reports.

    Returns:
        The ONNXProgram with the exported IR graph.

    Raises:
        TorchExportError: If the export process fails with torch.export.
        OnnxConversionError: If the ExportedProgram to ONNX translation fails.
    """
    # Set up the error reporting facilities
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S-%f")
    profiler = _maybe_start_profiler(profile)

    # Create the artifacts directory if it does not exist
    artifacts_dir = pathlib.Path(artifacts_dir)
    if error_report or profile or dump_exported_program:
        artifacts_dir.mkdir(parents=True, exist_ok=True)

    # Step 0: Export the model with torch.export.export if the model is not already an ExportedProgram
    if isinstance(model, torch.export.ExportedProgram):
        program = model
    elif isinstance(model, (torch.jit.ScriptModule, torch.jit.ScriptFunction)):
        model_repr = _take_first_line(repr(model))
        print(
            f"Obtain model graph for `{model_repr}` with `TorchScript to ExportedProgram converter` because model is jit'ed..."
        )
        try:
            program = _torchscript_converter.TS2EPConverter(
                model, args, kwargs
            ).convert()
            print(
                f"Obtain model graph for `{model_repr}` with `TorchScript to ExportedProgram converter` because model is jit'ed... ✅"
            )
        except Exception as e:
            print(
                f"Obtain model graph for `{model_repr}` with `TorchScript to ExportedProgram converter` because model is jit'ed... ❌"
            )
            profile_result = _maybe_stop_profiler_and_get_result(profiler)

            if error_report:
                error_report_path = (
                    artifacts_dir / f"onnx_export_{timestamp}_pt_export.md"
                )
                _reporting.create_torch_export_error_report(
                    error_report_path,
                    # TODO(justinchuby): Check if we need to format the previous exception separately
                    _format_exception(e),
                    profile_result=profile_result,
                )
            else:
                error_report_path = None

            raise errors.TorchScriptConverterError(
                textwrap.dedent(f"""\
                    Failed to export the model with TorchScript converter. {_BLUE}This is step 1/2{_END} of exporting the model to ONNX. Next steps:
                    - Create an issue in the PyTorch GitHub repository against the {_BLUE}*torch.export*{_END} component and attach the full error stack as well as reproduction scripts.
                    """)
                + f"Error report has been saved to '{error_report_path}'."
                if error_report
                else ""
            ) from e
    else:
        model_repr = _take_first_line(repr(model))
        print(f"Obtain model graph for `{model_repr}` with `torch.export.export`...")
        try:
            program = torch.export.export(
                model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes
            )
            print(
                f"Obtain model graph for `{model_repr}` with `torch.export.export`... ✅"
            )
        except Exception as e_export:
            print(
                f"Obtain model graph for `{model_repr}` with `torch.export.export` but failed. Falling back to use torch.jit.trace... ⚠️"
            )
            # If torch.export.export fails, fall back to torchscript and try again.
            try:
                jit_model = torch.jit.trace(
                    model, example_inputs=args, check_trace=False, strict=False
                )
                if dump_exported_program:
                    program_path = artifacts_dir / f"onnx_export_{timestamp}.pt"
                    jit_model.save(program_path)
                    print(f"Torch Script model has been saved to '{program_path}'.")
                program = _torchscript_converter.TS2EPConverter(
                    jit_model, args, kwargs
                ).convert()
            except Exception as e_trace:
                print(
                    f"Obtain model graph for `{model_repr}` with `TorchScript to ExportedProgram converter` fails as well... ❌"
                )
                profile_result = _maybe_stop_profiler_and_get_result(profiler)

                if error_report:
                    error_report_path = (
                        artifacts_dir / f"onnx_export_{timestamp}_pt_export.md"
                    )
                    _reporting.create_torch_export_error_report(
                        error_report_path,
                        _format_exception(e_trace),
                        profile_result=profile_result,
                    )
                else:
                    error_report_path = None

                # NOTE: We throw e_export and not e_trace because we want to
                # focus on the torch.export.export error. The torch.jit.trace
                # error is due to the fallback and can be confusing to users.
                # e_trace is still saved in the error report and will include
                # e_export as well.
                raise errors.TorchExportError(
                    textwrap.dedent(f"""\
                        Failed to export the model with torch.export. {_BLUE}This is step 1/2{_END} of exporting the model to ONNX. Next steps:
                        - Modify the model code for `torch.export.export` to succeed. Refer to https://pytorch.org/docs/stable/generated/exportdb/index.html for more information.
                        - Debug `torch.export.export` and summit a PR to PyTorch.
                        - Create an issue in the PyTorch GitHub repository against the {_BLUE}*torch.export*{_END} component and attach the full error stack as well as reproduction scripts.
                        """)
                    + f"Error report has been saved to '{error_report_path}'."
                    if error_report
                    else ""
                ) from e_export

    if dump_exported_program:
        print("Dumping ExportedProgram because `dump_exported_program=True`...")
        program_path = f"onnx_export_{timestamp}.pt2"
        torch.export.save(program, program_path)
        print(f"ExportedProgram has been saved to '{program_path}'.")

    # Step 1: Convert the exported program to an ONNX model
    try:
        print("Translate the graph into ONNX...")
        ir_model = exported_program_to_ir(program, registry=registry)

        if input_names:
            _ir_passes.rename_inputs(ir_model, input_names)
        if output_names:
            _ir_passes.rename_outputs(ir_model, output_names)

        # TODO(justinchuby): Remove the hack
        _ir_passes.add_torchlib_common_imports(ir_model)

        onnx_program = _onnx_program.ONNXProgram(ir_model, program)
        print("Translate the graph into ONNX... ✅")

    except Exception as e:
        print("Translate the graph into ONNX... ❌")
        profile_result = _maybe_stop_profiler_and_get_result(profiler)

        if error_report:
            error_report_path = artifacts_dir / f"onnx_export_{timestamp}_conversion.md"

            # Run the analysis to get the error report
            _reporting.create_onnx_export_error_report(
                error_report_path,
                _format_exception(e),
                program,
                step=1,
                profile_result=profile_result,
                registry=registry,
            )
        else:
            error_report_path = None

        raise errors.OnnxConversionError(
            textwrap.dedent(f"""\
                Failed to convert the exported program to an ONNX model. {_BLUE}This is step 2/2{_END} of exporting the model to ONNX. Next steps:
                - If there is a missing ONNX function, implement it and register it to the registry.
                - If there is an internal error during ONNX conversion, debug the error and summit a PR to PyTorch.
                - Save the ExportedProgram as a pt2 file and create an error report with `export(error_report=True)`. Create an issue in the PyTorch GitHub repository against the {_BLUE}*onnx*{_END} component. Attach the pt2 model and the error report.
                """)
            + f"Error report has been saved to '{error_report_path}'."
            if error_report
            else ""
        ) from e

    profile_result = _maybe_stop_profiler_and_get_result(profiler)

    if dump_exported_program:
        program_path = artifacts_dir / f"onnx_export_{timestamp}.pt2"
        print("Dumping ExportedProgram because `dump_exported_program=True`...")
        torch.export.save(program, program_path)
        print(f"ExportedProgram has been saved to '{program_path}'.")

    if not error_report:
        # Return if error report is not requested
        if profile:
            assert profile_result is not None
            _reporting.crete_onnx_export_profile_report(
                artifacts_dir / f"onnx_export_{timestamp}_profile.md",
                onnx_program.exported_program,
                profile_result,
                step=1,
            )
        return onnx_program

    # Step 2: (When error report is requested) Check the ONNX model with ONNX checker
    try:
        print("Run `onnx.checker` on the ONNX model...")

        # TODO: Handle when model is >2GB

        model_proto = onnx_program.model_proto
        byte_size = model_proto.ByteSize()
        if byte_size < 2 * 1024 * 1024 * 1024:
            # The checker may segfault so we need to run it in a separate process
            _isolated.safe_call(
                onnx.checker.check_model, onnx_program.model_proto, full_check=True
            )
            print("Run `onnx.checker` on the ONNX model... ✅")
        else:
            print(
                f"Run `onnx.checker` on the ONNX model... ⚠️ Skipped because model is too large ({byte_size})."
            )
    except Exception as e:
        print("Run `onnx.checker` on the ONNX model... ❌")
        if error_report:
            _reporting.create_onnx_export_error_report(
                artifacts_dir / f"onnx_export_{timestamp}_checker.md",
                _format_exception(e),
                onnx_program.exported_program,
                step=2,
                profile_result=profile_result,
                model=onnx_program.model,
                registry=registry,
            )
        logger.warning(
            "Conversion successful but the ONNX model fails ONNX checker. "  # noqa: G004
            "Please create an issue "
            f"in the PyTorch GitHub repository against the {_BLUE}*onnx*{_END} component and "
            "attach the full error stack as well as reproduction scripts. ",
            exc_info=e,
        )
        return onnx_program

    # Step 3: (When error report is requested) Execute the model with ONNX Runtime
    # try:
    #     print("Execute the model with ONNX Runtime... ")
    #     print("✅")
    # except Exception as e:
    #     raise errors.OnnxConversionError(
    #         "Conversion successful but the ONNX model fails to execute with ONNX Runtime. "
    #         "Please create an issue "
    #         f"in the PyTorch GitHub repository against the {_BLUE}*onnx*{_END} component and "
    #         "attach the full error stack as well as reproduction scripts. "
    #     ) from e

    # Step 4: (When error report is requested) Validate the output values
    # TODO

    if profile:
        assert profile_result is not None
        _reporting.crete_onnx_export_profile_report(
            artifacts_dir / f"onnx_export_{timestamp}_profile.md",
            onnx_program.exported_program,
            profile_result,
            step=2,  # TODO: Update the step number to 4 when validation is implemented
        )

    return onnx_program
