# --------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------

from __future__ import annotations

from pathlib import Path
from textwrap import dedent
from typing import Annotated, Any, Iterable, Optional, Set, TextIO

import pygen as cg
from onnx.defs import (
    AttributeProto,
    OpSchema,
    get_all_schemas_with_history,
    onnx_opset_version,
)
from onnx.helper import get_attribute_value

__all__ = [
    "OpsetId",
    "parse_opsetid",
    "QualOpName",
    "UnsupportedOpError",
    "OpsetsBuilder",
    "parse_attr_type",
    "parse_input_output_type",
]

MODULE_ONNX = "onnx"
MODULE_ONNX_DEFS = "onnx.defs"
MODULE_ONNX_SCRIPT_TYPES = "onnxscript.onnx_types"
MODULE_ONNX_SCRIPT_VALUES = "onnxscript.values"


OpsetId = tuple[Annotated[str, "domain"], Annotated[int, "version"]]


def parse_opsetid(opsetid: str) -> OpsetId:
    split_index = opsetid.rfind("/")
    version = int(opsetid[split_index + 1 :])
    return ("" if split_index < 0 else opsetid[:split_index], version)


def format_opsetid(opsetid: OpsetId) -> str:
    domain, version = opsetid
    return f"{domain}/{version}" if domain else str(version)


class QualOpName:
    def __init__(self, domain: str, name: str, version: int):
        self.domain = domain
        self.name = name
        self.version = version
        self.docuri = (
            "https://onnx.ai/onnx/operators/onnx_"
            f"{domain.replace('.', '')}_{name}.html#{name.lower()}-{version}"
        )

    def __repr__(self) -> str:
        return (
            f"QualOpName(domain={self.domain!r}, version={self.version!r}, name={self.name!r})"
        )

    def __str__(self) -> str:
        domain_prefix = f"{self.domain}::" if self.domain else ""
        return f"{domain_prefix}{self.name}({self.version})"


class OpsetBaseTypeRef(cg.TypeRef):
    def __init__(self):
        super().__init__(MODULE_ONNX_SCRIPT_VALUES, "Opset")


class TensorTypeRef(cg.TypeRef):
    def __init__(self):
        super().__init__(MODULE_ONNX_SCRIPT_TYPES, "Tensor")


class ConstraintTypeRef(cg.TypeRef):
    def __init__(self, name: str, *typeargs: cg.TypeRef, is_output: bool):
        super().__init__(None, name, *typeargs)
        self.is_output = is_output


class UnsupportedOpError(NotImplementedError):
    def __init__(self, op: QualOpName, message: str):
        super().__init__(self, message)
        self.op = op
        self.message = message


def _make_suffix(str: str) -> str:
    return f"_{str.replace('.', '_')}" if str else ""


def _make_class_name(domain: str, version: int) -> str:
    return f"Opset{_make_suffix(domain)}{version}"


def _make_module_name(base_name: str, domain: str, version: int) -> str:
    return f"{base_name}._impl.opset{_make_suffix(domain)}{version}"


class OpsetModule(cg.Module):
    def __init__(self, base_name: str, domain: str, version: int, *members: cg.Stmt):
        self.domain = domain
        self.version = version
        super().__init__(*members, name=_make_module_name(base_name, domain, version))


class OpsetsBuilder:
    class Result:
        def __init__(self):
            self.all_ops_count: int = 0
            self.all_modules: list[cg.Module] = []
            self.unsupported_ops: dict[str, list[UnsupportedOpError]] = {}
            self.included_opsets: set[OpsetId] = set()
            self.excluded_opsets: set[OpsetId] = set()

        def write(self, base_path: Path) -> list[Path]:
            return sorted(
                [self._write_module(base_path, module) for module in self.all_modules]
            )

        def _write_module(self, base_path: Path, module: cg.Module) -> Path:
            qual_name = module.name.split(".")
            base_path = base_path.joinpath(*qual_name[:-1])
            base_path.mkdir(parents=True, exist_ok=True)
            path = base_path.joinpath(qual_name[-1] + ".py")
            with open(path, "w", encoding="utf-8") as writer:
                self._write_header(writer)
                module.accept(cg.PythonWriter(writer))
            return path

        def _write_header(self, writer: TextIO):
            dashline = f"# {'-' * 74}\n"
            writer.write(dashline)
            writer.write("# ⚠️ WARNING - AUTO-GENERATED CODE - DO NOT EDIT ⚠️ \n")
            writer.write("# ⚙️ Generated by 'python -m opgen'\n")
            writer.write(dashline)
            writer.write("# Copyright (c) Microsoft Corporation. ")
            writer.write("All rights reserved.\n")
            writer.write("# Licensed under the MIT License.\n")
            writer.write(dashline)
            writer.write("# pylint: disable=W0221,W0222,R0901,W0237\n")
            writer.write("# mypy: disable-error-code=override\n")
            writer.write("# ruff: noqa: N801,E741\n")
            writer.write("# ruff: noqa: D214,D402,D405,D411,D412,D416,D417\n")
            writer.write(dashline)
            writer.write("\n")
            writer.write("from __future__ import annotations\n")

    def __init__(
        self,
        *,
        module_base_name: str,
        min_default_opset_version: int,
        include_opsets: Optional[Set[OpsetId]] = None,
        exclude_opsets: Optional[Set[OpsetId]] = None,
    ):
        self.module_base_name = module_base_name
        self.min_default_opset_version = min_default_opset_version
        self.include_opsets = include_opsets
        self.exclude_opsets = exclude_opsets

    def build(self) -> OpsetsBuilder.Result:
        self._result = OpsetsBuilder.Result()
        self._make_opset_modules()
        self._make_init_module()
        self._make_imports()
        return self._result

    def _log_unsupported(self, error: UnsupportedOpError):
        self._result.unsupported_ops.setdefault(error.message, []).append(error)

    def _make_opset_module(self, domain: str, version: int) -> Optional[OpsetModule]:
        opsetid = (domain, version)

        if (self.include_opsets and opsetid not in self.include_opsets) or (
            self.exclude_opsets and opsetid in self.exclude_opsets
        ):
            self._result.excluded_opsets.add(opsetid)
            return None

        self._result.included_opsets.add(opsetid)

        if version > 1:
            base_type = cg.TypeRef(
                _make_module_name(self.module_base_name, domain, version - 1),
                _make_class_name(domain, version - 1),
            )
        else:
            base_type = OpsetBaseTypeRef()

        opset = OpsetModule(
            self.module_base_name,
            domain,
            version,
            cg.ImportFrom("typing", cg.Alias("TypeVar")),
            cg.ClassDef(
                _make_class_name(domain, version),
                cg.FunctionDef(
                    "__new__",
                    cg.Arg("cls"),
                    body=cg.ThunkStmt(f"return Opset.__new__(cls, {domain!r}, {version!r})"),
                ),
                bases=[base_type],
            ),
        )

        self._result.all_modules.append(opset)
        return opset

    def _make_opset_modules(self):
        domains = {}
        schemas: list[OpSchema] = sorted(
            get_all_schemas_with_history(),
            key=lambda op: (op.domain, op.since_version, op.name),
        )

        for schema in schemas:
            qualname = QualOpName(schema.domain, schema.name, schema.since_version)
            domain: str = schema.domain
            version: int = schema.since_version
            domain_opsets = domains.setdefault(domain, {})

            if schema.deprecated:
                self._log_unsupported(UnsupportedOpError(qualname, "deprecated"))
                continue

            if version in domain_opsets:
                opset = domain_opsets[version]
            else:
                if opset := self._make_opset_module(domain, version):
                    domain_opsets[version] = opset
                else:
                    continue

            try:
                opset_class = cg.first_or_none(opset.get_children_of_type(cg.ClassDef))
                if opset_class is not None:
                    for stmt in self._make_function(qualname, schema):
                        opset_class.append_body(stmt)
                    self._result.all_ops_count += 1
            except NotImplementedError as error:
                if not isinstance(error, UnsupportedOpError):
                    error = UnsupportedOpError(qualname, str(error))
                self._log_unsupported(error)

        if onnx_opset_version() not in domains[""]:
            self._make_opset_module("", onnx_opset_version())

        for module in self._result.all_modules:
            module.accept(cg.DocCommentBuilder())

        self._result.all_modules.sort(key=lambda m: (m.domain, m.version, m.name))

    def _make_init_module(self):
        all_list = cg.ListExpr(cg.Constant("all_opsets"))
        init_module = cg.Module(
            cg.ImportFrom(MODULE_ONNX_DEFS, cg.Alias("onnx_opset_version")),
            cg.Assign(cg.Name("__all__"), all_list),
            cg.If(
                cg.BinOp(
                    cg.Call(cg.Name("onnx_opset_version")),
                    "<",
                    cg.Constant(self.min_default_opset_version),
                ),
                cg.Raise(
                    cg.Call(
                        cg.Name("ImportError"),
                        cg.ThunkExpr(
                            'f"ONNX Script requires ONNX opset >= '
                            f"{self.min_default_opset_version} "
                            'but {onnx_opset_version()} is detected."'
                        ),
                    )
                ),
            ),
            name=f"{self.module_base_name}.__init__",
        )

        all_opsets = cg.DictExpr()
        for opset_module in filter(
            lambda m: isinstance(m, OpsetModule), self._result.all_modules
        ):
            opset_module: OpsetModule
            opset_class = cg.first_or_none(opset_module.get_children_of_type(cg.ClassDef))
            if opset_class is not None:
                opset_export_name = opset_module.name.split(".")[-1]
                all_opsets.append_element(
                    cg.DictElem(
                        cg.TupleExpr(
                            cg.Constant(opset_module.domain), cg.Constant(opset_module.version)
                        ),
                        cg.Name(opset_export_name),
                    )
                )
                all_list.append_child(
                    cg.Constant(opset_export_name), cg.ListExpr.Roles.Elements
                )
                init_module.append_body(
                    cg.Assign(cg.Name(opset_export_name), cg.Call(opset_class.make_typeref()))
                )
        all_opsets_type = cg.TypeRef.make_composite_if_multiple(
            cg.TypingRefs.Mapping,
            cg.TypeRef.make_composite_if_multiple(
                cg.TypingRefs.Tuple, cg.StrTypeRef(), cg.IntTypeRef()
            ),
            cg.TypeRef(MODULE_ONNX_SCRIPT_VALUES, "Opset"),
        )
        init_module.append_body(cg.Assign(cg.Name("all_opsets"), all_opsets, all_opsets_type))

        self._result.all_modules.append(init_module)

    def _make_imports(self):
        for module in self._result.all_modules:
            if isinstance(module, OpsetModule):
                module.prepend_child(
                    cg.ImportFrom(MODULE_ONNX_DEFS, cg.Alias("get_schema")),
                    cg.Module.Roles.Body,
                )
                module.prepend_child(
                    cg.ImportFrom(MODULE_ONNX_SCRIPT_VALUES, cg.Alias("Op, Opset")),
                    cg.Module.Roles.Body,
                )
            module.accept(cg.ImportAdjuster())

    def _make_function_type_constraints(
        self, schema: OpSchema, typerefs: Iterable[cg.TypeRef]
    ) -> Iterable[cg.Assign]:
        input_constraints: dict[str, list[cg.TypeRef]] = {}
        output_constraints: dict[str, list[cg.TypeRef]] = {}

        def constraint_is_compatible(
            constraint_name: str, constraint_types: list[cg.TypeRef]
        ) -> bool:
            """Ensure that if we have already recoreded ``constraint_name`` in either
            ``input_constraints`` or ``output_constraints`` that the constrained types
            match by comparing the generated code that results from the types.
            """
            for existing_constraints in input_constraints, output_constraints:
                if (existing := existing_constraints.get(constraint_name, None)) is not None:
                    if len(existing) != len(constraint_types):
                        return False  # differing number of constraints, can't be compatible
                    for a, b in zip(existing, constraint_types):
                        if str(a) != str(b):
                            return False  # a constrained type does not match
            return True

        def collect_constraints(typerefs: Iterable[cg.TypeRef]):
            """Collect and validate all ConstraintTypeRef types recursively. Once a
            constraint has been identified, it becomes the 'canonical' definition of
            all constraints that will be collected with the same name.
            """
            for typeref in typerefs:
                if isinstance(typeref, ConstraintTypeRef):
                    typeref.name = f"{typeref.name}_{schema.name}"
                    constraint_types = [ta.remove() for ta in typeref.typeargs]

                    assert constraint_is_compatible(typeref.name, constraint_types), (
                        "This should not happen since constraints are "
                        "expected to be correct at the OpSchema spec level"
                    )

                    (output_constraints if typeref.is_output else input_constraints)[
                        typeref.name
                    ] = constraint_types
                else:
                    collect_constraints(typeref.typeargs)

        collect_constraints(typerefs)

        def make_type_alias(name: str, aliased_type: cg.TypeRef):
            return cg.Assign(
                cg.Name(name),
                aliased_type,
                cg.TypeRef("typing_extensions", "TypeAlias"),
            )

        # Pass 1: process input constraints first; if a constraint is bound to
        # a single type, generate a TypeAlias, otherwise a TypeVar.
        for constraint_name, constraint_typerefs in input_constraints.items():
            if len(constraint_typerefs) == 1:
                yield make_type_alias(constraint_name, constraint_typerefs[0])
            else:
                yield cg.Assign(
                    cg.Name(constraint_name),
                    cg.Call(
                        cg.Name("TypeVar"),
                        cg.Constant(constraint_name),
                        *constraint_typerefs,
                    ),
                )

        # Pass 2: process output constraints, only yielding constraints as
        # TypeAliases when an output constraint is not also used as an
        # input constraint. See https://github.com/microsoft/onnxscript/pull/778
        # for details. When we can rely on Python 3.12, we can have output-only
        # constraints as TypeVars, e.g. when the following syntax is allowed:
        #    def func[T]() -> T: ...
        for constraint_name, constraint_typerefs in output_constraints.items():
            if constraint_name not in input_constraints:
                yield make_type_alias(
                    constraint_name,
                    cg.TypeRef.make_composite_if_multiple(
                        cg.TypingRefs.Union,
                        *constraint_typerefs,
                    ),
                )

    def _make_function(self, qualname: QualOpName, schema: OpSchema) -> Iterable[cg.Stmt]:
        attr_args = self._make_function_attr_args(schema)
        input_args = self._make_function_input_args(schema)
        args = [cg.Arg("self"), *input_args, *attr_args]
        output_types = [
            self._make_input_output_type(output, is_output=True) for output in schema.outputs
        ]

        yield from self._make_function_type_constraints(
            schema,
            [arg.type for arg in input_args] + output_types,
        )

        op_inputs: list[cg.Expr] = []
        op_attrs: list[cg.Expr] = []

        for arg in args:
            if arg.name in {"self", "*"}:
                continue
            if arg.is_vararg:
                op_inputs.append(cg.Starred(cg.Name(arg.name)))
            elif arg.is_kwarg:
                op_attrs.append(cg.Assign(cg.Name(arg.name), cg.Name(arg.name)))
            else:
                op_inputs.append(cg.Name(arg.name))

        if len(op_inputs) > 0:
            op_call = cg.Call(
                cg.Name("op"),
                cg.Starred(
                    cg.Call(cg.Name("self._prepare_inputs"), cg.Name("schema"), *op_inputs)
                ),
                *op_attrs,
            )
        else:
            op_call = cg.Call(cg.Name("op"), *op_attrs)

        doc = f'[🌐 {qualname}]({qualname.docuri} "Online Documentation")\n\n{schema.doc}'

        yield cg.FunctionDef(
            qualname.name,
            *args,
            return_type=cg.TypeRef.make_composite_if_multiple(
                cg.TypingRefs.Tuple,
                *output_types,
            ),
            doc=_process_documentation(doc),
            body=[
                cg.Assign(
                    cg.Name("schema"),
                    cg.Call(
                        cg.Name("get_schema"),
                        cg.Constant(qualname.name),
                        cg.Constant(qualname.version),
                        cg.Constant(qualname.domain),
                    ),
                ),
                cg.Assign(
                    cg.Name("op"),
                    cg.Call(
                        cg.Name("Op"),
                        cg.Name("self"),
                        cg.Constant(qualname.name),
                        cg.Name("schema"),
                    ),
                ),
                cg.Return(op_call),
            ],
        )

    def _make_input_arg_name(self, input_name: str, schema: OpSchema):
        """ONNX allows for an op to have an input and an attribute with the same name.
        Attribute names have contextual meaning however, so detect this case and disambiguate
        the input name. See Split(1) for the only offending OpSchema as of opset 18.
        """
        for attr in schema.attributes.values():
            if attr.name == input_name:
                return f"{input_name}_"
        return input_name

    def _make_function_input_args(self, schema: OpSchema) -> Iterable[cg.Arg]:
        args: list[cg.Arg] = []
        for input in schema.inputs:
            optional = input.option == OpSchema.FormalParameterOption.Optional
            variadic = input.option == OpSchema.FormalParameterOption.Variadic
            heterogeneous = not input.is_homogeneous
            differentiable = (
                input.differentiation_category
                == OpSchema.DifferentiationCategory.Differentiable
            )
            non_differentiable = (
                input.differentiation_category
                == OpSchema.DifferentiationCategory.NonDifferentiable
            )

            doctags = []
            if optional:
                doctags.append("optional")
            elif variadic:
                # If we encounter a variadic input, previous inputs cannot have default
                # values as this allows for ambiguity at the call site.
                #
                # Specifically this avoids pylint W1113 (keyword-arg-before-vararg):
                #   https://pylint.pycqa.org/en/latest/user_guide/messages/warning/keyword-arg-before-vararg.html
                #
                # As of 2023-06-08, only Loop(1, 11, 13, 16) and Scan(8) exhibit this issue.
                # Ex:
                #
                #     def Loop(
                #         self,
                #         M: Optional[I] = None,        <- ambiguity with *v_initial
                #         cond: Optional[B] = None,     <- ambiguity with *v_initial
                #         *v_initial: V,
                #         body: Optional[GraphProto] = None,
                #     ) -> V: ...
                #
                for prev_arg in args:
                    prev_arg.default_value = None
                doctags.append("variadic")
            if heterogeneous:
                doctags.append("heterogeneous")
            if differentiable:
                doctags.append("differentiable")
            elif non_differentiable:
                doctags.append("non-differentiable")

            doc = input.description.strip()
            if len(doctags) > 0:
                doc = f"({', '.join(doctags)}) {doc}"

            type = self._make_input_output_type(input, is_output=False)
            if optional and not isinstance(type, cg.TypingRefs.Optional):
                type = cg.TypingRefs.Optional(type)

            args.append(
                cg.Arg(
                    self._make_input_arg_name(input.name, schema),
                    type=type,
                    doc=_process_documentation(doc),
                    is_vararg=variadic,
                    default_value=cg.Constant(None) if optional else None,
                )
            )

        return args

    def _make_function_attr_args(self, schema: OpSchema) -> Iterable[cg.Arg]:
        generate_kwonly_sentinel = True
        for attr in schema.attributes.values():
            attr_type = parse_attr_type(attr.type)
            default_value = None

            if attr.default_value.name:
                default_value = get_attribute_value(attr.default_value)

                def fmt(value: Any) -> str:
                    if isinstance(value, (bytes, bytearray)):
                        return str(value.decode("utf-8"))
                    return value

                if isinstance(default_value, list):
                    default_value = tuple(fmt(val) for val in default_value)
                else:
                    default_value = fmt(default_value)
            else:
                default_value = None

            if not attr.required and default_value is None:
                attr_type = cg.TypingRefs.Optional(attr_type)

            if generate_kwonly_sentinel and not any(
                i.option == OpSchema.FormalParameterOption.Variadic for i in schema.inputs
            ):
                generate_kwonly_sentinel = False
                yield cg.Arg("*")

            yield cg.Arg(
                attr.name,
                type=attr_type,
                default_value=None if attr.required else cg.Constant(default_value),
                doc=attr.description,
                is_kwarg=True,
            )

    def _make_input_output_type(
        self,
        parameter: OpSchema.FormalParameter,
        is_output: bool,
    ) -> cg.TypeRef:
        py_types = [parse_input_output_type(type) for type in sorted(parameter.types)]
        try:
            # input.type_str will either be a valid ONNX type (e.g. 'tensor(int)')
            # or the name of a type constraint. If it parses as a type, it's not
            # constrained; otherwise bind the underlying types to the constraint.
            parse_input_output_type(parameter.type_str)
            return cg.TypeRef.make_composite_if_multiple(cg.TypingRefs.Union, *py_types)
        except NotImplementedError:
            return ConstraintTypeRef(parameter.type_str, *py_types, is_output=is_output)


def parse_input_output_type(onnx_type: str) -> cg.TypeRef:
    def error(message: Optional[str] = None):
        return NotImplementedError(
            f"input/output type not implemented: {onnx_type!r}"
            + (f" ({message!r})" if message else "")
        )

    default_value_map = {
        "BOOL": bool(),  # noqa: UP018
        "FLOAT": float(),  # noqa: UP018
        "FLOAT16": float(),  # noqa: UP018
        "BFLOAT16": float(),  # noqa: UP018
        "DOUBLE": float(),  # noqa: UP018
        "INT8": int(),  # noqa: UP018
        "INT16": int(),  # noqa: UP018
        "INT32": int(),  # noqa: UP018
        "INT64": int(),  # noqa: UP018
        "UINT8": int(),  # noqa: UP018
        "UINT16": int(),  # noqa: UP018
        "UINT32": int(),  # noqa: UP018
        "UINT64": int(),  # noqa: UP018
        "COMPLEX64": complex(),
        "COMPLEX128": complex(),
    }

    id = ""
    stack: list[cg.TypeRef] = []
    for c in onnx_type:
        if c == "(":
            if id == "tensor":
                type = TensorTypeRef()
            elif id == "seq":
                type = cg.TypingRefs.Sequence()
            elif id == "map":
                type = cg.TypingRefs.Mapping()
            elif id == "optional":
                type = cg.TypingRefs.Optional()
            else:
                raise error(id)
            if len(stack) > 0:
                stack[-1].append_typearg(type)
            stack.append(type)
            id = ""
        elif c in (")", ","):
            type = stack.pop() if c == ")" else stack[-1]
            if isinstance(type, TensorTypeRef):
                type.name = id.upper()
                type.default_value = cg.Constant(default_value_map.get(type.name))
            elif id and isinstance(type, cg.TypingRefs.Mapping):
                if id == "int64":
                    type.append_typearg(cg.IntTypeRef())
                elif id == "string":
                    type.append_typearg(cg.StrTypeRef())
                else:
                    raise error(id)
            elif id:
                break
            id = ""
            if len(stack) == 0:
                return type
        else:
            id += c
    raise error()


def parse_attr_type(type) -> cg.TypeRef:
    if type == AttributeProto.FLOAT:
        return cg.FloatTypeRef()
    if type == AttributeProto.INT:
        return cg.IntTypeRef()
    if type == AttributeProto.STRING:
        return cg.StrTypeRef()
    if type == AttributeProto.TENSOR:
        return cg.TypeRef(MODULE_ONNX, "TensorProto")
    if type == AttributeProto.SPARSE_TENSOR:
        return cg.TypeRef(MODULE_ONNX, "SparseTensorProto")
    if type == AttributeProto.GRAPH:
        return cg.TypeRef(MODULE_ONNX, "GraphProto")
    if type == AttributeProto.TYPE_PROTO:
        return cg.TypeRef(MODULE_ONNX, "TypeProto")
    if type == AttributeProto.FLOATS:
        return cg.TypingRefs.Sequence(cg.FloatTypeRef())
    if type == AttributeProto.INTS:
        return cg.TypingRefs.Sequence(cg.IntTypeRef())
    if type == AttributeProto.STRINGS:
        return cg.TypingRefs.Sequence(cg.StrTypeRef())
    if type == AttributeProto.TENSORS:
        return cg.TypingRefs.Sequence(cg.TypeRef(MODULE_ONNX, "TensorProto"))
    if type == AttributeProto.SPARSE_TENSORS:
        return cg.TypingRefs.Sequence(cg.TypeRef(MODULE_ONNX, "SparseTensorProto"))
    if type == AttributeProto.GRAPHS:
        return cg.TypingRefs.Sequence(cg.TypeRef(MODULE_ONNX, "GraphProto"))
    if type == AttributeProto.TYPE_PROTOS:
        return cg.TypingRefs.Sequence(cg.TypeRef(MODULE_ONNX, "TypeProto"))
    raise NotImplementedError(f"attribute type not implemented: {type}")


def _process_documentation(doc: str):
    # Lifted from ONNX's docsgen:
    # https://github.com/onnx/onnx/blob/3fd41d249bb8006935aa0031a332dd945e61b7e5/docs/docsgen/source/onnx_sphinx.py#L414
    doc = dedent(doc or "")
    main_docs_url = "https://github.com/onnx/onnx/blob/master/"
    rep = {
        "[the doc](IR.md)": "`ONNX <{0}docs/IR.md>`_",
        "[the doc](Broadcasting.md)": "`Broadcasting in ONNX <{0}docs/Broadcasting.md>`_",
        "<dl>": "",
        "</dl>": "",
        "<dt>": "* ",
        "<dd>": "  ",
        "</dt>": "",
        "</dd>": "",
        "<tt>": "``",
        "</tt>": "``",
        "<br>": "\n",
    }
    for k, v in rep.items():
        doc = doc.replace(k, v.format(main_docs_url))
    move = 0
    lines = []
    for line in doc.split("\n"):
        if line.startswith("```"):
            if move > 0:
                move -= 4
                lines.append("\n")
            else:
                lines.append("::\n")
                move += 4
        elif move > 0:
            lines.append(" " * move + line)
        else:
            lines.append(line)
    return "\n".join(lines)