// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "core/providers/webgpu/shader_helper.h"
#include "core/providers/webgpu/webgpu_supported_types.h"
#include "contrib_ops/webgpu/webgpu_contrib_kernels.h"
#include "contrib_ops/webgpu/bert/rotary_embedding.h"

namespace onnxruntime {
namespace contrib {
namespace webgpu {

ONNX_OPERATOR_KERNEL_EX(
    RotaryEmbedding,
    kMSDomain,
    1,
    kWebGpuExecutionProvider,
    (*KernelDefBuilder::Create())
        .TypeConstraint("T", WebGpuSupportedFloatTypes())
        .TypeConstraint("M", DataTypeImpl::GetTensorType<int64_t>()),
    RotaryEmbedding);

Status RotaryEmbeddingProgram::GenerateShaderCode(ShaderHelper& shader) const {
  const auto& input = shader.AddInput("input", ShaderUsage::UseUniform);
  // The second input is either seqlens (use_seqlens_for_position_) or position_ids (legacy path).
  // Declared here so the input order matches the caller's AddInputs order:
  // [input, seqlens|position_ids, cos_cache, sin_cache].
  const auto& position_ids_or_seqlens = use_seqlens_for_position_
                                            ? shader.AddInput("seqlens", ShaderUsage::UseUniform)
                                            : shader.AddInput("position_ids", ShaderUsage::UseUniform);
  const auto& cos_cache = shader.AddInput("cos_cache", ShaderUsage::UseUniform);
  const auto& sin_cache = shader.AddInput("sin_cache", ShaderUsage::UseUniform);
  const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform);
  const auto interleaved_str = interleaved_ ? "true" : "false";
  if (use_seqlens_for_position_) {
    // Seqlens path (GQA): inputs are [input, seqlens, cos_cache, sin_cache].
    // Compute per-batch past_seqlen from seqlens[batch_idx] = total_seqlen - 1.
    shader.MainFunctionBody() << "  let half_rotary_emb_dim = uniforms.cos_cache_shape[1];\n"
                                 "  let bsnh = global_idx / uniforms.global_stride % uniforms.global_shape;\n"
                                 "  let size = uniforms.global_shape[0] * uniforms.global_stride[0];\n"
                                 "  if (global_idx >= size) { return; }\n"
                                 "  if (bsnh[3] < half_rotary_emb_dim) {\n"
                                 "    let batch_idx = bsnh[0];\n"
                              << "    let seqlen_i = " << position_ids_or_seqlens.GetByOffset("batch_idx") << ";\n"
                              << "    let seqlen = u32(seqlen_i);\n"
                                 "    let total_seqlen = seqlen + 1u;\n"
                                 "    let past_seqlen = total_seqlen - uniforms.global_shape[1];\n"
                                 "    let position_id = past_seqlen + bsnh[1];\n"
                              << "    let i = dot(bsnh, uniforms.input_output_stride) + select(0u, bsnh[3], " << interleaved_str << ");\n"
                              << "    let j = i + select(half_rotary_emb_dim, 1u, " << interleaved_str << ");\n"
                                                                                                          "    let max_position = uniforms.cos_cache_shape[0];\n"
                                                                                                          "    if (position_id >= max_position) {\n"
                              << "      " << output.SetByOffset("i", input.GetByOffset("i")) << "\n"
                              << "      " << output.SetByOffset("j", input.GetByOffset("j")) << "\n"
                                                                                                "    } else {\n"
                              << "      let re = " << input.GetByOffset("i") << " * " << cos_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << " - " << input.GetByOffset("j") << " * " << sin_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << ";\n"
                              << "      " << output.SetByOffset("i", "re") << "\n"
                              << "      let im = " << input.GetByOffset("i") << " * " << sin_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << " + " << input.GetByOffset("j") << " * " << cos_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << ";\n"
                              << "      " << output.SetByOffset("j", "im") << "\n"
                                                                              "    }\n"
                              << "  } else {\n"
                                 "    let k = dot(bsnh, uniforms.input_output_stride) + half_rotary_emb_dim;\n"
                              << "    " << output.SetByOffset("k", input.GetByOffset("k")) << "\n"
                              << "  }";
  } else {
    // Original path: inputs are [input, position_ids, cos_cache, sin_cache].
    const auto& position_ids = position_ids_or_seqlens;
    // TODO: remove output_indices.
    const auto& output_indices = shader.AddIndices("output_indices", ShaderUsage::None);
    shader.MainFunctionBody() << "  let half_rotary_emb_dim = uniforms.cos_cache_shape[1];\n"
                                 "  let bsnh = global_idx / uniforms.global_stride % uniforms.global_shape;\n"
                                 "  let size = uniforms.global_shape[0] * uniforms.global_stride[0];\n"
                                 "  if (global_idx >= size) { return; }\n"
                                 "  if (bsnh[3] < half_rotary_emb_dim) {\n"
                              << "    let position_ids_idx = " << position_ids.BroadcastedIndicesToOffset("bsnh.xy", output_indices) << ";\n"
                              << "    let raw_pos = " << position_ids.GetByOffset("position_ids_idx") << ";\n"
                              << "    let i = dot(bsnh, uniforms.input_output_stride) + select(0, bsnh[3], " << interleaved_str << ");\n"
                              << "    let j = i + select(half_rotary_emb_dim, 1, " << interleaved_str << ");\n"
                                                                                                         "    let max_position = uniforms.cos_cache_shape[0];\n"
                                                                                                         // Bounds check: raw_pos < 0 catches negative position_ids (i32 from truncated int64).
                                                                                                         // After u32 conversion + offset, check >= max_position catches too-large values.
                                                                                                         // On OOB, pass through input unchanged (same as CUDA kernel behavior).
                                                                                                         "    if (raw_pos < 0) {\n"
                              << "      " << output.SetByOffset("i", input.GetByOffset("i")) << "\n"
                              << "      " << output.SetByOffset("j", input.GetByOffset("j")) << "\n"
                                                                                                "    } else {\n"
                                                                                                "      let position_id = u32(raw_pos) + select(0, bsnh[1], position_ids_idx == 0);\n"
                                                                                                "      if (position_id >= max_position) {\n"
                              << "        " << output.SetByOffset("i", input.GetByOffset("i")) << "\n"
                              << "        " << output.SetByOffset("j", input.GetByOffset("j")) << "\n"
                                                                                                  "      } else {\n"
                              << "        let re = " << input.GetByOffset("i") << " * " << cos_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << " - " << input.GetByOffset("j") << " * " << sin_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << ";\n"
                              << "        " << output.SetByOffset("i", "re") << "\n"
                              << "        let im = " << input.GetByOffset("i") << " * " << sin_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << " + " << input.GetByOffset("j") << " * " << cos_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << ";\n"
                              << "        " << output.SetByOffset("j", "im") << "\n"
                                                                                "      }\n"
                                                                                "    }\n"
                              << "  } else { \n"
                                 "    let k = dot(bsnh, uniforms.input_output_stride) + half_rotary_emb_dim;\n"
                              << "    " << output.SetByOffset("k", input.GetByOffset("k")) << "\n"
                              << "  }";
  }

  return Status::OK();
}

Status FusedQKRotaryEmbeddingProgram::GenerateShaderCode(ShaderHelper& shader) const {
  // Inputs
  const auto& q_input = shader.AddInput("q_input", ShaderUsage::UseUniform);
  const auto& k_input = shader.AddInput("k_input", ShaderUsage::UseUniform);
  const auto& seqlens = shader.AddInput("seqlens", ShaderUsage::UseUniform);
  const auto& cos_cache = shader.AddInput("cos_cache", ShaderUsage::UseUniform);
  const auto& sin_cache = shader.AddInput("sin_cache", ShaderUsage::UseUniform);
  // Outputs
  const auto& q_output = shader.AddOutput("q_output", ShaderUsage::UseUniform);
  const auto& k_output = shader.AddOutput("k_output", ShaderUsage::UseUniform);

  const auto interleaved_str = interleaved_ ? "true" : "false";

  shader.MainFunctionBody()
      << "  if (global_idx >= uniforms.q_domain_size) { return; }\n"
      << "  let half_rotary_dim = uniforms.cos_cache_shape[1];\n"
      << "  let bsnh = global_idx / uniforms.q_global_stride % uniforms.q_global_shape;\n"
      << "  if (bsnh[3] < half_rotary_dim) {\n"
      << "    let batch_idx = bsnh[0];\n"
      << "    let sequence_idx = bsnh[1];\n"
      << "    let seqlen_i = " << seqlens.GetByOffset("batch_idx") << ";\n"
      << "    let seqlen = u32(seqlen_i);\n"
      << "    let total_seqlen = seqlen + 1u;\n"
      << "    let past_seqlen = total_seqlen - uniforms.q_global_shape[1];\n"
      // position_id is derived from past_seqlen + sequence_idx (always non-negative).
      << "    let position_id = past_seqlen + sequence_idx;\n"
      << "    let qi = dot(bsnh, uniforms.q_input_output_stride) + select(0u, bsnh[3], " << interleaved_str << ");\n"
      << "    let qj = qi + select(half_rotary_dim, 1u, " << interleaved_str << ");\n"
                                                                                // Bounds check: position_id must be within cos/sin cache range.
                                                                                // On OOB, pass through input unchanged (same as CUDA kernel behavior).
                                                                                "    let max_position = uniforms.cos_cache_shape[0];\n"
                                                                                "    if (position_id >= max_position) {\n"
      << "      " << q_output.SetByOffset("qi", q_input.GetByOffset("qi")) << "\n"
      << "      " << q_output.SetByOffset("qj", q_input.GetByOffset("qj")) << "\n"
      << "      if (bsnh[2] < uniforms.k_global_shape[2]) {\n"
      << "        let ki = dot(bsnh, uniforms.k_input_output_stride) + select(0u, bsnh[3], " << interleaved_str << ");\n"
      << "        let kj = ki + select(half_rotary_dim, 1u, " << interleaved_str << ");\n"
      << "        " << k_output.SetByOffset("ki", k_input.GetByOffset("ki")) << "\n"
      << "        " << k_output.SetByOffset("kj", k_input.GetByOffset("kj")) << "\n"
                                                                                "      }\n"
                                                                                "    } else {\n"
      << "      let cos_v = " << cos_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << ";\n"
      << "      let sin_v = " << sin_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << ";\n"
      << "      let q_re = " << q_input.GetByOffset("qi") << " * cos_v - " << q_input.GetByOffset("qj") << " * sin_v;\n"
      << "      " << q_output.SetByOffset("qi", "q_re") << "\n"
      << "      let q_im = " << q_input.GetByOffset("qi") << " * sin_v + " << q_input.GetByOffset("qj") << " * cos_v;\n"
      << "      " << q_output.SetByOffset("qj", "q_im") << "\n"
      // Conditionally process Key (only for heads that exist in K domain)
      << "      if (bsnh[2] < uniforms.k_global_shape[2]) {\n"
      << "        let ki = dot(bsnh, uniforms.k_input_output_stride) + select(0u, bsnh[3], " << interleaved_str << ");\n"
      << "        let kj = ki + select(half_rotary_dim, 1u, " << interleaved_str << ");\n"
      << "        let k_re = " << k_input.GetByOffset("ki") << " * cos_v - " << k_input.GetByOffset("kj") << " * sin_v;\n"
      << "        " << k_output.SetByOffset("ki", "k_re") << "\n"
      << "        let k_im = " << k_input.GetByOffset("ki") << " * sin_v + " << k_input.GetByOffset("kj") << " * cos_v;\n"
      << "        " << k_output.SetByOffset("kj", "k_im") << "\n"
                                                             "      }\n"
                                                             "    }\n"
      << "  } else {\n"
      << "    let qk = dot(bsnh, uniforms.q_input_output_stride) + half_rotary_dim;\n"
      << "    " << q_output.SetByOffset("qk", q_input.GetByOffset("qk")) << "\n"
      // Conditionally process Key (only for heads that exist in K domain)
      << "    if (bsnh[2] < uniforms.k_global_shape[2]) {\n"
      << "      let kk = dot(bsnh, uniforms.k_input_output_stride) + half_rotary_dim;\n"
      << "      " << k_output.SetByOffset("kk", k_input.GetByOffset("kk")) << "\n"
      << "    }\n"
      << "  }\n";
  return Status::OK();
}

RotaryEmbedding::RotaryEmbedding(const OpKernelInfo& info) : WebGpuKernel(info) {
  scale_ = info.GetAttrOrDefault<float>("scale", 1.0);
  rotary_embedding_dim_ = static_cast<int>(info.GetAttrOrDefault<int64_t>("rotary_embedding_dim", 0));
  num_heads_ = static_cast<int>(info.GetAttrOrDefault<int64_t>("num_heads", 0));
  interleaved_ = (info.GetAttrOrDefault<int64_t>("interleaved", 0) == 1);
  is_packed_batching_ = (info.GetAttrOrDefault<int64_t>("is_packed_batching", 0) == 1);
}

Status RunRotaryEmbedding(onnxruntime::webgpu::ComputeContext& context,
                          const Tensor* input,
                          const Tensor* position_ids_or_seqlens,
                          const Tensor* cos_cache,
                          const Tensor* sin_cache,
                          Tensor* output,
                          int batch_size,
                          int sequence_length,
                          int hidden_size,
                          int head_size,
                          float scale,
                          bool rotary_interleaved,
                          bool use_seqlens_for_position,
                          const std::vector<uint32_t>& input_output_strides) {
  const auto half_rotary_embedding_dim = onnxruntime::narrow<uint32_t>(cos_cache->Shape()[1]);
  const auto num_heads = hidden_size / head_size;

  // Rotary embeddings are calculated in a pair-wise fashion. Use the shape
  // [batch, sequence, heads, half_rotary_dim_complement] to unfold the global index in shader.
  const TensorShape global_shape({static_cast<int64_t>(batch_size),
                                  static_cast<int64_t>(sequence_length),
                                  static_cast<int64_t>(num_heads),
                                  static_cast<int64_t>(head_size - half_rotary_embedding_dim)});
  const auto rank = global_shape.NumDimensions();
  std::vector<uint32_t> global_dims(rank);
  std::vector<uint32_t> global_strides(rank);
  for (size_t j = 0; j < rank; ++j) {
    global_dims[j] = onnxruntime::narrow<uint32_t>(global_shape[j]);
    global_strides[j] = onnxruntime::narrow<uint32_t>(global_shape.SizeFromDimension(j + 1));
  }
  const auto output_size = onnxruntime::narrow<uint32_t>(global_shape.Size());

  RotaryEmbeddingProgram program(rotary_interleaved, use_seqlens_for_position);
  program
      .CacheHint(rotary_interleaved, use_seqlens_for_position)
      .AddInputs({{input, ProgramTensorMetadataDependency::TypeAndRank},
                  {position_ids_or_seqlens, ProgramTensorMetadataDependency::TypeAndRank},
                  {cos_cache, ProgramTensorMetadataDependency::Rank},
                  {sin_cache, ProgramTensorMetadataDependency::Rank}})
      .AddOutput({output, ProgramTensorMetadataDependency::None})
      .SetDispatchGroupSize((output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE)
      .AddUniformVariables({{scale},
                            {gsl::make_span(global_dims)},
                            {gsl::make_span(global_strides)},
                            {gsl::make_span(input_output_strides)}});
  if (!use_seqlens_for_position) {
    program.AddIndices(TensorShape{1, 1});
  }
  return context.RunProgram(program);
}

Status RotaryEmbedding::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const {
  const auto* input = context.Input<Tensor>(0);
  const auto input_shape = input->Shape();
  const auto* position_ids = context.Input<Tensor>(1);
  const auto* cos_cache = context.Input<Tensor>(2);
  const auto* sin_cache = context.Input<Tensor>(3);
  auto* output = context.Output(0, input_shape);

  const auto batch_size = onnxruntime::narrow<uint32_t>(input->Shape()[0]);
  const auto batch_stride = onnxruntime::narrow<uint32_t>(input_shape.SizeFromDimension(1));
  const auto sequence_length = onnxruntime::narrow<uint32_t>(input_shape[input_shape.NumDimensions() - 2]);
  const auto hidden_size = batch_stride / sequence_length;
  const auto half_rotary_embedding_dim = onnxruntime::narrow<uint32_t>(cos_cache->Shape()[1]);
  const auto head_size = rotary_embedding_dim_ == 0 ? half_rotary_embedding_dim * 2 : hidden_size / num_heads_;

  // position_ids bounds validation is handled by shader-side defense-in-depth checks
  // (OOB position_ids → pass-through input unchanged). Host-side value scanning is not possible
  // because WebGPU program inputs must be GPU buffers (InputMemoryType(OrtMemTypeCPUInput) is
  // incompatible with AddInputs).

  const auto input_output_strides =
      input_shape.NumDimensions() == 3
          ? std::vector<uint32_t>({batch_stride, hidden_size, head_size, 1})
          : (input_shape.NumDimensions() == 4
                 ? std::vector<uint32_t>({batch_stride, head_size, sequence_length * head_size, 1})
                 : std::vector<uint32_t>({}));

  return RunRotaryEmbedding(context, input, position_ids, cos_cache, sin_cache, output,
                            static_cast<int>(batch_size), static_cast<int>(sequence_length),
                            static_cast<int>(hidden_size), static_cast<int>(head_size),
                            scale_, interleaved_, /*use_seqlens_for_position=*/false, input_output_strides);
}

}  // namespace webgpu
}  // namespace contrib
}  // namespace onnxruntime
