#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/core/DeviceType.h>
#include <c10/cuda/CUDACachingAllocator.h>
#include <c10/cuda/CUDAException.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <cuda_runtime_api.h>
#include <torch/library.h>

#include <cstdint>
#include <limits>
#include <memory>
#include <mutex>
#include <optional>
#include <stdexcept>
#include <unordered_map>
#include <utility>
#include <vector>

#include "nvfp4_linear.h"

namespace {

constexpr int kAbiVersion = 1;
constexpr int64_t kFp4BlockElements = 16;
constexpr int64_t kScaleTileOuter = 128;
constexpr int64_t kScaleTileInner = 4;

int64_t round_up_int64(int64_t value, int64_t multiple) {
  TORCH_CHECK(value > 0, "round_up requires a positive value");
  TORCH_CHECK(multiple > 0, "round_up requires a positive multiple");
  TORCH_CHECK(
      value <= std::numeric_limits<int64_t>::max() - (multiple - 1),
      "round_up overflow");
  return ((value + multiple - 1) / multiple) * multiple;
}

int checked_int_arg(int64_t value, const char* label) {
  TORCH_CHECK(value > 0, label, " must be positive");
  TORCH_CHECK(
      value <= static_cast<int64_t>(std::numeric_limits<int>::max()),
      label,
      " exceeds 32-bit ABI bound");
  return static_cast<int>(value);
}

size_t expected_packed_weight_bytes(int out_features, int in_features) {
  TORCH_CHECK(
      out_features > 0 && out_features % 8 == 0,
      "resident NVFP4 requires out_features divisible by 8, got ",
      out_features);
  TORCH_CHECK(
      in_features > 0 && in_features % 32 == 0,
      "resident NVFP4 requires in_features divisible by 32, got ",
      in_features);
  return (static_cast<size_t>(out_features) * static_cast<size_t>(in_features)) /
      2;
}

size_t expected_scale_bytes(int out_features, int in_features) {
  TORCH_CHECK(
      in_features > 0 && in_features % kFp4BlockElements == 0,
      "resident NVFP4 requires in_features divisible by 16, got ",
      in_features);
  TORCH_CHECK(out_features > 0, "out_features must be positive");
  const int64_t inner_dim =
      round_up_int64(in_features / kFp4BlockElements, kScaleTileInner);
  const int64_t outer_tiles =
      (out_features + kScaleTileOuter - 1) / kScaleTileOuter;
  return static_cast<size_t>(inner_dim) * static_cast<size_t>(outer_tiles) *
      static_cast<size_t>(kScaleTileOuter);
}

[[noreturn]] void throw_abi_error(const char* operation) {
  const char* message = mage_nvfp4_last_error();
  TORCH_CHECK(
      false,
      operation,
      ": ",
      message == nullptr ? "unknown native resident error" : message);
}

void ensure_abi_version() {
  static std::once_flag once;
  std::call_once(once, []() {
    const int version = mage_nvfp4_abi_version();
    TORCH_CHECK(
        version == kAbiVersion,
        "resident NVFP4 ABI mismatch: compiled torch op expects ",
        kAbiVersion,
        " but native library reports ",
        version);
  });
}

struct ContextKey {
  int device_index = -1;
  uintptr_t stream = 0;

  bool operator==(const ContextKey& other) const noexcept {
    return device_index == other.device_index && stream == other.stream;
  }
};

struct ContextKeyHash {
  size_t operator()(const ContextKey& key) const noexcept {
    const size_t left = std::hash<int>{}(key.device_index);
    const size_t right = std::hash<uintptr_t>{}(key.stream);
    return left ^ (right + 0x9e3779b97f4a7c15ULL + (left << 6) + (left >> 2));
  }
};

struct NativeContextDeleter {
  void operator()(void* context) const noexcept {
    if (context == nullptr) {
      return;
    }
    (void)mage_nvfp4_destroy_context(context);
  }
};

class ContextRegistry {
 public:
  void* get(int device_index, uintptr_t stream) {
    std::lock_guard<std::mutex> guard(mutex_);
    const ContextKey key{device_index, stream};
    auto it = contexts_.find(key);
    if (it != contexts_.end()) {
      return it->second.get();
    }
    void* context = nullptr;
    const int status = mage_nvfp4_create_context(device_index, &context);
    if (status != 0) {
      throw_abi_error("creating resident NVFP4 context");
    }
    auto inserted = contexts_.emplace(
        key, std::unique_ptr<void, NativeContextDeleter>(context));
    return inserted.first->second.get();
  }

  void clear() {
    decltype(contexts_) retired;
    {
      std::lock_guard<std::mutex> guard(mutex_);
      retired.swap(contexts_);
    }
    // Destruction synchronizes each bound stream through the resident ABI.
    // Keep it outside the registry mutex.
  }

 private:
  std::mutex mutex_;
  std::unordered_map<
      ContextKey,
      std::unique_ptr<void, NativeContextDeleter>,
      ContextKeyHash>
      contexts_;
};

ContextRegistry& context_registry() {
  // Deliberately avoid a static destructor that could call CUDA after runtime
  // teardown. Long-lived processes must call the explicit close op; otherwise
  // the OS reclaims these process resources at exit.
  static ContextRegistry* registry = new ContextRegistry();
  return *registry;
}

void clear_native_contexts() {
  context_registry().clear();
}

void record_tensor_stream(
    const at::Tensor& tensor, c10::cuda::CUDAStream stream) {
  if (!tensor.defined() || !tensor.is_cuda()) {
    return;
  }
  c10::cuda::CUDACachingAllocator::recordStream(
      tensor.storage().data_ptr(), stream);
}

void validate_common(
    const at::Tensor& input,
    const at::Tensor& packed_weight,
    const at::Tensor& weight_scales,
    const at::Tensor& weight_scale,
    const std::optional<at::Tensor>& bias,
    int64_t in_features,
    int64_t out_features) {
  TORCH_CHECK(input.dim() >= 1, "resident NVFP4 input must have at least one dimension");
  TORCH_CHECK(
      input.scalar_type() == at::kBFloat16,
      "resident NVFP4 input must be bfloat16");
  TORCH_CHECK(!input.requires_grad(), "resident NVFP4 torch op is inference-only");
  TORCH_CHECK(
      input.size(-1) == in_features,
      "resident NVFP4 expected last dimension ",
      in_features,
      " but got ",
      input.size(-1));
  TORCH_CHECK(
      packed_weight.scalar_type() == at::kByte &&
          packed_weight.dim() == 1 && packed_weight.is_contiguous(),
      "packed_weight must be a contiguous 1D uint8 tensor");
  TORCH_CHECK(
      weight_scales.scalar_type() == at::kByte &&
          weight_scales.dim() == 1 && weight_scales.is_contiguous(),
      "weight_scales must be a contiguous 1D uint8 tensor");
  TORCH_CHECK(
      weight_scale.scalar_type() == at::kFloat &&
          weight_scale.numel() == 1 && weight_scale.is_contiguous(),
      "weight_scale must be one contiguous float32 value");
  TORCH_CHECK(
      packed_weight.device() == input.device(),
      "packed_weight and input must be on the same device");
  TORCH_CHECK(
      weight_scales.device() == input.device(),
      "weight_scales and input must be on the same device");
  TORCH_CHECK(
      weight_scale.device() == input.device(),
      "weight_scale and input must be on the same device");
  if (bias.has_value()) {
    const at::Tensor& bias_value = *bias;
    TORCH_CHECK(
        bias_value.scalar_type() == at::kBFloat16 &&
            bias_value.dim() == 1 && bias_value.is_contiguous(),
        "bias must be a contiguous 1D bfloat16 tensor");
    TORCH_CHECK(
        bias_value.numel() == out_features,
        "bias length must match out_features");
    TORCH_CHECK(
        bias_value.device() == input.device(),
        "bias and input must be on the same device");
  }

  const int in_features_i = checked_int_arg(in_features, "in_features");
  const int out_features_i = checked_int_arg(out_features, "out_features");
  const size_t expected_weight = expected_packed_weight_bytes(
      out_features_i, in_features_i);
  const size_t expected_scales = expected_scale_bytes(
      out_features_i, in_features_i);
  TORCH_CHECK(
      static_cast<size_t>(packed_weight.numel()) == expected_weight,
      "packed_weight size mismatch: expected ",
      expected_weight,
      " bytes but got ",
      packed_weight.numel());
  TORCH_CHECK(
      static_cast<size_t>(weight_scales.numel()) == expected_scales,
      "weight_scales size mismatch: expected ",
      expected_scales,
      " bytes but got ",
      weight_scales.numel());

  const size_t native_weight =
      mage_nvfp4_packed_weight_bytes(out_features_i, in_features_i);
  const size_t native_scales =
      mage_nvfp4_weight_scale_bytes(out_features_i, in_features_i);
  TORCH_CHECK(
      native_weight == expected_weight,
      "native resident library packed-weight metadata disagrees with torch op");
  TORCH_CHECK(
      native_scales == expected_scales,
      "native resident library scale metadata disagrees with torch op");
}

at::Tensor sm120_linear_native_cuda(
    const at::Tensor& input,
    const at::Tensor& packed_weight,
    const at::Tensor& weight_scales,
    const at::Tensor& weight_scale,
    const std::optional<at::Tensor>& bias,
    int64_t in_features,
    int64_t out_features) {
  ensure_abi_version();
  validate_common(
      input, packed_weight, weight_scales, weight_scale, bias, in_features,
      out_features);

  TORCH_CHECK(input.is_cuda(), "resident NVFP4 CUDA implementation requires a CUDA input");
  const auto device = input.device();
  c10::cuda::CUDAGuard guard(device);
  const int device_index = device.index();
  cudaDeviceProp properties{};
  C10_CUDA_CHECK(cudaGetDeviceProperties(&properties, device_index));
  TORCH_CHECK(
      properties.major == 12 && properties.minor == 0,
      "resident NVFP4 CUDA implementation is labeled sm_120-only; selected device reports ",
      properties.major,
      ".",
      properties.minor);

  const c10::cuda::CUDAStream stream =
      c10::cuda::getCurrentCUDAStream(device_index);
  TORCH_CHECK(
      !stream.is_capturing(),
      "resident NVFP4 C ABI torch op does not support CUDA graph capture");

  const int in_features_i = checked_int_arg(in_features, "in_features");
  const int out_features_i = checked_int_arg(out_features, "out_features");
  at::Tensor contiguous =
      input.reshape({-1, in_features_i}).contiguous();
  const int64_t logical_m64 = contiguous.size(0);
  TORCH_CHECK(logical_m64 > 0, "resident NVFP4 torch op does not support empty inputs");
  const int logical_m = checked_int_arg(logical_m64, "logical_m");
  const int64_t padded_m64 = round_up_int64(logical_m64, 8);
  at::Tensor padded_output = at::empty(
      {padded_m64, out_features_i},
      input.options().dtype(at::kBFloat16));

  record_tensor_stream(contiguous, stream);
  record_tensor_stream(packed_weight, stream);
  record_tensor_stream(weight_scales, stream);
  record_tensor_stream(weight_scale, stream);
  if (bias.has_value()) {
    record_tensor_stream(*bias, stream);
  }
  record_tensor_stream(padded_output, stream);

  void* context =
      context_registry().get(device_index, reinterpret_cast<uintptr_t>(stream.stream()));
  const void* bias_pointer =
      bias.has_value() ? bias->data_ptr() : nullptr;
  const int status = mage_nvfp4_linear_forward(
      context,
      contiguous.data_ptr(),
      packed_weight.data_ptr(),
      static_cast<size_t>(packed_weight.numel()),
      weight_scales.data_ptr(),
      static_cast<size_t>(weight_scales.numel()),
      weight_scale.data_ptr(),
      bias_pointer,
      padded_output.data_ptr(),
      logical_m,
      in_features_i,
      out_features_i,
      reinterpret_cast<uintptr_t>(stream.stream()));
  if (status != 0) {
    throw_abi_error("running resident NVFP4 linear forward");
  }

  at::Tensor logical = padded_output.narrow(0, 0, logical_m64);
  std::vector<int64_t> output_sizes = input.sizes().vec();
  output_sizes.back() = out_features_i;
  return logical.view(output_sizes);
}

} // namespace

TORCH_LIBRARY_FRAGMENT(mage_nvfp4, m) {
  m.def(
      "sm120_linear_native(Tensor input, Tensor packed_weight, Tensor weight_scales, Tensor weight_scale, Tensor? bias, int in_features, int out_features) -> Tensor");
  m.def("clear_native_contexts() -> ()", TORCH_FN(clear_native_contexts));
}

TORCH_LIBRARY_IMPL(mage_nvfp4, CUDA, m) {
  m.impl("sm120_linear_native", TORCH_FN(sm120_linear_native_cuda));
}