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

#include "core/tensor.h"

#include <functional>
#include <memory>
#include <vector>

namespace forge {

// A Var is a Tensor plus (lazily allocated) gradient storage. Copies share
// the same impl, so parameters handed to modules, the optimizer and the
// tape all see one .grad — which is also what makes tied embeddings and
// gradient accumulation across micro-batches work for free.
class Var {
public:
    Var() = default;
    explicit Var(Tensor value, bool requires_grad = false)
        : impl_(std::make_shared<Impl>(Impl{std::move(value), Tensor{}, requires_grad})) {}

    bool defined() const { return impl_ != nullptr; }
    // Stable identity of the underlying storage — two Vars with the same id
    // are the same parameter (tied embeddings/lm_head).
    const void* id() const { return impl_.get(); }
    Tensor& value() const { return impl_->value; }
    bool requires_grad() const { return impl_ && impl_->requires_grad; }

    // Gradient, allocated as zeros on first touch. Grads are f32 always.
    Tensor& grad() const {
        if (!impl_->grad.defined()) {
            impl_->grad = Tensor::zeros(impl_->value.shape(), DType::F32);
        }
        return impl_->grad;
    }
    bool has_grad() const { return impl_ && impl_->grad.defined(); }
    void zero_grad() const {
        if (impl_ && impl_->grad.defined()) impl_->grad.zero_();
    }

    // Shape-only view: shares BOTH value and grad storage with this Var, so
    // gradient flow needs no tape node.
    Var reshaped(std::vector<int64_t> shape) const {
        Var out(value().view(shape), requires_grad());
        if (requires_grad()) out.impl_->grad = grad().view(std::move(shape));
        return out;
    }

private:
    struct Impl {
        Tensor value;
        Tensor grad;
        bool requires_grad;
    };
    std::shared_ptr<Impl> impl_;
};

// Dynamic tape. Ops that produce grad-requiring outputs push a backward
// lambda; backward() runs them in reverse and clears the tape. Not
// thread-safe by design: one training thread.
class Tape {
public:
    static Tape& get();

    bool enabled() const { return enabled_; }
    void set_enabled(bool e) { enabled_ = e; }

    void record(std::function<void()> backward_fn) {
        if (enabled_) nodes_.push_back(std::move(backward_fn));
    }

    // Seeds d(loss)/d(loss) = 1 and walks the tape in reverse.
    void backward(const Var& loss);

    void clear() { nodes_.clear(); }
    size_t size() const { return nodes_.size(); }

private:
    Tape() = default;
    std::vector<std::function<void()>> nodes_;
    bool enabled_ = true;
};

// RAII scope that disables tape recording (inference / generation).
struct NoGrad {
    NoGrad() : prev_(Tape::get().enabled()) { Tape::get().set_enabled(false); }
    ~NoGrad() { Tape::get().set_enabled(prev_); }
    NoGrad(const NoGrad&) = delete;
    NoGrad& operator=(const NoGrad&) = delete;

private:
    bool prev_;
};

} // namespace forge