// Author: Simon-Pierre Boucher — contact@spboucher.ai
#pragma once

#include "nn/config.h"
#include "nn/linear.h"
#include "ops/ops.h"

#include <cmath>
#include <memory>

namespace forge::nn {

// Interface so linear-attention / sliding-window variants can slot in
// without touching TransformerBlock.
class AttentionBase : public Module {
public:
    virtual ~AttentionBase() = default;
    // x: [B, T, C] → [B, T, C]
    virtual Var forward(const Var& x) const = 0;
};

// Causal multi-head self-attention with GQA and (optional) RoPE.
// Per-layer flavour comes from `layer_idx` + config: NoPE layers skip RoPE
// (SmolLM3), global layers may use a different rope theta (Gemma 3), and
// head_dim may be decoupled from d_model/n_heads (Qwen3). Optional QKV bias
// (Qwen2.5, on q/k/v only — llama convention keeps wo bias-free).
class CausalSelfAttention : public AttentionBase {
public:
    CausalSelfAttention(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng,
                        int64_t layer_idx = 0)
        : n_heads_(cfg.n_heads),
          n_kv_heads_(cfg.n_kv_heads),
          head_dim_(cfg.head_dim()),
          use_rope_(cfg.layer_uses_rope(layer_idx)),
          qk_norm_(cfg.qk_norm),
          norm_eps_(cfg.norm_eps),
          window_(cfg.layer_is_global(layer_idx) ? 0 : cfg.sliding_window),
          attn_softcap_(cfg.attn_softcap) {
        const float std = 0.02f;
        const int64_t C = cfg.d_model;
        const int64_t Cq = cfg.n_heads * head_dim_;
        const int64_t Ckv = cfg.n_kv_heads * head_dim_;
        const bool bias = cfg.attention_bias;
        wq_ = std::make_unique<Linear>(C, Cq, bias, std, rng);
        wk_ = std::make_unique<Linear>(C, Ckv, bias, std, rng);
        wv_ = std::make_unique<Linear>(C, Ckv, bias, std, rng);
        wo_ = std::make_unique<Linear>(Cq, C, false, proj_std, rng); // residual projection
        const ops::QuantMode qm = quant_mode_from(cfg.quant);
        wq_->set_quant(qm);
        wk_->set_quant(qm);
        wv_->set_quant(qm);
        wo_->set_quant(qm);
        absorb("wq", *wq_);
        absorb("wk", *wk_);
        absorb("wv", *wv_);
        absorb("wo", *wo_);
        if (qk_norm_) {
            q_norm_w_ = register_param("q_norm.weight",
                                       Tensor::full({head_dim_}, 1.0f));
            k_norm_w_ = register_param("k_norm.weight",
                                       Tensor::full({head_dim_}, 1.0f));
        }
        if (use_rope_) {
            const bool global = cfg.layer_is_global(layer_idx);
            const float theta = (global && cfg.rope_theta_global > 0.0f)
                                    ? cfg.rope_theta_global
                                    : cfg.rope_theta;
            freqs_ = ops::rope_freqs(head_dim_, theta, cfg.rope_scale_factor,
                                     cfg.rope_scale_low, cfg.rope_scale_high,
                                     cfg.rope_scale_orig_ctx);
        }
    }

    Var forward(const Var& x) const override {
        const int64_t B = x.value().size(0), T = x.value().size(1), C = x.value().size(2);
        const int64_t Cq = n_heads_ * head_dim_;
        const int64_t Ckv = n_kv_heads_ * head_dim_;

        Var x2d = x.reshaped({B * T, C});
        Var q = wq_->forward(x2d).reshaped({B, T, Cq});
        Var k = wk_->forward(x2d).reshaped({B, T, Ckv});
        Var v = wv_->forward(x2d).reshaped({B, T, Ckv});

        if (qk_norm_) {
            // Per-head RMSNorm before RoPE (Qwen3/Gemma3): normalize over
            // head_dim by viewing heads as rows.
            q = ops::rmsnorm(q.reshaped({B * T * n_heads_, head_dim_}), q_norm_w_,
                             norm_eps_).reshaped({B, T, Cq});
            k = ops::rmsnorm(k.reshaped({B * T * n_kv_heads_, head_dim_}), k_norm_w_,
                             norm_eps_).reshaped({B, T, Ckv});
        }

        if (use_rope_) {
            q = ops::rope(q, n_heads_, freqs_, 0);
            k = ops::rope(k, n_kv_heads_, freqs_, 0);
        }

        const float scale = 1.0f / std::sqrt(float(head_dim_));
        Var o = ops::attention(q, k, v, n_heads_, n_kv_heads_, /*causal=*/true, scale,
                               window_, attn_softcap_);
        return wo_->forward(o.reshaped({B * T, Cq})).reshaped({B, T, C});
    }

private:
    int64_t n_heads_, n_kv_heads_, head_dim_;
    bool use_rope_;
    bool qk_norm_;
    float norm_eps_;
    int64_t window_;     // 0 = full attention (global layer or SWA off)
    float attn_softcap_;
    std::unique_ptr<Linear> wq_, wk_, wv_, wo_;
    Var q_norm_w_, k_norm_w_; // defined only when qk_norm
    Tensor freqs_;            // [head_dim/2], defined only when use_rope
};

} // namespace forge::nn