/*
 * SPDX-License-Identifier: Apache-2.0
 */

#include <algorithm>
#include <optional>

#include "onnx_lib/defs/function.h"
#include "onnx_proto/onnx_alias.h"

namespace ONNX_LIGHT_NAMESPACE {

// ---------------------------------------------------------------------------
// FlexAttention Constants
// ---------------------------------------------------------------------------

// Expected number of inputs for modifier subgraphs: (tensor_in)
static constexpr int kModifierInputCount = 1;

static constexpr const char *FlexAttention_ver1_doc = R"DOC(
Computes scaled dot-product attention over rank-4 (batched, multi-head) inputs,
with optional user-provided customization subgraphs at two stages:

1. score_mod: Modify the attention score tensor after Q·K^T
2. prob_mod: Modify the probability tensor after Softmax

This operator mirrors the capabilities of PyTorch's flex_attention:
https://docs.pytorch.org/docs/stable/nn.attention.flex_attention.html

Input Shapes (MUST be rank-4 tensors):
- Q: `(batch_size, q_num_heads, q_sequence_length, head_size)`
- K: `(batch_size, kv_num_heads, kv_sequence_length, head_size)`
- V: `(batch_size, kv_num_heads, kv_sequence_length, v_head_size)`

Output Shape:
- Y: `(batch_size, q_num_heads, q_sequence_length, v_head_size)`

FlexAttention Computation:
```
Scores = (Q @ K^T) * scale
Scores = score_mod(Scores)             # if 'score_mod' is provided
Probs = Softmax(Scores, axis=-1)
Probs = prob_mod(Probs)                # if 'prob_mod' is provided
Y = Probs @ V
```

Grouped Query Attention (GQA):
When `q_num_heads != kv_num_heads`, each K/V head is shared by a contiguous
group of query heads in head-index order. Let
`group_size = q_num_heads / kv_num_heads`; then query head `h` uses K/V head
`floor(h / group_size)`. `q_num_heads` must be a multiple of
`kv_num_heads`.

Modifier Subgraphs (score_mod, prob_mod):
Each modifier subgraph takes exactly one rank-4 tensor input and must produce
exactly one rank-4 tensor output of the same shape and element type.
- score_mod input/output shape: `(batch_size, q_num_heads, q_sequence_length, kv_sequence_length)`
- prob_mod  input/output shape: `(batch_size, q_num_heads, q_sequence_length, kv_sequence_length)`
The element type is determined by softmax_precision (defaults to float32 for
non-double inputs, otherwise double).

Masking can be expressed in score_mod by writing masked positions as -inf (or a
large negative value appropriate for the target precision).
)DOC";

// ---------------------------------------------------------------------------
// Type Utilities
// ---------------------------------------------------------------------------

// Returns true if the element type supports softmax computation.
// Softmax requires floating-point types for numerical stability.
static constexpr bool IsValidSoftmaxElementType(int32_t elem_type) {
  switch (elem_type) {
  case TensorProto::FLOAT:
  case TensorProto::FLOAT16:
  case TensorProto::BFLOAT16:
  case TensorProto::DOUBLE:
    return true;
  default:
    return false;
  }
}

// Determines the element type for softmax computation.
// Priority:
//   1. Use explicit softmax_precision attribute if provided
//   2. Default to float32 for all non-double input types
//   3. Keep double as double
static int32_t GetSoftmaxElementType(int32_t input_elem_type,
                                     const AttributeProto *softmax_precision_attr) {
  // Honor explicit precision setting
  if (softmax_precision_attr != nullptr) {
    return static_cast<int32_t>(softmax_precision_attr->i());
  }
  return input_elem_type == TensorProto::DOUBLE ? TensorProto::DOUBLE : TensorProto::FLOAT;
}

// ---------------------------------------------------------------------------
// Modifier Graph Validation
// ---------------------------------------------------------------------------

// Validates a modifier graph attribute for function builder context.
// Returns true if the graph is valid or not provided.
static bool ValidateModifierGraphForBuilder(const AttributeProto *attr) {
  if (attr == nullptr) {
    return true; // Optional attribute not provided.
  }
  if (!attr->has_g()) {
    return false; // Attribute exists but is not a graph.
  }
  const auto &graph = attr->g();
  return graph.input_size() == kModifierInputCount && graph.output_size() == 1;
}

// Checks if a modifier graph is a trivial identity (passes through first input unchanged).
static bool IsIdentityModifierGraph(const AttributeProto *attr) {
  if (attr == nullptr || !attr->has_g()) {
    return false;
  }
  const auto &graph = attr->g();
  return graph.node_size() == 0 && graph.input_size() == kModifierInputCount &&
         graph.output_size() == 1 && graph.input(0).name().sv() == graph.output(0).name().sv();
}

// ---------------------------------------------------------------------------
// Shape Inference
// ---------------------------------------------------------------------------

// Validates a modifier graph attribute during shape inference.
// Checks I/O counts, element types, and scalar requirements.
static void ValidateModifierGraph(const AttributeProto *attr, const std::string &attr_name,
                                  std::optional<int32_t> expected_output_elem_type) {
  if (!attr) {
    return;
  }
  if (!attr->has_g()) {
    fail_shape_inference("Attribute ", attr_name, " must be a graph.");
  }

  const auto &graph = attr->g();

  // Validate I/O counts.
  if (graph.input_size() != kModifierInputCount) {
    fail_shape_inference("Attribute ", attr_name, " expected ", kModifierInputCount,
                         " input but graph has ", graph.input_size(), ".");
  }
  if (graph.output_size() != 1) {
    fail_shape_inference("Attribute ", attr_name, " must have exactly one output.");
  }

  // Helper: validate element type.
  auto validate_dtype = [&](const ValueInfoProto &vi, int32_t expected, const char *desc) {
    if (!vi.has_type() || !vi.type().has_tensor_type())
      return;
    if (vi.type().tensor_type().elem_type() != expected) {
      fail_shape_inference("Attribute ", attr_name, " ", desc, " type mismatch. Expected ",
                           expected, ", got ", vi.type().tensor_type().elem_type(), ".");
    }
  };

  // Validate input element types.
  if (expected_output_elem_type.has_value() && graph.input_size() >= 1) {
    validate_dtype(graph.input(0), expected_output_elem_type.value(), "input(0)");
  }

  // Validate output type and shape.
  if (graph.output(0).has_type()) {
    const auto &type = graph.output(0).type();
    if (!type.has_tensor_type()) {
      fail_shape_inference("Attribute ", attr_name, " output must be a tensor.");
    }
    const auto &tensor = type.tensor_type();
    if (expected_output_elem_type.has_value() &&
        tensor.elem_type() != expected_output_elem_type.value()) {
      fail_shape_inference("Attribute ", attr_name, " output type mismatch. Expected ",
                           expected_output_elem_type.value(), ", got ", tensor.elem_type(), ".");
    }
  }

  // Validate modifier tensor ranks and compatibility when specified.
  if (graph.input(0).has_type() && graph.input(0).type().has_tensor_type()) {
    const auto &input_tensor = graph.input(0).type().tensor_type();
    if (input_tensor.has_shape() && input_tensor.shape().dim_size() != 4) {
      fail_shape_inference("Attribute ", attr_name,
                           " input must be rank-4 tensor when shape is specified.");
    }
  }
  if (graph.output(0).has_type() && graph.output(0).type().has_tensor_type()) {
    const auto &output_tensor = graph.output(0).type().tensor_type();
    if (output_tensor.has_shape() && output_tensor.shape().dim_size() != 4) {
      fail_shape_inference("Attribute ", attr_name,
                           " output must be rank-4 tensor when shape is specified.");
    }
  }
  if (graph.input(0).has_type() && graph.output(0).has_type() &&
      graph.input(0).type().has_tensor_type() && graph.output(0).type().has_tensor_type()) {
    const auto &in_shape = graph.input(0).type().tensor_type().shape();
    const auto &out_shape = graph.output(0).type().tensor_type().shape();
    const int dim_count = std::min(in_shape.dim_size(), out_shape.dim_size());
    for (int i = 0; i < dim_count; ++i) {
      if (in_shape.dim(i).has_dim_value() && out_shape.dim(i).has_dim_value() &&
          in_shape.dim(i).dim_value() != out_shape.dim(i).dim_value()) {
        fail_shape_inference("Attribute ", attr_name,
                             " output shape must match input shape when specified.");
      }
    }
  }
}

static void FlexAttentionShapeInference(InferenceContext &ctx) {
  // Require Q, K, V inputs.
  if (ctx.getNumInputs() < 3) {
    fail_type_inference("FlexAttention requires query, key, and value inputs.");
  }

  const auto *q_type = ctx.getInputType(0);
  const auto *k_type = ctx.getInputType(1);
  const auto *v_type = ctx.getInputType(2);
  if (!q_type || !k_type || !v_type) {
    return;
  }
  if (!q_type->has_tensor_type() || !k_type->has_tensor_type() || !v_type->has_tensor_type()) {
    fail_type_inference("Inputs 0, 1, and 2 are required tensor types.");
  }

  // Validate element type consistency.
  const auto q_elem_type = q_type->tensor_type().elem_type();
  const auto k_elem_type = k_type->tensor_type().elem_type();
  const auto v_elem_type = v_type->tensor_type().elem_type();
  if (q_elem_type != k_elem_type || q_elem_type != v_elem_type) {
    fail_type_inference("Inputs query, key, and value must have the same element type.");
  }

  // Set output element type.
  auto *output_type = ctx.getOutputType(0)->mutable_tensor_type();
  output_type->set_elem_type(q_elem_type);

  // Validate input ranks.
  const auto &q_shape = q_type->tensor_type().shape();
  const auto &k_shape = k_type->tensor_type().shape();
  const auto &v_shape = v_type->tensor_type().shape();
  if (q_shape.dim_size() != 4 || k_shape.dim_size() != 4 || v_shape.dim_size() != 4) {
    fail_shape_inference("FlexAttention requires rank-4 inputs: "
                         "Q (B, Hq, L, Dqk), K (B, Hkv, S, Dqk), V (B, Hkv, S, Dv).");
  }

  // Validate K and V share the same number of heads.
  if (k_shape.dim(1).has_dim_value() && v_shape.dim(1).has_dim_value() &&
      k_shape.dim(1).dim_value() != v_shape.dim(1).dim_value()) {
    fail_shape_inference("Key and value must share the same head dimension.");
  }

  // Validate K and V share the same sequence length.
  if (k_shape.dim(2).has_dim_value() && v_shape.dim(2).has_dim_value() &&
      k_shape.dim(2).dim_value() != v_shape.dim(2).dim_value()) {
    fail_shape_inference("Key and value must share the same sequence length.");
  }

  // Validate Q and K share the same embedding dimension.
  if (q_shape.dim(3).has_dim_value() && k_shape.dim(3).has_dim_value() &&
      q_shape.dim(3).dim_value() != k_shape.dim(3).dim_value()) {
    fail_shape_inference("Query and key must share the same embedding dimension.");
  }

  // Validate Grouped Query Attention (GQA) configuration.
  if (q_shape.dim(1).has_dim_value() && k_shape.dim(1).has_dim_value()) {
    const auto hq = q_shape.dim(1).dim_value();
    const auto hkv = k_shape.dim(1).dim_value();
    if (hq != hkv && (hkv <= 0 || (hq % hkv) != 0)) {
      fail_shape_inference("q_num_heads must be a multiple of kv_num_heads when they differ.");
    }
  }

  // Set output shape: (B, Hq, L, Dv).
  auto *output_shape = output_type->mutable_shape();
  output_shape->clear_dim();
  *output_shape->add_dim() = q_shape.dim(0); // Batch
  *output_shape->add_dim() = q_shape.dim(1); // Query heads
  *output_shape->add_dim() = q_shape.dim(2); // Query sequence length
  *output_shape->add_dim() = v_shape.dim(3); // Value head size

  // Merge batch dimension info from all inputs.
  mergeInDimensionInfo(k_shape.dim(0), *output_shape->mutable_dim(0), 0);
  mergeInDimensionInfo(v_shape.dim(0), *output_shape->mutable_dim(0), 0);
  if (q_shape.dim(1).has_dim_value() && k_shape.dim(1).has_dim_value() &&
      q_shape.dim(1).dim_value() == k_shape.dim(1).dim_value()) {
    mergeInDimensionInfo(k_shape.dim(1), *output_shape->mutable_dim(1), 1);
  }

  // Validate softmax precision.
  const int32_t softmax_elem_type =
      GetSoftmaxElementType(q_elem_type, ctx.getAttribute("softmax_precision"));
  if (!IsValidSoftmaxElementType(softmax_elem_type)) {
    fail_type_inference(
        "softmax_precision must be specified when inputs are not float/float16/bfloat16/double.");
  }

  // Validate modifier graphs.
  ValidateModifierGraph(ctx.getAttribute("score_mod"), "score_mod", softmax_elem_type);
  ValidateModifierGraph(ctx.getAttribute("prob_mod"), "prob_mod", softmax_elem_type);
}

// ---------------------------------------------------------------------------
// Function Body Builder
// ---------------------------------------------------------------------------

// Builds the function body for FlexAttention.
// The following pattern is applied:
//
//      Q          K          V
//      |          |          |
//      |       Transpose     |
//      |          |          |
//      ---MatMul---          |
//            |               |
//         * scale            |
//            |               |
//      [score_mod]           |
//            |               |
//         Softmax            |
//            |               |
//       [prob_mod]           |
//            |               |
//            -----MatMul------
//                   |
//                   Y
//
static bool BuildFlexAttentionFunctionBody(const FunctionBodyBuildContext &ctx,
                                           const OpSchema &schema, FunctionProto &functionProto) {
  // Validate and extract input types
  const auto *q_type = ctx.getInputType(0);
  const auto *k_type = ctx.getInputType(1);
  const auto *v_type = ctx.getInputType(2);

  if (!q_type || !k_type || !v_type || !q_type->has_tensor_type() || !k_type->has_tensor_type() ||
      !v_type->has_tensor_type()) {
    return false;
  }

  const int64_t input_elem_type = q_type->tensor_type().elem_type();
  if (k_type->tensor_type().elem_type() != input_elem_type ||
      v_type->tensor_type().elem_type() != input_elem_type) {
    return false;
  }

  // Determine softmax precision
  const int32_t softmax_precision = GetSoftmaxElementType(static_cast<int32_t>(input_elem_type),
                                                          ctx.getAttribute("softmax_precision"));
  if (!IsValidSoftmaxElementType(softmax_precision)) {
    return false;
  }

  const auto *scale_attr = ctx.getAttribute("scale");
  const float scale = scale_attr ? scale_attr->f() : 1.0f;

  // Validate modifier graphs
  const auto *score_mod_attr = ctx.getAttribute("score_mod");
  const auto *prob_mod_attr = ctx.getAttribute("prob_mod");

  if (!ValidateModifierGraphForBuilder(score_mod_attr) ||
      !ValidateModifierGraphForBuilder(prob_mod_attr)) {
    return false;
  }

  // Check which modifier graphs are needed.
  const bool need_score_mod = score_mod_attr && !IsIdentityModifierGraph(score_mod_attr);
  const bool need_prob_mod = prob_mod_attr && !IsIdentityModifierGraph(prob_mod_attr);

  // Build the function graph
  constexpr int64_t kFloat32 = TensorProto_DataType_FLOAT;

  FunctionBuilder builder(functionProto);

  // Extract shape information.
  builder.Add("BatchSize = Shape <start = 0, end = 1> (Q)")
      .Add("KVSeqLen = Shape <start = -2, end = -1> (K)");

  builder.Add("QNumHeads = Shape <start = 1, end = 2> (Q)")
      .Add("KVNumHeads = Shape <start = 1, end = 2> (K)");

  // Handle Grouped Query Attention (GQA) by sharing each K/V head across
  // contiguous query-head groups when needed.
  builder.Add("KVRepeat = Div (QNumHeads, KVNumHeads)")
      .Const1D("Axis2", static_cast<int64_t>(2))
      // K: [B, Hkv, S, E] -> [B, Hq, S, E]
      .Add("KUnsqueezed = Unsqueeze(K, Axis2)")
      .Add("KHeadSize1D = Shape <start = 3, end = 4> (K)")
      .Add("KExpandShape = Concat <axis = 0> (BatchSize, KVNumHeads, KVRepeat, KVSeqLen, "
           "KHeadSize1D)")
      .Add("KExpanded = Expand(KUnsqueezed, KExpandShape)")
      .Add("KAlignedShape = Concat <axis = 0> (BatchSize, QNumHeads, KVSeqLen, KHeadSize1D)")
      .Add("KAligned = Reshape(KExpanded, KAlignedShape)")
      // V: same transformation
      .Add("VUnsqueezed = Unsqueeze(V, Axis2)")
      .Add("VHeadSize1D = Shape <start = 3, end = 4> (V)")
      .Add("VExpandShape = Concat <axis = 0> (BatchSize, KVNumHeads, KVRepeat, KVSeqLen, "
           "VHeadSize1D)")
      .Add("VExpanded = Expand(VUnsqueezed, VExpandShape)")
      .Add("VAlignedShape = Concat <axis = 0> (BatchSize, QNumHeads, KVSeqLen, VHeadSize1D)")
      .Add("VAligned = Reshape(VExpanded, VAlignedShape)");

  // Compute scaling factor.
  if (scale_attr != nullptr) {
    builder.Const("ScaleFactorF32", ToTensor<float>(scale));
  } else {
    builder.Add("QShapeAll = Shape(Q)")
        .Const("Idx3Head", ToTensor<int64_t>(3))
        .Add("QKHeadSize = Gather <axis = 0> (QShapeAll, Idx3Head)")
        .Add("QKHeadSizeF = Cast (QKHeadSize)", "to", kFloat32)
        .Add("SqrtHeadSize = Sqrt(QKHeadSizeF)")
        .Const("OneF", ToTensor<float>(1.0f))
        .Add("CalculatedScale = Div(OneF, SqrtHeadSize)")
        .Add("ScaleFactorF32 = Identity(CalculatedScale)");
  }
  builder.Add("ScaleFactorSqrtF32 = Sqrt(ScaleFactorF32)")
      .Add("ScaleFactorSqrt = Cast (ScaleFactorSqrtF32)", "to", input_elem_type);

  // Compute attention scores by pre-scaling Q and K with sqrt(scale).
  builder.Add("QScaled = Mul(Q, ScaleFactorSqrt)")
      .Add("KScaled = Mul(KAligned, ScaleFactorSqrt)")
      .Add("KTranspose = Transpose <perm = [0, 1, 3, 2]> (KScaled)")
      .Add("Score = MatMul(QScaled, KTranspose)")
      .Add("SoftmaxCast = Cast (Score)", "to", softmax_precision);

  // Apply score_mod.
  if (need_score_mod) {
    builder.AddInlinedCall({"ScoreAfterScoreMod"}, score_mod_attr->g(), {"SoftmaxCast"},
                           "ScoreMod_");
  } else {
    builder.Add("ScoreAfterScoreMod = Identity(SoftmaxCast)");
  }

  // Apply Softmax.
  builder.Add("Prob = Softmax <axis = 3> (ScoreAfterScoreMod)");

  // Apply prob_mod.
  if (need_prob_mod) {
    builder.AddInlinedCall({"ProbAfterProbMod"}, prob_mod_attr->g(), {"Prob"}, "ProbMod_");
  } else {
    builder.Add("ProbAfterProbMod = Identity(Prob)");
  }

  // Compute final output: Prob @ V.
  if (static_cast<int64_t>(softmax_precision) != input_elem_type) {
    builder.Add("VSp = Cast (VAligned)", "to", softmax_precision)
        .Add("YSp = MatMul (ProbAfterProbMod, VSp)")
        .Add("Y = Cast (YSp)", "to", input_elem_type);
  } else {
    builder.Add("Y = MatMul (ProbAfterProbMod, VAligned)");
  }

  // Finalize function.
  schema.BuildFunction(functionProto);

  return true;
}

// ---------------------------------------------------------------------------
// Operator Schema Registration
// ---------------------------------------------------------------------------

ONNX_PREVIEW_OPERATOR_SET_SCHEMA(
    FlexAttention, 1,
    OpSchema()
        .SetDoc(FlexAttention_ver1_doc)
        .Input(0, "Q", "Query tensor with shape `(batch_size, q_num_heads, q_seq_len, head_size)`.",
               "T1")
        .Input(1, "K", "Key tensor with shape `(batch_size, kv_num_heads, kv_seq_len, head_size)`.",
               "T1")
        .Input(2, "V",
               "Value tensor with shape `(batch_size, kv_num_heads, kv_seq_len, v_head_size)`.",
               "T1")
        .Output(0, "Y",
                "Output tensor with shape `(batch_size, q_num_heads, q_seq_len, v_head_size)`.",
                "T1")
        .Attr("scale", "Scaling factor for Q*K^T. Defaults to 1/sqrt(head_size).",
              AttributeProto::FLOAT, OPTIONAL_VALUE)
        .Attr("softmax_precision",
              "Floating-point precision for softmax computation. "
              "Defaults to float32 for non-double inputs, otherwise uses double. "
              "Must be explicitly specified for non-float types.",
              AttributeProto::INT, OPTIONAL_VALUE)
        .Attr("score_mod",
              "Optional score modifier subgraph with 1 rank-4 tensor input and 1 rank-4 tensor "
              "output "
              "of the same shape and element type: (scores) -> scores_out. "
              "scores has softmax_precision element type and shape (B, Hq, L, S). "
              "The output must preserve the input shape.",
              AttributeProto::GRAPH, OPTIONAL_VALUE)
        .Attr("prob_mod",
              "Optional probability modifier subgraph with 1 rank-4 tensor input and 1 rank-4 "
              "tensor output "
              "of the same shape and element type: (probs) -> probs_out. "
              "probs has softmax_precision element type and shape (B, Hq, L, S). "
              "The output must preserve the input shape.",
              AttributeProto::GRAPH, OPTIONAL_VALUE)
        .TypeConstraint("T1", OpSchema::all_float_types_ir4(),
                        "Constrain Q, K, V to float tensors.")
        .TypeAndShapeInferenceFunction(FlexAttentionShapeInference)
        .SetSupportLevel(OpSchema::SupportType::EXPERIMENTAL)
        .SetNodeDeterminism(OpSchema::NodeDeterminism::Deterministic)
        .SetContextDependentFunctionBodyBuilder(BuildFlexAttentionFunctionBody));

} // namespace ONNX_LIGHT_NAMESPACE
