// @nolint
#include "blackwell_fmha_fwd_template.cuh"
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)

std::tuple<at::Tensor, at::Tensor> dispatch_fmha_fwd(
    const at::Tensor& q,
    const at::Tensor& k, // (batch_size, KV_seqlen, num_KV_heads, head_dim) if non-paged or (num_blocks, page_block_size, num_KV_heads, head_dim) if paged
    const at::Tensor& v, // (batch_size, KV_seqlen, num_KV_heads, head_dim) if non-paged or (num_blocks, page_block_size, num_KV_heads, head_dim) if paged
    const std::optional<at::Tensor>& cu_seqlens_q,
    const std::optional<at::Tensor>& cu_seqlens_k,
    std::optional<int64_t> max_seq_len_q,
    std::optional<int64_t> max_seq_len_k,
    std::optional<double> softmax_scale,
    bool causal,
    const std::optional<at::Tensor>& seqlen_kv,
    const std::optional<at::Tensor>& page_table, // dim: (batch_size, max_num_pages_per_seq) , null if non-paged
    std::optional<int64_t> seqlen_k,
    int64_t window_size_left,
    int64_t window_size_right,
    bool bottom_right) {

  bool kIsPaged = false;
  if (page_table && page_table->defined()) {
    kIsPaged = true;
  }

  // Handle local attention parameters
  bool local = (window_size_left >= 0 || window_size_right >= 0);
  if (local) {
    // If causal is enabled, override window_size_right to 0 for causal+local
    // behavior
    if (causal) {
      window_size_right = 0;
      causal = false; // Use local attention instead of causal
    }
    // Expand -1 window sizes to full sequence length if available
    if (window_size_left < 0) {
      TORCH_CHECK(
          max_seq_len_k.has_value(),
          "window_size_left is negative but max_seq_len_k is not provided");
      window_size_left = max_seq_len_k.value();
    }
    if (window_size_right < 0) {
      TORCH_CHECK(
          max_seq_len_k.has_value(),
          "window_size_right is negative but max_seq_len_k is not provided");
      window_size_right = max_seq_len_k.value();
    }
  }
  if (causal){
    window_size_right = 0;
  }

  auto dispatch_fmha = [&](auto element,
                           auto element_out,
                           auto head_dim,
                           auto varlen,
                           auto mask) {
    return fmha_fwd<
        decltype(element),
        decltype(element_out),
        head_dim,
        varlen,
        decltype(mask)>(
        q,
        k,
        v,
        cu_seqlens_q,
        cu_seqlens_k,
        max_seq_len_q,
        max_seq_len_k,
        softmax_scale,
        seqlen_kv,
        page_table,
        seqlen_k,
        window_size_left,
        window_size_right);
  };

  auto dispatch_type = [&](auto varlen, auto mask, auto head_dim) {
    if (q.dtype() == torch::kFloat16) {
      return dispatch_fmha(
          cutlass::half_t{}, cutlass::half_t{}, head_dim, varlen, mask);
    } else if (q.dtype() == torch::kBFloat16) {
      return dispatch_fmha(
          cutlass::bfloat16_t{}, cutlass::bfloat16_t{}, head_dim, varlen, mask);
    } else if (q.dtype() == torch::kFloat8_e4m3fn) {
      // Return BF16 when input is FP8
      return dispatch_fmha(
          cutlass::float_e4m3_t{},
          cutlass::bfloat16_t{},
          head_dim,
          varlen,
          mask);
    }
    TORCH_CHECK(false, "Unsupported dtype for q: ", q.dtype());
  };

  auto dispatch_head_dim = [&](auto varlen, auto mask) {
    if (q.size(q.dim() - 1) == 128) {
      return dispatch_type(varlen, mask, std::integral_constant<int, 128>{});
    } else if (q.size(q.dim() - 1) == 64) {
      return dispatch_type(varlen, mask, std::integral_constant<int, 64>{});
    } else {
      TORCH_CHECK(false, "Unsupported head dim: ", q.size(q.dim() - 1));
    }
  };

  auto dispatch_mask = [&](auto varlen) {
    int seq_k = kIsPaged
        ? (varlen
            ? static_cast<int>(*max_seq_len_k)
            : static_cast<int>(*seqlen_k))
        : (varlen
            ? k.size(0)
            : k.size(1));
    if (causal) {
      if (bottom_right) {
        return dispatch_head_dim(varlen, CausalMask</*kIsQBegin=*/false>{});
      } else {
        return dispatch_head_dim(varlen, CausalMask</*kIsQBegin=*/true>{});
      }
    } else if (local) {
      if (bottom_right) {
          return dispatch_head_dim(varlen, LocalMask</*kIsQBegin=*/false>{});
      } else {
        return dispatch_head_dim(varlen, LocalMask</*kIsQBegin=*/true>{});
      }
    } else if (varlen || seq_k % 128 != 0) {
      // Use the residual mask for varlen or when K seqlen is not multiple of
      // blockN
      return dispatch_head_dim(varlen, ResidualMask{});
    } else {
      return dispatch_head_dim(varlen, NoMask{});
    }
  };

  if (max_seq_len_q.has_value()) {
    return dispatch_mask(std::bool_constant<true>{});
  } else {
    return dispatch_mask(std::bool_constant<false>{});
  }
}

std::tuple<at::Tensor, at::Tensor> dispatch_fmha_fwd_meta(
    const at::Tensor& q,
    const at::Tensor& k, // (batch_size, KV_seqlen, num_KV_heads, head_dim) if non-paged or (num_blocks, page_block_size, num_KV_heads, head_dim) if paged
    const at::Tensor& v, // (batch_size, KV_seqlen, num_KV_heads, head_dim) if non-paged or (num_blocks, page_block_size, num_KV_heads, head_dim) if paged
    const std::optional<at::Tensor>& cu_seqlens_q,
    const std::optional<at::Tensor>& cu_seqlens_k,
    std::optional<c10::SymInt> max_seq_len_q,
    std::optional<c10::SymInt> max_seq_len_k,
    std::optional<double> softmax_scale,
    bool causal,
    const std::optional<at::Tensor>& seqlen_kv,
    const std::optional<at::Tensor>& page_table, // dim: (batch_size, max_num_pages_per_seq) , null if non-paged
    std::optional<c10::SymInt> seqlen_k,
    c10::SymInt window_size_left,
    c10::SymInt window_size_right,
    bool bottom_right) {
  auto output = at::empty_like(q);
  bool k_is_varlen = max_seq_len_q.has_value();
  auto SQ = k_is_varlen ? q.sym_size(0) : q.sym_size(1);
  auto H_Q = k_is_varlen ? q.sym_size(1) : q.sym_size(2);
  auto B = k_is_varlen ? 1 : q.sym_size(0);
  auto logsumexp = q.new_empty_symint({B, H_Q, SQ}, q.options());
  return std::make_tuple(output, logsumexp);
}

// -------------------------------------------------------------------------------------------------
// Op registration
// -------------------------------------------------------------------------------------------------
TORCH_LIBRARY_FRAGMENT(mslk, m) {
  m.def(
      "fmha_fwd("
      "    Tensor query, "
      "    Tensor key, "
      "    Tensor value, "
      "    Tensor? cu_seqlens_q=None, "
      "    Tensor? cu_seqlens_k=None, "
      "    SymInt? max_seq_len_q=None, "
      "    SymInt? max_seq_len_k=None, "
      "    float? softmax_scale=None, "
      "    bool causal=False, "
      "    Tensor? seqlen_kv=None, "
      "    Tensor? page_table=None, "
      "    SymInt? seqlen_k=None, "
      "    SymInt window_size_left=-1, "
      "    SymInt window_size_right=-1, "
      "    bool bottom_right=True"
      ") -> (Tensor, Tensor)");
}

TORCH_LIBRARY_IMPL(mslk, CUDA, m) {
  m.impl("fmha_fwd", dispatch_fmha_fwd);
}
TORCH_LIBRARY_IMPL(mslk, Meta, m) {
  m.impl("fmha_fwd", dispatch_fmha_fwd_meta);
}
#endif // CUTLASS_ARCH_MMA_SM100_SUPPORTED
