// Copyright 2021 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include "iree-dialects-c/Dialects.h"
#include "iree-dialects-c/Utils.h"
#include "mlir-c/Bindings/Python/Interop.h"
#include "mlir-c/BuiltinAttributes.h"
#include "mlir-c/BuiltinTypes.h"
#include "mlir-c/Diagnostics.h"
#include "mlir-c/Registration.h"
#include "mlir/Bindings/Python/PybindAdaptors.h"

namespace py = pybind11;
using namespace mlir::python::adaptors;

namespace {

struct PyIREEPyDMSourceBundle {
  PyIREEPyDMSourceBundle(IREEPyDMSourceBundle wrapped) : wrapped(wrapped) {}
  PyIREEPyDMSourceBundle(PyIREEPyDMSourceBundle &&other)
      : wrapped(other.wrapped) {
    other.wrapped.ptr = nullptr;
  }
  PyIREEPyDMSourceBundle(const PyIREEPyDMSourceBundle &) = delete;
  ~PyIREEPyDMSourceBundle() {
    if (wrapped.ptr)
      ireePyDMSourceBundleDestroy(wrapped);
  }
  IREEPyDMSourceBundle wrapped;
};

struct PyIREEPyDMLoweringOptions {
  PyIREEPyDMLoweringOptions() : wrapped(ireePyDMLoweringOptionsCreate()) {}
  PyIREEPyDMLoweringOptions(PyIREEPyDMLoweringOptions &&other)
      : wrapped(other.wrapped) {
    other.wrapped.ptr = nullptr;
  }
  PyIREEPyDMLoweringOptions(const PyIREEPyDMLoweringOptions &) = delete;
  ~PyIREEPyDMLoweringOptions() {
    if (wrapped.ptr)
      ireePyDMLoweringOptionsDestroy(wrapped);
  }
  IREEPyDMLoweringOptions wrapped;
};

} // namespace

PYBIND11_MODULE(_ireeDialects, m) {
  m.doc() = "iree-dialects main python extension";

  auto irModule = py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"));
  auto typeClass = irModule.attr("Type");

  //===--------------------------------------------------------------------===//
  // Utils
  //===--------------------------------------------------------------------===//

  m.def(
      "lookup_nearest_symbol_from",
      [](MlirOperation fromOp, MlirAttribute symbol) {
        if (!mlirAttributeIsASymbolRef(symbol)) {
          throw std::invalid_argument("expected a SymbolRefAttr");
        }
        return ireeLookupNearestSymbolFrom(fromOp, symbol);
      },
      py::arg("fromOp"), py::arg("symbol"));

  // TODO: Upstream this into the main Python bindings.
  m.def(
      "emit_error",
      [](MlirLocation loc, std::string message) {
        mlirEmitError(loc, message.c_str());
      },
      py::arg("loc"), py::arg("message"));

  //===--------------------------------------------------------------------===//
  // IREEDialect
  //===--------------------------------------------------------------------===//
  auto iree_m = m.def_submodule("iree_input");
  iree_m.def(
      "register_dialect",
      [](MlirContext context, bool load) {
        MlirDialectHandle handle = mlirGetDialectHandle__iree_input__();
        mlirDialectHandleRegisterDialect(handle, context);
        if (load) {
          mlirDialectHandleLoadDialect(handle, context);
        }
      },
      py::arg("context") = py::none(), py::arg("load") = true);

  //===--------------------------------------------------------------------===//
  // IREELinalgExt
  //===--------------------------------------------------------------------===//
  auto iree_linalg_ext_m = m.def_submodule("iree_linalg_ext");
  iree_linalg_ext_m.def(
      "register_dialect",
      [](MlirContext context, bool load) {
        MlirDialectHandle handle = mlirGetDialectHandle__iree_linalg_ext__();
        mlirDialectHandleRegisterDialect(handle, context);
        if (load) {
          mlirDialectHandleLoadDialect(handle, context);
        }
      },
      py::arg("context") = py::none(), py::arg("load") = true);

  //===--------------------------------------------------------------------===//
  // LinalgTransform
  //===--------------------------------------------------------------------===//
  auto iree_linalg_transform_m = m.def_submodule("iree_linalg_transform");
  mlirIREELinalgTransformRegisterPasses();
  iree_linalg_transform_m.def(
      "register_dialect",
      [](MlirContext context, bool load) {
        MlirDialectHandle handle =
            mlirGetDialectHandle__iree_linalg_transform__();
        mlirDialectHandleRegisterDialect(handle, context);
        if (load) {
          mlirDialectHandleLoadDialect(handle, context);
        }
      },
      py::arg("context") = py::none(), py::arg("load") = true);

  //===--------------------------------------------------------------------===//
  // TransformDialect
  //===--------------------------------------------------------------------===//
  auto transform_m = m.def_submodule("transform");
  mlirIREETransformRegisterPasses();

  transform_m.def(
      "register_dialect",
      [](MlirContext context, bool load) {
        MlirDialectHandle handle = mlirGetDialectHandle__transform__();
        mlirDialectHandleRegisterDialect(handle, context);
        ireeRegisterTransformExtensions(context);
        if (load) {
          mlirDialectHandleLoadDialect(handle, context);
        }
      },
      py::arg("context") = py::none(), py::arg("load") = true);

  //===--------------------------------------------------------------------===//
  // IREEPyDMDialect
  //===--------------------------------------------------------------------===//
  auto iree_pydm_m = m.def_submodule("iree_pydm");
  mlirIREEPyDMRegisterPasses();

  py::class_<PyIREEPyDMSourceBundle>(
      iree_pydm_m, "SourceBundle", py::module_local(),
      "Contains raw assembly source or a reference to a file")
      .def_static(
          "from_asm",
          [](std::string asmBlob) {
            return PyIREEPyDMSourceBundle(ireePyDMSourceBundleCreateAsm(
                {asmBlob.data(), asmBlob.size()}));
          },
          py::arg("asm_blob"),
          "Creates a SourceBundle from an ASM blob (string or bytes)")
      .def_static(
          "from_file",
          [](std::string asmFile) {
            return PyIREEPyDMSourceBundle(ireePyDMSourceBundleCreateFile(
                {asmFile.data(), asmFile.size()}));
          },
          py::arg("asm_file"),
          "Creates a SourceBundle from a file containing ASM");
  py::class_<PyIREEPyDMLoweringOptions>(iree_pydm_m, "LoweringOptions",
                                        py::module_local(),
                                        "Lowering options to compile to IREE")
      .def(py::init<>())
      .def(
          "link_rtl",
          [](PyIREEPyDMLoweringOptions &self,
             PyIREEPyDMSourceBundle &sourceBundle) {
            ireePyDMLoweringOptionsLinkRtl(self.wrapped, sourceBundle.wrapped);
          },
          "Enables linking against a runtime-library module");

  iree_pydm_m.def(
      "register_dialect",
      [](MlirContext context, bool load) {
        MlirDialectHandle handle = mlirGetDialectHandle__iree_pydm__();
        mlirDialectHandleRegisterDialect(handle, context);
        if (load) {
          mlirDialectHandleLoadDialect(handle, context);
        }
      },
      py::arg("context") = py::none(), py::arg("load") = true);

  iree_pydm_m.def(
      "build_lower_to_iree_pass_pipeline",
      [](MlirPassManager passManager, PyIREEPyDMLoweringOptions &options) {
        MlirOpPassManager opPassManager =
            mlirPassManagerGetAsOpPassManager(passManager);
        mlirIREEPyDMBuildLowerToIREEPassPipeline(opPassManager,
                                                 options.wrapped);
      },
      py::arg("pass_manager"), py::arg("link_rtl_asm") = py::none());

  iree_pydm_m.def(
      "build_post_import_pass_pipeline",
      [](MlirPassManager passManager) {
        MlirOpPassManager opPassManager =
            mlirPassManagerGetAsOpPassManager(passManager);
        mlirIREEPyDMBuildPostImportPassPipeline(opPassManager);
      },
      py::arg("pass_manager"));

#define DEFINE_IREEPYDM_NULLARY_TYPE(Name)                                     \
  mlir_type_subclass(iree_pydm_m, #Name "Type", mlirTypeIsAIREEPyDM##Name,     \
                     typeClass)                                                \
      .def_classmethod(                                                        \
          "get",                                                               \
          [](py::object cls, MlirContext context) {                            \
            return cls(mlirIREEPyDM##Name##TypeGet(context));                  \
          },                                                                   \
          py::arg("cls"), py::arg("context") = py::none());

  DEFINE_IREEPYDM_NULLARY_TYPE(Bool)
  DEFINE_IREEPYDM_NULLARY_TYPE(Bytes)
  DEFINE_IREEPYDM_NULLARY_TYPE(ExceptionResult)
  DEFINE_IREEPYDM_NULLARY_TYPE(FreeVarRef)
  DEFINE_IREEPYDM_NULLARY_TYPE(List)
  DEFINE_IREEPYDM_NULLARY_TYPE(None)
  DEFINE_IREEPYDM_NULLARY_TYPE(Str)
  DEFINE_IREEPYDM_NULLARY_TYPE(Tuple)
  DEFINE_IREEPYDM_NULLARY_TYPE(Type)

  // IntegerType.
  mlir_type_subclass(iree_pydm_m, "IntegerType", mlirTypeIsAIREEPyDMInteger,
                     typeClass)
      .def_classmethod(
          "get",
          [](py::object cls, MlirContext context) {
            return cls(mlirIREEPyDMIntegerTypeGet(context));
          },
          py::arg("cls"), py::arg("context") = py::none())
      .def_classmethod(
          "get_explicit",
          [](py::object cls, int bitWidth, bool isSigned, MlirContext context) {
            return cls(mlirIREEPyDMIntegerTypeGetExplicit(context, bitWidth,
                                                          isSigned));
          },
          py::arg("cls"), py::arg("bit_width"), py::arg("is_signed") = true,
          py::arg("context") = py::none());

  // RealType.
  mlir_type_subclass(iree_pydm_m, "RealType", mlirTypeIsAIREEPyDMReal,
                     typeClass)
      .def_classmethod(
          "get",
          [](py::object cls, MlirContext context) {
            return cls(mlirIREEPyDMRealTypeGet(context));
          },
          py::arg("cls"), py::arg("context") = py::none())
      .def_classmethod(
          "get_explicit",
          [](py::object cls, MlirType fpType) {
            // TODO: Add a C-API for generically checking for FloatType.
            if (!mlirTypeIsAF32(fpType) && !mlirTypeIsAF64(fpType) &&
                !mlirTypeIsAF16(fpType) && !mlirTypeIsABF16(fpType)) {
              throw std::invalid_argument("expected a floating point type");
            }
            return cls(mlirIREEPyDMRealTypeGetExplicit(fpType));
          },
          py::arg("cls"), py::arg("fp_type"));

  // ObjectType.
  mlir_type_subclass(iree_pydm_m, "ObjectType", mlirTypeIsAIREEPyDMObject,
                     typeClass)
      .def_classmethod(
          "get",
          [](py::object cls, MlirContext context) {
            return cls(mlirIREEPyDMObjectTypeGet(context, {nullptr}));
          },
          py::arg("cls"), py::arg("context") = py::none())
      .def_classmethod(
          "get_typed",
          [](py::object cls, MlirType type) {
            if (!mlirTypeIsAIREEPyDMPrimitiveType(type)) {
              throw std::invalid_argument(
                  "expected a primitive type when constructing object");
            }
            MlirContext context = mlirTypeGetContext(type);
            return cls(mlirIREEPyDMObjectTypeGet(context, type));
          },
          py::arg("cls"), py::arg("type"));
}