// Copyright 2023 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 "./io.h"

#include <iostream>
#include <string_view>
#include <unordered_map>

#include "./buffer_interop.h"
#include "./vm.h"
#include "iree/base/internal/file_io.h"
#include "iree/base/internal/path.h"
#include "iree/io/formats/gguf/gguf_parser.h"
#include "iree/io/formats/irpa/irpa_builder.h"
#include "iree/io/formats/irpa/irpa_parser.h"
#include "iree/io/formats/safetensors/safetensors_parser.h"
#include "iree/io/parameter_index_provider.h"
#include "iree/modules/io/parameters/module.h"
#include "iree/schemas/parameter_archive.h"

namespace iree::python {

namespace {

VmModule CreateIoParametersModule(VmInstance &instance, py::args providers) {
  iree_vm_module_t *module = nullptr;
  std::vector<iree_io_parameter_provider_t *> c_providers;
  iree_host_size_t size = providers.size();
  c_providers.resize(size);
  for (iree_host_size_t i = 0; i < size; ++i) {
    ParameterProvider *provider = py::cast<ParameterProvider *>(providers[i]);
    c_providers[i] = provider->raw_ptr();
  }
  CheckApiStatus(iree_io_parameters_module_create(
                     instance.raw_ptr(), size, c_providers.data(),
                     iree_allocator_system(), &module),
                 "Error creating io_parameters module");
  return VmModule::StealFromRawPtr(module);
}

FileHandle FileHandleWrapMemory(py::object host_buffer, bool readable,
                                bool writable, size_t &out_buffer_size) {
  struct Retained {
    Retained(py::object host_buffer)
        : buffer_request(host_buffer, PyBUF_SIMPLE),
          host_buffer(std::move(host_buffer)) {}
    PyBufferRequest buffer_request;
    py::object host_buffer;
  };
  std::unique_ptr<Retained> outer_retained =
      std::make_unique<Retained>(std::move(host_buffer));
  iree_io_file_access_t access = 0;
  if (readable) access |= IREE_IO_FILE_ACCESS_READ;
  if (writable) access |= IREE_IO_FILE_ACCESS_WRITE;
  iree_io_file_handle_t *created_handle;
  out_buffer_size = outer_retained->buffer_request.view().len;
  CheckApiStatus(
      iree_io_file_handle_wrap_host_allocation(
          access,
          iree_byte_span_t{
              static_cast<uint8_t *>(outer_retained->buffer_request.view().buf),
              static_cast<iree_host_size_t>(
                  outer_retained->buffer_request.view().len)},
          iree_io_file_handle_release_callback_t{
              +[](void *user_data, iree_io_file_handle_primitive_t primitive) {
                Retained *inner_retained = static_cast<Retained *>(user_data);
                delete inner_retained;
              },
              (void *)outer_retained.get(),
          },
          iree_allocator_system(), &created_handle),
      "Could not wrap host memory into a file handle");
  outer_retained.release();
  return FileHandle::StealFromRawPtr(created_handle);
}

void ParameterIndexAddFromFileHandle(ParameterIndex &self, std::string &key,
                                     FileHandle &file_handle, uint64_t length,
                                     uint64_t offset,
                                     std::optional<std::string> metadata) {
  iree_io_parameter_index_entry_t entry;
  memset(&entry, 0, sizeof(entry));
  entry.key = iree_make_string_view(key.data(), key.size());
  if (metadata) {
    entry.metadata.data = reinterpret_cast<const uint8_t *>(metadata->data());
    entry.metadata.data_length = metadata->size();
  }
  entry.length = length;
  entry.type = IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE;
  entry.storage.file.handle = file_handle.raw_ptr();
  entry.storage.file.offset = offset;
  CheckApiStatus(iree_io_parameter_index_add(self.raw_ptr(), &entry),
                 "Could not add parameter index entry");
}

void ParameterIndexParseFileHandle(ParameterIndex &self,
                                   FileHandle &file_handle,
                                   std::string &format) {
  if (format == "gguf") {
    CheckApiStatus(
        iree_io_parse_gguf_index(file_handle.raw_ptr(), self.raw_ptr()),
        "Could not parse gguf file into index");
  } else if (format == "irpa") {
    CheckApiStatus(
        iree_io_parse_irpa_index(file_handle.raw_ptr(), self.raw_ptr()),
        "Could not parse IREE parameter archive file into index");
  } else if (format == "safetensors") {
    CheckApiStatus(
        iree_io_parse_safetensors_index(file_handle.raw_ptr(), self.raw_ptr()),
        "Could not parse safetensors file into index");
  } else {
    throw std::invalid_argument(
        "Unrecognized file format. Expected one of: 'gguf', 'irpa', "
        "'safetensors'");
  }
}

void ParameterIndexLoadFile(ParameterIndex &self, std::string &file_path,
                            std::optional<std::string> format, bool readable,
                            bool writable, bool mmap) {
  // Default format from extension.
  if (!format) {
    iree_string_view_t path_ext = iree_file_path_extension(
        iree_make_string_view(file_path.data(), file_path.size()));
    format.emplace(path_ext.data, path_ext.size);
  }

  // Open file.
  iree_file_read_flags_t read_flags = IREE_FILE_READ_FLAG_DEFAULT;
  if (mmap) {
    read_flags = IREE_FILE_READ_FLAG_MMAP;
  } else {
    read_flags = IREE_FILE_READ_FLAG_PRELOAD;
  }
  iree_file_contents_t *file_contents = nullptr;
  CheckApiStatus(
      iree_file_read_contents(file_path.c_str(), read_flags,
                              iree_allocator_system(), &file_contents),
      "Error opening parameter file");
  iree_io_file_handle_release_callback_t release_callback = {
      +[](void *user_data, iree_io_file_handle_primitive_t handle_primitive) {
        iree_file_contents_t *file_contents = (iree_file_contents_t *)user_data;
        iree_file_contents_free(file_contents);
      },
      file_contents,
  };

  // Wrap contents.
  iree_io_file_handle_t *raw_file_handle = nullptr;
  iree_status_t status = iree_io_file_handle_wrap_host_allocation(
      IREE_IO_FILE_ACCESS_READ, file_contents->buffer, release_callback,
      iree_allocator_system(), &raw_file_handle);
  if (!iree_status_is_ok(status)) {
    iree_file_contents_free(file_contents);
    CheckApiStatus(status, "Error accessing parameter memory");
  }

  // Parse.
  FileHandle file_handle = FileHandle::StealFromRawPtr(raw_file_handle);
  ParameterIndexParseFileHandle(self, file_handle, *format);
}

}  // namespace

void SetupIoBindings(py::module_ &m) {
  m.def("create_io_parameters_module", &CreateIoParametersModule);

  py::class_<FileHandle>(m, "FileHandle")
      .def_static(
          "wrap_memory",
          [](py::object host_buffer, bool readable, bool writable) {
            size_t unused_len;
            return FileHandleWrapMemory(std::move(host_buffer), readable,
                                        writable, unused_len);
          },
          py::arg("host_buffer"), py::arg("readable") = true,
          py::arg("writable") = false);
  py::class_<ParameterProvider>(m, "ParameterProvider");
  py::class_<ParameterIndex>(m, "ParameterIndex")
      .def("__init__",
           [](ParameterIndex *new_self) {
             iree_io_parameter_index_t *created;
             CheckApiStatus(iree_io_parameter_index_create(
                                iree_allocator_system(), &created),
                            "Could not create IO parameter index");
             new (new_self) ParameterIndex();
             *new_self = ParameterIndex::StealFromRawPtr(created);
           })
      .def("__len__",
           [](ParameterIndex &self) {
             return iree_io_parameter_index_count(self.raw_ptr());
           })
      .def(
          "reserve",
          [](ParameterIndex &self, iree_host_size_t new_capacity) {
            CheckApiStatus(
                iree_io_parameter_index_reserve(self.raw_ptr(), new_capacity),
                "Could not reserve capacity");
          },
          py::arg("new_capacity"))
      .def(
          "add_splat",
          [](ParameterIndex &self, std::string key, py::object pattern,
             uint64_t total_length, std::optional<std::string> metadata) {
            iree_io_parameter_index_entry_t entry;
            memset(&entry, 0, sizeof(entry));
            entry.key = iree_make_string_view(key.data(), key.size());
            if (metadata) {
              entry.metadata.data =
                  reinterpret_cast<const uint8_t *>(metadata->data());
              entry.metadata.data_length = metadata->size();
            }
            entry.length = total_length;
            entry.type = IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_SPLAT;
            PyBufferRequest pattern_info(pattern, PyBUF_SIMPLE);
            auto pattern_size = pattern_info.view().len;
            if (pattern_size > sizeof(entry.storage.splat.pattern)) {
              throw std::invalid_argument(
                  "pattern must be limited to 16 bytes");
            }
            entry.storage.splat.pattern_length = pattern_size;
            std::memcpy(entry.storage.splat.pattern, pattern_info.view().buf,
                        pattern_size);
            CheckApiStatus(iree_io_parameter_index_add(self.raw_ptr(), &entry),
                           "Could not add parameter index entry");
          },
          py::arg("key"), py::arg("pattern"), py::arg("total_length"),
          py::arg("metadata") = py::none())
      .def("add_from_file_handle", ParameterIndexAddFromFileHandle,
           py::arg("key"), py::arg("file_handle"), py::arg("length"),
           py::arg("offset") = 0, py::arg("metadata") = py::none())
      .def(
          "add_buffer",
          [](ParameterIndex &self, std::string key, py::object buffer,
             bool readable, bool writable,
             std::optional<std::string> metadata) {
            size_t buffer_size;
            FileHandle file_handle = FileHandleWrapMemory(
                std::move(buffer), readable, writable, buffer_size);
            ParameterIndexAddFromFileHandle(self, key, file_handle, buffer_size,
                                            /*offset=*/0, std::move(metadata));
          },
          py::arg("key"), py::arg("buffer"), py::arg("readable") = true,
          py::arg("writable") = false, py::arg("metadata") = py::none())
      .def("load_from_file_handle", ParameterIndexParseFileHandle,
           py::arg("file_handle"), py::arg("format"))
      .def("load", ParameterIndexLoadFile, py::arg("file_path"),
           py::arg("format") = py::none(), py::arg("readable") = true,
           py::arg("writable") = false, py::arg("mmap") = true)
      .def(
          "create_provider",
          [](ParameterIndex &self, std::string scope,
             std::optional<iree_host_size_t> max_concurrent_operations) {
            if (!max_concurrent_operations) {
              max_concurrent_operations =
                  IREE_IO_PARAMETER_INDEX_PROVIDER_DEFAULT_MAX_CONCURRENT_OPERATIONS;
            }
            iree_io_parameter_provider_t *created;
            CheckApiStatus(
                iree_io_parameter_index_provider_create(
                    iree_make_string_view(scope.data(), scope.size()),
                    self.raw_ptr(), *max_concurrent_operations,
                    iree_allocator_system(), &created),
                "Could not create parameter provider from index");
            return ParameterProvider::StealFromRawPtr(created);
          },
          py::arg("scope") = std::string(),
          py::arg("max_concurrent_operations") = py::none())
      .def(
          "create_archive_file",
          [](ParameterIndex &self, std::string file_path,
             iree_io_physical_offset_t file_offset,
             ParameterIndex *explicit_target_index) {
            // If no target index was given, RAII manage a local target index.
            iree_io_parameter_index_t *target_index = nullptr;
            ParameterIndex default_target_index;
            if (explicit_target_index) {
              target_index = explicit_target_index->raw_ptr();
            } else {
              iree_io_parameter_index_t *created;
              CheckApiStatus(iree_io_parameter_index_create(
                                 iree_allocator_system(), &created),
                             "Could not create IO parameter index");
              default_target_index = ParameterIndex::StealFromRawPtr(created);
              target_index = default_target_index.raw_ptr();
            }

            // Open the file via callback.
            struct OpenParams {
              const char *path;
            };
            OpenParams file_open_user_data{file_path.c_str()};
            auto file_open_callback =
                +[](void *user_data, iree_io_physical_offset_t archive_offset,
                    iree_io_physical_size_t archive_length,
                    iree_io_file_handle_t **out_file_handle) -> iree_status_t {
              OpenParams *params = static_cast<OpenParams *>(user_data);
              iree_file_contents_t *file_contents = NULL;
              IREE_RETURN_IF_ERROR(iree_file_create_mapped(
                  params->path, archive_offset + archive_length, archive_offset,
                  (iree_host_size_t)archive_length, iree_allocator_system(),
                  &file_contents));
              iree_io_file_handle_release_callback_t release_callback;
              memset(&release_callback, 0, sizeof(release_callback));
              release_callback.fn =
                  +[](void *user_data,
                      iree_io_file_handle_primitive_t handle_primitive) {
                    iree_file_contents_free(
                        static_cast<iree_file_contents_t *>(user_data));
                  };
              release_callback.user_data = file_contents;
              iree_status_t status = iree_io_file_handle_wrap_host_allocation(
                  IREE_IO_FILE_ACCESS_WRITE, file_contents->buffer,
                  release_callback, iree_allocator_system(), out_file_handle);
              if (!iree_status_is_ok(status)) {
                iree_file_contents_free(file_contents);
              }
              return status;
            };

            // Write the archive.
            CheckApiStatus(iree_io_build_parameter_archive(
                               self.raw_ptr(), target_index,
                               iree_io_parameter_archive_file_open_callback_t{
                                   file_open_callback,
                                   &file_open_user_data,
                               },
                               file_offset, iree_allocator_system()),
                           "Error building parameter archive");

            // Return the target index.
            return ParameterIndex::BorrowFromRawPtr(target_index);
          },
          py::arg("file_path"), py::arg("file_offset") = 0,
          py::arg("target_index") = nullptr);
}

}  // namespace iree::python