#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>

#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>

#include <tvm/ffi/container/tensor.h>

#include <cmath>
#include <cstdint>
#include <limits>
#include <string>

namespace {

enum class ActivationKind : uint32_t {
  kSiLU,
  kGELU,
  kGELUTanh,
  kReLU2,
};

template <ActivationKind kAct>
SGL_DEVICE float apply_activation_f32(float x_f32) {
  if constexpr (kAct == ActivationKind::kSiLU) {
    return x_f32 / (1.0f + expf(-x_f32));
  } else if constexpr (kAct == ActivationKind::kGELU) {
    constexpr auto kSqrt1Over2 = 0.7071067811865475f;
    return x_f32 * (0.5f * (1.0f + erff(x_f32 * kSqrt1Over2)));
  } else if constexpr (kAct == ActivationKind::kGELUTanh) {
    constexpr auto kGeluTanhAlpha = 0.044715f;
    constexpr auto kGeluTanhBeta = 0.7978845608028654f;
    const float cdf = 0.5f * (1.0f + tanhf(kGeluTanhBeta * (x_f32 + kGeluTanhAlpha * x_f32 * x_f32 * x_f32)));
    return x_f32 * cdf;
  } else if constexpr (kAct == ActivationKind::kReLU2) {
    const float relu = x_f32 > 0.0f ? x_f32 : 0.0f;
    return relu * relu;
  } else {
    static_assert(host::dependent_false_v<decltype(kAct)>, "unsupported activation kind");
    return 0.0f;
  }
}

struct ActivationParams {
  const void* __restrict__ input;
  void* __restrict__ out;
  uint32_t hidden_dim;
  uint32_t num_tokens;
  // Optional MoE expert filtering: when expert_ids != nullptr, a token is
  // skipped if expert_ids[token_id / expert_step] == -1. expert_step is 1
  // for per-token routing and BLOCK_SIZE_M for sorted/TMA routing.
  const int32_t* __restrict__ expert_ids;
  uint32_t expert_step;
};

template <typename T, ActivationKind kAct, bool kUsePDL, bool kFilterExpert>
__global__ void act_and_mul_kernel(const __grid_constant__ ActivationParams params) {
  using namespace device;
  constexpr auto kVecSize = kMaxVecBytes / sizeof(T);
  using vec_t = AlignedVector<T, kMaxVecBytes / sizeof(T)>;
  const auto num_vecs = params.hidden_dim / kVecSize;  // per token
  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
  const auto token_id = tid / num_vecs;

  if (token_id >= params.num_tokens) return;
  if constexpr (kFilterExpert) {
    if (params.expert_ids[token_id / params.expert_step] == -1) return;
  }
  const auto offset = tid % num_vecs;
  const auto input_offset = token_id * (num_vecs * 2) + offset;
  const auto output_offset = tid;
  PDLWaitPrimary<kUsePDL>();
  const auto gate = device::load_as<vec_t>(params.input, input_offset);
  const auto up = device::load_as<vec_t>(params.input, input_offset + num_vecs);
  vec_t out;
#pragma unroll
  for (int i = 0; i < kVecSize; ++i) {
    const float gate_f32 = device::cast<fp32_t>(gate[i]);
    const float up_f32 = device::cast<fp32_t>(up[i]);
    out[i] = device::cast<T>(apply_activation_f32<kAct>(gate_f32) * up_f32);
  }
  device::store_as<vec_t>(params.out, out, output_offset);
  PDLTriggerSecondary<kUsePDL>();
}

struct UnaryActivationParams {
  const void* __restrict__ input;
  void* __restrict__ out;
  uint32_t num_vecs;
};

template <typename T, ActivationKind kAct, bool kUsePDL>
__global__ void act_kernel(const __grid_constant__ UnaryActivationParams params) {
  using namespace device;
  constexpr auto kVecSize = kMaxVecBytes / sizeof(T);
  using vec_t = AlignedVector<T, kMaxVecBytes / sizeof(T)>;
  const auto vec_id = blockIdx.x * blockDim.x + threadIdx.x;
  if (vec_id >= params.num_vecs) return;
  PDLWaitPrimary<kUsePDL>();
  const auto in = device::load_as<vec_t>(params.input, vec_id);
  vec_t out;
#pragma unroll
  for (int i = 0; i < kVecSize; ++i) {
    out[i] = device::cast<T>(apply_activation_f32<kAct>(device::cast<fp32_t>(in[i])));
  }
  device::store_as<vec_t>(params.out, out, vec_id);
  PDLTriggerSecondary<kUsePDL>();
}

template <typename T, bool kUsePDL>
struct ActivationKernel {
  static constexpr auto kVecSize = device::kMaxVecBytes / sizeof(T);
  static constexpr auto kBlockSize = 256u;

  using kernel_fn_t = decltype(&act_and_mul_kernel<T, ActivationKind::kSiLU, kUsePDL, false>);
  using unary_kernel_fn_t = decltype(&act_kernel<T, ActivationKind::kReLU2, kUsePDL>);

  template <ActivationKind kAct, bool kFilterExpert>
  static constexpr kernel_fn_t activation_kernel = act_and_mul_kernel<T, kAct, kUsePDL, kFilterExpert>;

  static_assert(device::kMaxVecBytes % sizeof(T) == 0, "unsupported data type");

  template <bool kFilterExpert>
  static kernel_fn_t select_kernel(const std::string& type) {
    using namespace host;
    if (type == "silu") {
      return activation_kernel<ActivationKind::kSiLU, kFilterExpert>;
    } else if (type == "gelu") {
      return activation_kernel<ActivationKind::kGELU, kFilterExpert>;
    } else if (type == "gelu_tanh") {
      return activation_kernel<ActivationKind::kGELUTanh, kFilterExpert>;
    } else {
      Panic("unsupported activation type: ", type);
    }
    return nullptr;
  }

  static void launch(
      const tvm::ffi::TensorView& input,
      const tvm::ffi::TensorView& out,
      const std::string& type,
      const int32_t* expert_ids,
      uint32_t expert_step) {
    using namespace host;

    auto N = SymbolicSize{"num_tokens"};
    auto D_in = SymbolicSize{"input_width"};
    auto D_out = SymbolicSize{"output_width"};
    auto device_ = SymbolicDevice{};
    device_.set_options<kDLCUDA>();

    TensorMatcher({N, D_out})  //
        .with_dtype<T>()
        .with_device(device_)
        .verify(out);
    TensorMatcher({N, D_in})  //
        .with_dtype<T>()
        .with_device(device_)
        .verify(input);

    const auto hidden_size = static_cast<uint32_t>(D_out.unwrap());
    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
    const auto device = device_.unwrap();
    if (num_tokens == 0) return;
    RuntimeCheck(hidden_size * 2 == D_in.unwrap(), "invalid activation dimension");
    RuntimeCheck(hidden_size % kVecSize == 0, "hidden size must be divisible by vector size");
    // only get once to avoid overhead
    const auto num_total_items = num_tokens * (hidden_size / kVecSize);
    RuntimeCheck(num_total_items <= std::numeric_limits<uint32_t>::max(), "too many items for 32-bit indexing");
    const auto num_blocks = div_ceil(static_cast<uint32_t>(num_total_items), kBlockSize);
    const auto params = ActivationParams{
        .input = input.data_ptr(),
        .out = out.data_ptr(),
        .hidden_dim = hidden_size,
        .num_tokens = num_tokens,
        .expert_ids = expert_ids,
        .expert_step = expert_step,
    };
    if (expert_ids != nullptr) {
      RuntimeCheck(expert_step > 0, "expert_step must be positive");
      const auto kernel = select_kernel<true>(type);
      LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
    } else {
      const auto kernel = select_kernel<false>(type);
      LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
    }
  }

  static void run_activation(const tvm::ffi::TensorView input, const tvm::ffi::TensorView out, std::string type) {
    launch(input, out, type, /*expert_ids=*/nullptr, /*expert_step=*/1);
  }

  static void run_activation_filtered(
      const tvm::ffi::TensorView input,
      const tvm::ffi::TensorView out,
      const tvm::ffi::TensorView expert_ids,
      int64_t expert_step,
      std::string type) {
    using namespace host;
    RuntimeCheck(is_type<int32_t>(expert_ids.dtype()), "expert_ids must have dtype int32");
    RuntimeCheck(expert_step >= 1, "expert_step must be positive");
    launch(input, out, type, static_cast<const int32_t*>(expert_ids.data_ptr()), static_cast<uint32_t>(expert_step));
  }

  template <ActivationKind kAct>
  static constexpr auto unary_kernel = act_kernel<T, kAct, kUsePDL>;

  // Use the explicit non-const function-pointer type (mirrors select_kernel's
  // kernel_fn_t) rather than a trailing `decltype(unary_kernel<...>)` return,
  // which deduces a const-qualified pointer that clang-HIP (gfx942) refuses to
  // initialize from an lvalue / nullptr. nvcc accepts both; this form works for
  // CUDA and ROCm alike.
  static unary_kernel_fn_t select_unary_kernel(const std::string& type) {
    using namespace host;
    if (type == "relu2") {
      return ActivationKernel::template unary_kernel<ActivationKind::kReLU2>;
    } else {
      Panic("unsupported unary activation type: ", type);
    }
    return nullptr;
  }

  static void run_unary_activation(const tvm::ffi::TensorView input, const tvm::ffi::TensorView out, std::string type) {
    using namespace host;

    auto N = SymbolicSize{"num_tokens"};
    auto D = SymbolicSize{"hidden"};
    auto device_ = SymbolicDevice{};
    device_.set_options<kDLCUDA>();

    TensorMatcher({N, D})  //
        .with_dtype<T>()
        .with_device(device_)
        .verify(out)
        .verify(input);

    const auto num_elems = static_cast<int64_t>(N.unwrap()) * D.unwrap();
    const auto device = device_.unwrap();
    if (num_elems == 0) return;
    RuntimeCheck(num_elems % kVecSize == 0, "num elements must be divisible by vector size");
    const auto num_vecs = num_elems / kVecSize;
    RuntimeCheck(num_vecs <= std::numeric_limits<uint32_t>::max(), "too many items for 32-bit indexing");
    const auto num_blocks = div_ceil(static_cast<uint32_t>(num_vecs), kBlockSize);
    const auto params = UnaryActivationParams{
        .input = input.data_ptr(),
        .out = out.data_ptr(),
        .num_vecs = static_cast<uint32_t>(num_vecs),
    };
    const auto kernel = select_unary_kernel(type);
    LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
  }
};

}  // namespace