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

#include "core/dtype.h"

#include <cstdint>
#include <memory>
#include <string>
#include <vector>

namespace MTL {
class Buffer;
}

namespace forge {

// Refcounted view over a pooled MTLBuffer (MTLStorageModeShared), so the
// same memory is addressable from CPU code and GPU kernels with zero copies.
// Row-major, element strides. Views (reshape/view/slice) share storage.
class Tensor {
public:
    Tensor() = default;

    static Tensor empty(std::vector<int64_t> shape, DType dtype = DType::F32);
    static Tensor zeros(std::vector<int64_t> shape, DType dtype = DType::F32);
    static Tensor full(std::vector<int64_t> shape, float value, DType dtype = DType::F32);

    // Wrap an externally-owned MTLBuffer region (e.g. an mmapped .fmodel
    // shard bridged with newBuffer(bytesNoCopy)). The buffer is retained for
    // the storage's lifetime and released — NOT returned to the allocator
    // pool — when the last view drops. byte_offset must be a multiple of the
    // dtype size (fmodel aligns to the 16 KB page, far stricter).
    static Tensor from_buffer(MTL::Buffer* buffer, size_t byte_offset,
                              std::vector<int64_t> shape, DType dtype = DType::F32);

    bool defined() const { return storage_ != nullptr; }
    DType dtype() const { return dtype_; }
    int64_t ndim() const { return int64_t(shape_.size()); }
    const std::vector<int64_t>& shape() const { return shape_; }
    const std::vector<int64_t>& strides() const { return strides_; }
    int64_t size(int64_t dim) const;
    int64_t numel() const;
    size_t itemsize() const { return dtype_size(dtype_); }
    size_t nbytes() const { return size_t(numel()) * itemsize(); }
    bool is_contiguous() const;

    // CPU access. Valid at all times (unified memory) — but only touch it
    // when no in-flight command buffer may write the same storage.
    void* raw() const;
    template <typename T>
    T* data() const { return static_cast<T*>(raw()); }

    // GPU access.
    MTL::Buffer* buffer() const;
    size_t buffer_offset() const; // bytes from the start of the MTLBuffer

    // Views (no copy). Both require contiguous layout.
    Tensor view(std::vector<int64_t> new_shape) const;
    Tensor reshape(std::vector<int64_t> new_shape) const { return view(std::move(new_shape)); }
    // Slice along dim 0: rows [start, start+len). Contiguous only.
    Tensor slice0(int64_t start, int64_t len) const;

    void fill_(float value);
    void zero_() { fill_(0.0f); }

    // Read/write a single element as float, whatever the dtype (test/debug).
    float item_at(int64_t linear_index) const;
    void set_item(int64_t linear_index, float value);

    std::string describe() const;

private:
    struct Storage;

    std::shared_ptr<Storage> storage_;
    std::vector<int64_t> shape_;
    std::vector<int64_t> strides_; // in elements
    int64_t offset_ = 0;           // in elements
    DType dtype_ = DType::F32;

    static std::vector<int64_t> contiguous_strides(const std::vector<int64_t>& shape);
};

} // namespace forge