// Author: Simon-Pierre Boucher — contact@spboucher.ai
#include "core/tensor.h"

#include "core/device.h"

#include <Metal/Metal.hpp>

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <sstream>

namespace forge {

namespace {
[[noreturn]] void die(const std::string& msg) {
    std::fprintf(stderr, "forge: %s\n", msg.c_str());
    std::abort();
}
} // namespace

// Owns one pooled MTLBuffer; returns it to the pool when the last Tensor
// view drops it. External storages (fmodel mmaps) retain a caller-owned
// buffer instead and release it directly — the allocator never sees it.
struct Tensor::Storage {
    MTL::Buffer* buffer = nullptr;
    bool external = false;

    explicit Storage(size_t nbytes) {
        buffer = Device::get().allocator().acquire(nbytes);
    }
    Storage(MTL::Buffer* ext, bool) : buffer(ext), external(true) {
        buffer->retain();
    }
    ~Storage() {
        if (external) buffer->release();
        else Device::get().allocator().release(buffer);
    }
    Storage(const Storage&) = delete;
    Storage& operator=(const Storage&) = delete;
};

std::vector<int64_t> Tensor::contiguous_strides(const std::vector<int64_t>& shape) {
    std::vector<int64_t> strides(shape.size());
    int64_t acc = 1;
    for (int64_t i = int64_t(shape.size()) - 1; i >= 0; --i) {
        strides[size_t(i)] = acc;
        acc *= shape[size_t(i)];
    }
    return strides;
}

Tensor Tensor::empty(std::vector<int64_t> shape, DType dtype) {
    Tensor t;
    t.dtype_ = dtype;
    t.shape_ = std::move(shape);
    t.strides_ = contiguous_strides(t.shape_);
    int64_t n = 1;
    for (int64_t d : t.shape_) {
        if (d <= 0) die("Tensor::empty: non-positive dim");
        n *= d;
    }
    t.storage_ = std::make_shared<Storage>(size_t(n) * dtype_size(dtype));
    return t;
}

Tensor Tensor::from_buffer(MTL::Buffer* buffer, size_t byte_offset,
                           std::vector<int64_t> shape, DType dtype) {
    Tensor t;
    t.dtype_ = dtype;
    t.shape_ = std::move(shape);
    t.strides_ = contiguous_strides(t.shape_);
    if (byte_offset % dtype_size(dtype) != 0)
        die("Tensor::from_buffer: offset not aligned to dtype size");
    t.offset_ = int64_t(byte_offset / dtype_size(dtype));
    t.storage_ = std::make_shared<Storage>(buffer, true);
    return t;
}

Tensor Tensor::zeros(std::vector<int64_t> shape, DType dtype) {
    Tensor t = empty(std::move(shape), dtype);
    std::memset(t.raw(), 0, t.nbytes());
    return t;
}

Tensor Tensor::full(std::vector<int64_t> shape, float value, DType dtype) {
    Tensor t = empty(std::move(shape), dtype);
    t.fill_(value);
    return t;
}

int64_t Tensor::size(int64_t dim) const {
    if (dim < 0) dim += ndim();
    if (dim < 0 || dim >= ndim()) die("Tensor::size: dim out of range");
    return shape_[size_t(dim)];
}

int64_t Tensor::numel() const {
    int64_t n = 1;
    for (int64_t d : shape_) n *= d;
    return n;
}

bool Tensor::is_contiguous() const {
    return strides_ == contiguous_strides(shape_);
}

void* Tensor::raw() const {
    if (!storage_) die("Tensor::raw: undefined tensor");
    return static_cast<char*>(storage_->buffer->contents()) + size_t(offset_) * itemsize();
}

MTL::Buffer* Tensor::buffer() const {
    if (!storage_) die("Tensor::buffer: undefined tensor");
    return storage_->buffer;
}

size_t Tensor::buffer_offset() const {
    return size_t(offset_) * itemsize();
}

Tensor Tensor::view(std::vector<int64_t> new_shape) const {
    if (!is_contiguous()) die("Tensor::view: tensor not contiguous");
    int64_t n = 1;
    for (int64_t d : new_shape) n *= d;
    if (n != numel()) die("Tensor::view: numel mismatch");
    Tensor t = *this;
    t.shape_ = std::move(new_shape);
    t.strides_ = contiguous_strides(t.shape_);
    return t;
}

Tensor Tensor::slice0(int64_t start, int64_t len) const {
    if (!is_contiguous()) die("Tensor::slice0: tensor not contiguous");
    if (ndim() < 1 || start < 0 || len <= 0 || start + len > shape_[0])
        die("Tensor::slice0: range out of bounds");
    Tensor t = *this;
    t.shape_[0] = len;
    t.offset_ = offset_ + start * strides_[0];
    return t;
}

void Tensor::fill_(float value) {
    const int64_t n = numel();
    switch (dtype_) {
        case DType::F32: {
            float* p = data<float>();
            for (int64_t i = 0; i < n; ++i) p[i] = value;
            break;
        }
        case DType::F16: {
            f16_t* p = data<f16_t>();
            const f16_t v = f16_t(value);
            for (int64_t i = 0; i < n; ++i) p[i] = v;
            break;
        }
        case DType::BF16: {
            uint16_t* p = data<uint16_t>();
            const uint16_t v = float_to_bf16(value);
            for (int64_t i = 0; i < n; ++i) p[i] = v;
            break;
        }
        case DType::U16: {
            uint16_t* p = data<uint16_t>();
            const uint16_t v = uint16_t(value);
            for (int64_t i = 0; i < n; ++i) p[i] = v;
            break;
        }
        case DType::I32: {
            int32_t* p = data<int32_t>();
            const int32_t v = int32_t(value);
            for (int64_t i = 0; i < n; ++i) p[i] = v;
            break;
        }
    }
}

float Tensor::item_at(int64_t i) const {
    switch (dtype_) {
        case DType::F32:  return data<float>()[i];
        case DType::F16:  return float(data<f16_t>()[i]);
        case DType::BF16: return bf16_to_float(data<uint16_t>()[i]);
        case DType::U16:  return float(data<uint16_t>()[i]);
        case DType::I32:  return float(data<int32_t>()[i]);
    }
    return 0.0f;
}

void Tensor::set_item(int64_t i, float value) {
    switch (dtype_) {
        case DType::F32:  data<float>()[i] = value; break;
        case DType::F16:  data<f16_t>()[i] = f16_t(value); break;
        case DType::BF16: data<uint16_t>()[i] = float_to_bf16(value); break;
        case DType::U16:  data<uint16_t>()[i] = uint16_t(value); break;
        case DType::I32:  data<int32_t>()[i] = int32_t(value); break;
    }
}

std::string Tensor::describe() const {
    std::ostringstream os;
    os << "Tensor(" << dtype_name(dtype_) << ", [";
    for (size_t i = 0; i < shape_.size(); ++i) {
        if (i) os << ", ";
        os << shape_[i];
    }
    os << "])";
    return os.str();
}

} // namespace forge