#include <torch/extension.h>

#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cublasLt.h>
#include <cuda_runtime_api.h>

#include <array>
#include <cstdint>
#include <stdexcept>
#include <string>
#include <vector>

namespace py = pybind11;

#define LT_CHECK(call)                                                        \
  do {                                                                        \
    const cublasStatus_t status_ = (call);                                     \
    if (status_ != CUBLAS_STATUS_SUCCESS)                                      \
      throw std::runtime_error(std::string(#call) + " failed: " +             \
                               std::to_string(static_cast<int>(status_)));     \
  } while (0)

namespace {

thread_local cublasLtHandle_t handle = nullptr;

cublasLtHandle_t get_handle() {
  if (!handle) LT_CHECK(cublasLtCreate(&handle));
  return handle;
}

void require_cuda_contiguous(const torch::Tensor& tensor, const char* name) {
  TORCH_CHECK(tensor.is_cuda(), name, " must be CUDA");
  TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
}

void require_same_device(const torch::Tensor& tensor,
                         const torch::Tensor& reference,
                         const char* name) {
  TORCH_CHECK(tensor.device() == reference.device(), name,
              " must be on the same CUDA device as activation_qdata");
}

struct Problem {
  cublasLtMatmulDesc_t operation = nullptr;
  cublasLtMatrixLayout_t a = nullptr, b = nullptr, c = nullptr, d = nullptr;

  Problem(const torch::Tensor& activation_qdata,
          const torch::Tensor& activation_block_scale,
          const torch::Tensor& weight_qdata,
          const torch::Tensor& weight_block_scale) {
    require_cuda_contiguous(activation_qdata, "activation_qdata");
    require_cuda_contiguous(activation_block_scale, "activation_block_scale");
    require_cuda_contiguous(weight_qdata, "weight_qdata");
    require_cuda_contiguous(weight_block_scale, "weight_block_scale");
    require_same_device(activation_block_scale, activation_qdata,
                        "activation_block_scale");
    require_same_device(weight_qdata, activation_qdata, "weight_qdata");
    require_same_device(weight_block_scale, activation_qdata,
                        "weight_block_scale");
    TORCH_CHECK(activation_qdata.scalar_type() == at::kByte &&
                weight_qdata.scalar_type() == at::kByte,
                "packed NVFP4 operands must use uint8 storage");
    TORCH_CHECK(activation_qdata.dim() == 2 && weight_qdata.dim() == 2,
                "packed NVFP4 operands must be rank two");
    TORCH_CHECK(activation_block_scale.element_size() == 1 &&
                weight_block_scale.element_size() == 1,
                "NVFP4 block scales must use one-byte E4M3 storage");
    TORCH_CHECK(activation_qdata.size(1) == weight_qdata.size(1),
                "packed K dimensions differ");

    // This is deliberately the same column-major reinterpretation used by
    // Comfy Kitchen 0.2.31: weight is Lt A, activation is Lt B, and D is D^T.
    const int64_t m = weight_qdata.size(0);       // row-major N
    const int64_t n = activation_qdata.size(0);  // row-major padded M
    const int64_t k = activation_qdata.size(1) * 2;
    TORCH_CHECK(activation_block_scale.numel() >= n * (k / 16),
                "activation_block_scale storage is too small");
    TORCH_CHECK(weight_block_scale.numel() >= m * (k / 16),
                "weight_block_scale storage is too small");
    LT_CHECK(cublasLtMatmulDescCreate(&operation, CUBLAS_COMPUTE_32F,
                                     CUDA_R_32F));
    cublasLtMatmulMatrixScale_t scale_mode =
        CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3;
    LT_CHECK(cublasLtMatmulDescSetAttribute(
        operation, CUBLASLT_MATMUL_DESC_A_SCALE_MODE, &scale_mode,
        sizeof(scale_mode)));
    LT_CHECK(cublasLtMatmulDescSetAttribute(
        operation, CUBLASLT_MATMUL_DESC_B_SCALE_MODE, &scale_mode,
        sizeof(scale_mode)));
    const cublasOperation_t transa = CUBLAS_OP_T;
    const cublasOperation_t transb = CUBLAS_OP_N;
    LT_CHECK(cublasLtMatmulDescSetAttribute(
        operation, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa)));
    LT_CHECK(cublasLtMatmulDescSetAttribute(
        operation, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transb)));
    const void* a_scale = weight_block_scale.data_ptr();
    const void* b_scale = activation_block_scale.data_ptr();
    LT_CHECK(cublasLtMatmulDescSetAttribute(
        operation, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, &a_scale,
        sizeof(a_scale)));
    LT_CHECK(cublasLtMatmulDescSetAttribute(
        operation, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &b_scale,
        sizeof(b_scale)));
    const cublasDataType_t scale_type = CUDA_R_32F;
    LT_CHECK(cublasLtMatmulDescSetAttribute(
        operation, CUBLASLT_MATMUL_DESC_SCALE_TYPE, &scale_type,
        sizeof(scale_type)));
    const cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_DEVICE;
    LT_CHECK(cublasLtMatmulDescSetAttribute(
        operation, CUBLASLT_MATMUL_DESC_POINTER_MODE, &pointer_mode,
        sizeof(pointer_mode)));
    const cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_DEFAULT;
    LT_CHECK(cublasLtMatmulDescSetAttribute(
        operation, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue,
        sizeof(epilogue)));

    LT_CHECK(cublasLtMatrixLayoutCreate(&a, CUDA_R_4F_E2M1, k, m, k));
    LT_CHECK(cublasLtMatrixLayoutCreate(&b, CUDA_R_4F_E2M1, k, n, k));
    LT_CHECK(cublasLtMatrixLayoutCreate(&c, CUDA_R_16BF, m, n, m));
    LT_CHECK(cublasLtMatrixLayoutCreate(&d, CUDA_R_16BF, m, n, m));
  }

  ~Problem() {
    if (d) cublasLtMatrixLayoutDestroy(d);
    if (c) cublasLtMatrixLayoutDestroy(c);
    if (b) cublasLtMatrixLayoutDestroy(b);
    if (a) cublasLtMatrixLayoutDestroy(a);
    if (operation) cublasLtMatmulDescDestroy(operation);
  }
};

template <typename T>
bool config_get(const cublasLtMatmulAlgo_t& algo,
                cublasLtMatmulAlgoConfigAttributes_t attr, T* value) {
  size_t written = 0;
  return cublasLtMatmulAlgoConfigGetAttribute(&algo, attr, value,
                                               sizeof(T), &written) ==
             CUBLAS_STATUS_SUCCESS &&
         written == sizeof(T);
}

template <typename T>
void put_config(py::dict& result, const char* name,
                const cublasLtMatmulAlgo_t& algo,
                cublasLtMatmulAlgoConfigAttributes_t attr) {
  T value{};
  if (config_get(algo, attr, &value)) result[name] = value;
}

template <typename T>
void put_cap_scalar(py::dict& caps, const char* name,
                    const cublasLtMatmulAlgo_t& algo,
                    cublasLtMatmulAlgoCapAttributes_t attr) {
  T value{};
  size_t written = 0;
  if (cublasLtMatmulAlgoCapGetAttribute(&algo, attr, &value, sizeof(value),
                                        &written) == CUBLAS_STATUS_SUCCESS &&
      written == sizeof(value))
    caps[name] = value;
}

void put_cap_array(py::dict& caps, const char* name,
                   const cublasLtMatmulAlgo_t& algo,
                   cublasLtMatmulAlgoCapAttributes_t attr) {
  size_t bytes = 0;
  if (cublasLtMatmulAlgoCapGetAttribute(&algo, attr, nullptr, 0, &bytes) !=
          CUBLAS_STATUS_SUCCESS ||
      bytes == 0)
    return;
  std::vector<uint32_t> values((bytes + sizeof(uint32_t) - 1) /
                               sizeof(uint32_t));
  size_t written = 0;
  if (cublasLtMatmulAlgoCapGetAttribute(&algo, attr, values.data(), bytes,
                                        &written) != CUBLAS_STATUS_SUCCESS)
    return;
  values.resize(written / sizeof(uint32_t));
  caps[name] = values;
}

py::dict describe(const cublasLtMatmulAlgo_t& algo,
                  const cublasLtMatmulHeuristicResult_t& checked,
                  cublasStatus_t api_status) {
  py::dict result;
  for (const char* name : {"algorithm_id", "tile_id", "stages_id", "split_k",
                           "reduction_scheme", "custom_option", "cta_swizzle",
                           "inner_shape", "cluster_shape"})
    result[name] = py::none();
  put_config<int>(result, "algorithm_id", algo, CUBLASLT_ALGO_CONFIG_ID);
  put_config<uint32_t>(result, "tile_id", algo, CUBLASLT_ALGO_CONFIG_TILE_ID);
  put_config<uint32_t>(result, "stages_id", algo,
                       CUBLASLT_ALGO_CONFIG_STAGES_ID);
  // CUDA 13 documents SPLITK_NUM as int32_t. Preserve negative library
  // sentinel values instead of wrapping them into fictitious huge factors.
  put_config<int32_t>(result, "split_k", algo,
                      CUBLASLT_ALGO_CONFIG_SPLITK_NUM);
  put_config<uint32_t>(result, "reduction_scheme", algo,
                       CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME);
  put_config<uint32_t>(result, "custom_option", algo,
                       CUBLASLT_ALGO_CONFIG_CUSTOM_OPTION);
  put_config<uint32_t>(result, "cta_swizzle", algo,
                       CUBLASLT_ALGO_CONFIG_CTA_SWIZZLING);
#if CUDA_VERSION >= 12000
  put_config<uint32_t>(result, "inner_shape", algo,
                       CUBLASLT_ALGO_CONFIG_INNER_SHAPE_ID);
  put_config<uint32_t>(result, "cluster_shape", algo,
                       CUBLASLT_ALGO_CONFIG_CLUSTER_SHAPE_ID);
#endif
  result["required_workspace_bytes"] = checked.workspaceSize;
  result["waves"] = checked.wavesCount;
  result["state"] = static_cast<int>(checked.state);
  result["api_status"] = static_cast<int>(api_status);
  result["valid"] = api_status == CUBLAS_STATUS_SUCCESS &&
                    checked.state == CUBLAS_STATUS_SUCCESS;

  py::dict caps;
  put_cap_scalar<int>(caps, "split_k_support", algo,
                      CUBLASLT_ALGO_CAP_SPLITK_SUPPORT);
  put_cap_scalar<uint32_t>(caps, "reduction_scheme_mask", algo,
                           CUBLASLT_ALGO_CAP_REDUCTION_SCHEME_MASK);
  put_cap_scalar<uint32_t>(caps, "cta_swizzle_support", algo,
                           CUBLASLT_ALGO_CAP_CTA_SWIZZLING_SUPPORT);
  put_cap_scalar<int>(caps, "custom_option_max", algo,
                      CUBLASLT_ALGO_CAP_CUSTOM_OPTION_MAX);
  put_cap_scalar<int>(caps, "strided_batch_support", algo,
                      CUBLASLT_ALGO_CAP_STRIDED_BATCH_SUPPORT);
  put_cap_scalar<int>(caps, "out_of_place_result_support", algo,
                      CUBLASLT_ALGO_CAP_OUT_OF_PLACE_RESULT_SUPPORT);
  put_cap_array(caps, "tile_ids", algo, CUBLASLT_ALGO_CAP_TILE_IDS);
  put_cap_array(caps, "stages_ids", algo, CUBLASLT_ALGO_CAP_STAGES_IDS);
  caps["inner_cluster_shape_capability_note"] =
      "This CUDA 13 cublasLt.h exposes config IDs but no public capability "
      "attributes that enumerate inner/cluster shape IDs.";
  result["capabilities"] = caps;
  return result;
}

cublasLtMatmulAlgo_t init_algo(int algorithm_id) {
  cublasLtMatmulAlgo_t algo{};
  LT_CHECK(cublasLtMatmulAlgoInit(
      get_handle(), CUBLAS_COMPUTE_32F, CUDA_R_32F, CUDA_R_4F_E2M1,
      CUDA_R_4F_E2M1, CUDA_R_16BF, CUDA_R_16BF, algorithm_id, &algo));
  return algo;
}

template <typename T>
void maybe_set(cublasLtMatmulAlgo_t* algo, const py::dict& config,
               const char* key, cublasLtMatmulAlgoConfigAttributes_t attr) {
  if (!config.contains(key) || config[key].is_none()) return;
  const T value = config[key].cast<T>();
  LT_CHECK(cublasLtMatmulAlgoConfigSetAttribute(algo, attr, &value,
                                                sizeof(value)));
}

cublasLtMatmulAlgo_t configured_algo(const py::dict& config) {
  TORCH_CHECK(config.contains("algorithm_id"), "algorithm_id is required");
  auto algo = init_algo(config["algorithm_id"].cast<int>());
  maybe_set<uint32_t>(&algo, config, "tile_id", CUBLASLT_ALGO_CONFIG_TILE_ID);
  maybe_set<uint32_t>(&algo, config, "stages_id",
                      CUBLASLT_ALGO_CONFIG_STAGES_ID);
  maybe_set<int32_t>(&algo, config, "split_k",
                     CUBLASLT_ALGO_CONFIG_SPLITK_NUM);
  maybe_set<uint32_t>(&algo, config, "reduction_scheme",
                      CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME);
  maybe_set<uint32_t>(&algo, config, "custom_option",
                      CUBLASLT_ALGO_CONFIG_CUSTOM_OPTION);
  maybe_set<uint32_t>(&algo, config, "cta_swizzle",
                      CUBLASLT_ALGO_CONFIG_CTA_SWIZZLING);
#if CUDA_VERSION >= 12000
  maybe_set<uint32_t>(&algo, config, "inner_shape",
                      CUBLASLT_ALGO_CONFIG_INNER_SHAPE_ID);
  maybe_set<uint32_t>(&algo, config, "cluster_shape",
                      CUBLASLT_ALGO_CONFIG_CLUSTER_SHAPE_ID);
#endif
  return algo;
}

py::list enumerate(torch::Tensor activation_qdata,
                   torch::Tensor activation_block_scale,
                   torch::Tensor weight_qdata,
                   torch::Tensor weight_block_scale,
                   int64_t max_workspace, int requested_count) {
  c10::cuda::CUDAGuard guard(activation_qdata.device());
  Problem problem(activation_qdata, activation_block_scale, weight_qdata,
                  weight_block_scale);
  cublasLtMatmulPreference_t preference = nullptr;
  LT_CHECK(cublasLtMatmulPreferenceCreate(&preference));
  LT_CHECK(cublasLtMatmulPreferenceSetAttribute(
      preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &max_workspace,
      sizeof(max_workspace)));
  std::vector<cublasLtMatmulHeuristicResult_t> found(requested_count);
  int returned = 0;
  const auto status = cublasLtMatmulAlgoGetHeuristic(
      get_handle(), problem.operation, problem.a, problem.b, problem.c,
      problem.d, preference, requested_count, found.data(), &returned);
  cublasLtMatmulPreferenceDestroy(preference);
  LT_CHECK(status);
  py::list output;
  for (int i = 0; i < returned; ++i)
    output.append(describe(found[i].algo, found[i], CUBLAS_STATUS_SUCCESS));
  return output;
}

py::dict check(torch::Tensor activation_qdata,
               torch::Tensor activation_block_scale,
               torch::Tensor weight_qdata,
               torch::Tensor weight_block_scale, py::dict config) {
  c10::cuda::CUDAGuard guard(activation_qdata.device());
  Problem problem(activation_qdata, activation_block_scale, weight_qdata,
                  weight_block_scale);
  auto algo = configured_algo(config);
  cublasLtMatmulHeuristicResult_t result{};
  const auto status = cublasLtMatmulAlgoCheck(
      get_handle(), problem.operation, problem.a, problem.b, problem.c,
      problem.d, &algo, &result);
  return describe(algo, result, status);
}

void run(torch::Tensor activation_qdata,
         torch::Tensor activation_block_scale,
         torch::Tensor weight_qdata, torch::Tensor weight_block_scale,
         torch::Tensor alpha, torch::Tensor beta, torch::Tensor output,
         torch::Tensor workspace, py::dict config) {
  c10::cuda::CUDAGuard guard(activation_qdata.device());
  require_cuda_contiguous(alpha, "alpha");
  require_cuda_contiguous(beta, "beta");
  require_cuda_contiguous(output, "output");
  require_cuda_contiguous(workspace, "workspace");
  require_same_device(alpha, activation_qdata, "alpha");
  require_same_device(beta, activation_qdata, "beta");
  require_same_device(output, activation_qdata, "output");
  require_same_device(workspace, activation_qdata, "workspace");
  TORCH_CHECK(alpha.scalar_type() == at::kFloat && alpha.numel() == 1,
              "alpha must be one device FP32 value");
  TORCH_CHECK(beta.scalar_type() == at::kFloat && beta.numel() == 1,
              "beta must be one device FP32 value");
  TORCH_CHECK(output.scalar_type() == at::kBFloat16 && output.dim() == 2,
              "output must be rank-two BF16");
  TORCH_CHECK(workspace.scalar_type() == at::kByte,
              "workspace must use uint8 storage");
  TORCH_CHECK(output.size(0) == activation_qdata.size(0) &&
              output.size(1) == weight_qdata.size(0),
              "output must be [packed activation rows, weight rows]");
  Problem problem(activation_qdata, activation_block_scale, weight_qdata,
                  weight_block_scale);
  auto algo = configured_algo(config);
  cublasLtMatmulHeuristicResult_t checked{};
  LT_CHECK(cublasLtMatmulAlgoCheck(get_handle(), problem.operation, problem.a,
                                  problem.b, problem.c, problem.d, &algo,
                                  &checked));
  TORCH_CHECK(checked.state == CUBLAS_STATUS_SUCCESS,
              "selected algorithm failed AlgoCheck with state ",
              static_cast<int>(checked.state));
  TORCH_CHECK(checked.workspaceSize <= static_cast<size_t>(workspace.numel()),
              "selected algorithm requires ", checked.workspaceSize,
              " workspace bytes but caller supplied ", workspace.numel());
  const auto stream = at::cuda::getCurrentCUDAStream(
      activation_qdata.get_device()).stream();
  void* workspace_ptr = workspace.numel() ? workspace.data_ptr() : nullptr;
  LT_CHECK(cublasLtMatmul(
      get_handle(), problem.operation, alpha.data_ptr(), weight_qdata.data_ptr(),
      problem.a, activation_qdata.data_ptr(), problem.b, beta.data_ptr(),
      output.data_ptr(), problem.c, output.data_ptr(), problem.d, &algo,
      workspace_ptr, workspace.numel(), stream));
}

py::dict build_info() {
  py::dict result;
  result["cuda_version"] = CUDA_VERSION;
  result["cublas_version"] = CUBLAS_VERSION;
  result["stream_k_public_control"] = false;
  result["stream_k_note"] =
      "CUDA 13 cuBLASLt exposes no documented MatmulAlgoConfig attribute "
      "that directly selects Stream-K. Negative SPLITK_NUM values returned "
      "by heuristics are preserved as undocumented library sentinels, not "
      "claimed as public Stream-K control.";
  return result;
}

}  // namespace

PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
  module.def("enumerate", &enumerate);
  module.def("check", &check);
  module.def("run", &run);
  module.def("build_info", &build_info);
}