#include <algorithm>
#include <chrono>
#include <cstring>
#include <exception>
#include <memory>
#include <sstream>
#include <string>
#include <vector>

#include <dnn/hb_dnn.h>
#include <dnn/hb_dnn_ext.h>
#include <dnn/hb_dnn_status.h>

namespace {

constexpr int kOutputFloats = 12;

void set_error(char *err, int err_len, const std::string &msg) {
  if (err == nullptr || err_len <= 0) {
    return;
  }
  std::snprintf(err, static_cast<size_t>(err_len), "%s", msg.c_str());
}

std::string dnn_error(const std::string &where, int code) {
  std::ostringstream oss;
  oss << where << " failed: " << code;
  const char *desc = hbDNNGetErrorDesc(code);
  if (desc != nullptr) {
    oss << " (" << desc << ")";
  }
  return oss.str();
}

int element_count(const hbDNNTensorShape &shape) {
  int count = 1;
  for (int i = 0; i < shape.numDimensions; ++i) {
    count *= shape.dimensionSize[i];
  }
  return count;
}

void check_dnn(int code, const std::string &where) {
  if (code != 0) {
    throw std::runtime_error(dnn_error(where, code));
  }
}

class BpuDnnPolicy {
 public:
  BpuDnnPolicy(const char *model_path, int bpu_core, int priority)
      : bpu_core_(bpu_core), priority_(priority) {
    if (model_path == nullptr || model_path[0] == '\0') {
      throw std::runtime_error("empty model path");
    }

    const char *model_files[] = {model_path};
    check_dnn(hbDNNInitializeFromFiles(&packed_handle_, model_files, 1),
              "hbDNNInitializeFromFiles");

    const char **model_names = nullptr;
    int32_t model_count = 0;
    check_dnn(hbDNNGetModelNameList(&model_names, &model_count, packed_handle_),
              "hbDNNGetModelNameList");
    if (model_count <= 0 || model_names == nullptr || model_names[0] == nullptr) {
      throw std::runtime_error("model has no names");
    }
    model_name_ = model_names[0];
    check_dnn(hbDNNGetModelHandle(&dnn_handle_, packed_handle_, model_names[0]),
              "hbDNNGetModelHandle");

    int32_t input_count = 0;
    int32_t output_count = 0;
    check_dnn(hbDNNGetInputCount(&input_count, dnn_handle_), "hbDNNGetInputCount");
    check_dnn(hbDNNGetOutputCount(&output_count, dnn_handle_), "hbDNNGetOutputCount");
    if (input_count != 1 || output_count != 1) {
      std::ostringstream oss;
      oss << "expected 1 input and 1 output, got " << input_count << " inputs and "
          << output_count << " outputs";
      throw std::runtime_error(oss.str());
    }

    input_tensors_.resize(1);
    output_tensors_.resize(1);
    check_dnn(hbDNNGetInputTensorProperties(&input_tensors_[0].properties, dnn_handle_, 0),
              "hbDNNGetInputTensorProperties");
    check_dnn(hbDNNGetOutputTensorProperties(&output_tensors_[0].properties, dnn_handle_, 0),
              "hbDNNGetOutputTensorProperties");

    input_floats_ = element_count(input_tensors_[0].properties.validShape);
    output_floats_ = element_count(output_tensors_[0].properties.validShape);
    validate_tensor(input_tensors_[0].properties, HB_DNN_TENSOR_TYPE_F32, "input");
    validate_tensor(output_tensors_[0].properties, HB_DNN_TENSOR_TYPE_F32, "output");

    alloc_tensor_mem(input_tensors_[0]);
    alloc_tensor_mem(output_tensors_[0]);
  }

  ~BpuDnnPolicy() {
    release_tensor_mem(input_tensors_);
    release_tensor_mem(output_tensors_);
    if (packed_handle_ != nullptr) {
      hbDNNRelease(packed_handle_);
      packed_handle_ = nullptr;
    }
  }

  void infer(const float *input, float *output) {
    if (input == nullptr || output == nullptr) {
      throw std::runtime_error("null input/output pointer");
    }

    auto &input_mem = input_tensors_[0].sysMem[0];
    std::memcpy(input_mem.virAddr, input, sizeof(float) * input_floats_);
    int code = hbSysFlushMem(&input_mem, HB_SYS_MEM_CACHE_CLEAN);
    if (code != 0) {
      throw std::runtime_error(dnn_error("hbSysFlushMem(input)", code));
    }

    hbDNNInferCtrlParam ctrl;
    HB_DNN_INITIALIZE_INFER_CTRL_PARAM(&ctrl);
    ctrl.bpuCoreId = bpu_core_;
    ctrl.priority = priority_;
    ctrl.more = false;

    hbDNNTaskHandle_t task_handle = nullptr;
    hbDNNTensor *output_ptr = output_tensors_.data();
    check_dnn(hbDNNInfer(&task_handle, &output_ptr, input_tensors_.data(), dnn_handle_, &ctrl),
              "hbDNNInfer");
    check_dnn(hbDNNWaitTaskDone(task_handle, 0), "hbDNNWaitTaskDone");
    check_dnn(hbDNNReleaseTask(task_handle), "hbDNNReleaseTask");
    task_handle = nullptr;

    auto &output_mem = output_tensors_[0].sysMem[0];
    code = hbSysFlushMem(&output_mem, HB_SYS_MEM_CACHE_INVALIDATE);
    if (code != 0) {
      throw std::runtime_error(dnn_error("hbSysFlushMem(output)", code));
    }
    std::memcpy(output, output_mem.virAddr, sizeof(float) * output_floats_);
  }

  int input_floats() const { return input_floats_; }
  int output_floats() const { return output_floats_; }

 private:
  static void validate_tensor(const hbDNNTensorProperties &props, int expected_type,
                              const char *name) {
    if (props.tensorType != expected_type) {
      std::ostringstream oss;
      oss << name << " tensor type " << props.tensorType << " != " << expected_type;
      throw std::runtime_error(oss.str());
    }
    if (props.alignedByteSize <= 0) {
      throw std::runtime_error(std::string(name) + " alignedByteSize <= 0");
    }
  }

  static void alloc_tensor_mem(hbDNNTensor &tensor) {
    std::memset(tensor.sysMem, 0, sizeof(tensor.sysMem));
    const uint32_t bytes = static_cast<uint32_t>(tensor.properties.alignedByteSize);
    const int code = hbSysAllocCachedMem(&tensor.sysMem[0], bytes);
    if (code != 0) {
      throw std::runtime_error(dnn_error("hbSysAllocCachedMem", code));
    }
  }

  static void release_tensor_mem(std::vector<hbDNNTensor> &tensors) {
    for (auto &tensor : tensors) {
      if (tensor.sysMem[0].virAddr != nullptr || tensor.sysMem[0].phyAddr != 0) {
        hbSysFreeMem(&tensor.sysMem[0]);
        std::memset(tensor.sysMem, 0, sizeof(tensor.sysMem));
      }
    }
  }

  hbPackedDNNHandle_t packed_handle_{nullptr};
  hbDNNHandle_t dnn_handle_{nullptr};
  std::string model_name_;
  std::vector<hbDNNTensor> input_tensors_;
  std::vector<hbDNNTensor> output_tensors_;
  int input_floats_{0};
  int output_floats_{0};
  int bpu_core_{HB_BPU_CORE_ANY};
  int priority_{HB_DNN_PRIORITY_LOWEST};
};

}  // namespace

extern "C" {

void *rlgym_bpu_create(const char *model_path, int bpu_core, int priority,
                       char *err, int err_len) {
  try {
    set_error(err, err_len, "");
    return new BpuDnnPolicy(model_path, bpu_core, priority);
  } catch (const std::exception &e) {
    set_error(err, err_len, e.what());
    return nullptr;
  }
}

int rlgym_bpu_infer(void *handle, const float *input, float *output, char *err,
                    int err_len) {
  try {
    set_error(err, err_len, "");
    if (handle == nullptr) {
      throw std::runtime_error("null policy handle");
    }
    static_cast<BpuDnnPolicy *>(handle)->infer(input, output);
    return 0;
  } catch (const std::exception &e) {
    set_error(err, err_len, e.what());
    return -1;
  }
}

int rlgym_bpu_input_floats(void *handle) {
  if (handle == nullptr) {
    return 0;
  }
  return static_cast<BpuDnnPolicy *>(handle)->input_floats();
}

int rlgym_bpu_output_floats(void *handle) {
  if (handle == nullptr) {
    return 0;
  }
  return static_cast<BpuDnnPolicy *>(handle)->output_floats();
}

void rlgym_bpu_destroy(void *handle) {
  delete static_cast<BpuDnnPolicy *>(handle);
}

const char *rlgym_bpu_version() {
  return "cpp_dnn_api_x5";
}

}
