#include "attention_internal.h"

namespace engine::modules {

using namespace attention::internal;

SelfAttentionModule::SelfAttentionModule(AttentionConfig config) : config_(config) {
    validate_attention_config(config_);
}

const AttentionConfig & SelfAttentionModule::config() const noexcept {
    return config_;
}

const core::ModuleSchema & SelfAttentionModule::schema() const noexcept {
    return static_schema();
}

core::TensorValue SelfAttentionModule::build(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & input,
    const AttentionWeights & weights) const {
    return build_attention_impl(ctx, input, input, config_, require_attention_weights(weights, config_.use_bias));
}

const core::ModuleSchema & SelfAttentionModule::static_schema() noexcept {
    return kSelfAttentionSchema;
}

StreamingSelfAttentionModule::StreamingSelfAttentionModule(AttentionConfig config) : config_(config) {
    validate_attention_config(config_);
}

const AttentionConfig & StreamingSelfAttentionModule::config() const noexcept {
    return config_;
}

StreamingAttentionOutputs StreamingSelfAttentionModule::build(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & input,
    const core::TensorValue & positions,
    const AttentionWeights & weights,
    const std::optional<core::TensorValue> & prefix_key,
    const std::optional<core::TensorValue> & prefix_value,
    const std::optional<core::TensorValue> & attention_mask) const {
    validate_sequence_input(input, config_.hidden_size, "input");
    if (prefix_key.has_value() != prefix_value.has_value()) {
        throw std::runtime_error("Streaming self-attention requires both prefix_key and prefix_value or neither");
    }
    if (prefix_key.has_value()) {
        core::validate_rank_between(*prefix_key, 4, 4, "prefix_key");
        core::validate_rank_between(*prefix_value, 4, 4, "prefix_value");
        if (prefix_key->shape.dims[0] != input.shape.dims[0] || prefix_value->shape.dims[0] != input.shape.dims[0]) {
            throw std::runtime_error("Streaming self-attention prefix batch size must match input");
        }
    }

    const int64_t head_dim = config_.hidden_size / config_.num_heads;
    const float scale = 1.0f / std::sqrt(static_cast<float>(head_dim));

    const LinearModule q_proj({config_.hidden_size, config_.hidden_size, config_.use_bias});
    const LinearModule k_proj({config_.hidden_size, config_.hidden_size, config_.use_bias});
    const LinearModule v_proj({config_.hidden_size, config_.hidden_size, config_.use_bias});
    const LinearModule out_proj({config_.hidden_size, config_.hidden_size, config_.use_bias});
    const MatMulModule matmul;

    auto q = q_proj.build(ctx, input, make_linear_weights(weights.q_weight, weights.q_bias));
    auto k = k_proj.build(ctx, input, make_linear_weights(weights.k_weight, weights.k_bias));
    auto v = v_proj.build(ctx, input, make_linear_weights(weights.v_weight, weights.v_bias));

    q = reshape_heads(ctx, q, config_.num_heads, head_dim);
    k = reshape_heads(ctx, k, config_.num_heads, head_dim);
    v = reshape_heads(ctx, v, config_.num_heads, head_dim);

    q = build_positions_rope(ctx, q, positions, head_dim);
    k = build_positions_rope(ctx, k, positions, head_dim);

    auto q_heads = permute_tensor(ctx, q, {0, 2, 1, 3});
    auto k_heads = permute_tensor(ctx, k, {0, 2, 1, 3});
    auto v_heads = permute_tensor(ctx, v, {0, 2, 1, 3});

    core::TensorValue context;
    if (prefix_key.has_value()) {
        auto all_k = concat_sequence_axis(ctx, prefix_key, k);
        auto all_v = concat_sequence_axis(ctx, prefix_value, v);
        auto all_k_heads = permute_tensor(ctx, all_k, {0, 2, 1, 3});
        auto all_v_heads = permute_tensor(ctx, all_v, {0, 2, 1, 3});
        auto k_transposed = permute_tensor(ctx, all_k_heads, {0, 1, 3, 2});
        auto scores = matmul.build(ctx, q_heads, k_transposed);
        core::TensorValue attn;
        if (attention_mask.has_value()) {
            attn = core::wrap_tensor(
                ggml_soft_max_ext(ctx.ggml, ensure_contiguous_layout(ctx, scores).tensor, attention_mask->tensor, scale, 0.0f),
                scores.shape,
                GGML_TYPE_F32);
        } else {
            scores = core::wrap_tensor(ggml_scale(ctx.ggml, scores.tensor, scale), scores.shape, GGML_TYPE_F32);
            scores = core::wrap_tensor(
                ggml_diag_mask_inf(ctx.ggml, scores.tensor, static_cast<int>(prefix_key->shape.dims[1])),
                scores.shape,
                GGML_TYPE_F32);
            attn = core::wrap_tensor(ggml_soft_max(ctx.ggml, ensure_contiguous_layout(ctx, scores).tensor), scores.shape, GGML_TYPE_F32);
        }
        context = matmul.build(ctx, attn, all_v_heads);
    } else {
        auto k_transposed = permute_tensor(ctx, k_heads, {0, 1, 3, 2});
        auto scores = matmul.build(ctx, q_heads, k_transposed);
        core::TensorValue attn;
        if (attention_mask.has_value()) {
            attn = core::wrap_tensor(
                ggml_soft_max_ext(ctx.ggml, ensure_contiguous_layout(ctx, scores).tensor, attention_mask->tensor, scale, 0.0f),
                scores.shape,
                GGML_TYPE_F32);
        } else {
            scores = core::wrap_tensor(ggml_scale(ctx.ggml, scores.tensor, scale), scores.shape, GGML_TYPE_F32);
            scores = core::wrap_tensor(ggml_diag_mask_inf(ctx.ggml, scores.tensor, 0), scores.shape, GGML_TYPE_F32);
            attn = core::wrap_tensor(ggml_soft_max(ctx.ggml, ensure_contiguous_layout(ctx, scores).tensor), scores.shape, GGML_TYPE_F32);
        }
        context = matmul.build(ctx, attn, v_heads);
    }

    context = permute_tensor(ctx, context, {0, 2, 1, 3});
    context = ensure_contiguous_layout(ctx, context);
    context = core::reshape_tensor(ctx, context, core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], config_.hidden_size}));

    auto output = out_proj.build(ctx, context, make_linear_weights(weights.out_weight, weights.out_bias));
    return {output, k, v};
}

}  // namespace engine::modules