/*
 * Copyright 2010 Google Inc.
 *
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#ifndef skgpu_BufferWriter_DEFINED
#define skgpu_BufferWriter_DEFINED

#include "include/core/SkImageInfo.h"
#include "include/core/SkRect.h"
#include "include/private/SkColorData.h"
#include "include/private/base/SkAssert.h"
#include "include/private/base/SkDebug.h"
#include "include/private/base/SkTemplates.h"
#include "include/private/base/SkTo.h"
#include "src/base/SkRectMemcpy.h"
#include "src/base/SkVx.h"
#include "src/core/SkConvertPixels.h"

#include <array>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <type_traits>
#include <utility>

namespace skgpu {

struct BufferWriter {
public:
    // Marks a read-only position in the underlying buffer
    struct Mark {
    public:
        Mark() : Mark(nullptr) {}
        Mark(void* ptr, size_t offset = 0)
                : fMark(reinterpret_cast<uintptr_t>(ptr) + offset) {
            SkASSERT(ptr || offset == 0);
        }

        bool operator< (const Mark& o) const { return fMark <  o.fMark; }
        bool operator<=(const Mark& o) const { return fMark <= o.fMark; }
        bool operator==(const Mark& o) const { return fMark == o.fMark; }
        bool operator!=(const Mark& o) const { return fMark != o.fMark; }
        bool operator>=(const Mark& o) const { return fMark >= o.fMark; }
        bool operator> (const Mark& o) const { return fMark >  o.fMark; }

        ptrdiff_t operator-(const Mark& o) const { return fMark - o.fMark; }

        explicit operator bool() const { return *this != Mark(); }
    private:
        uintptr_t fMark;
    };

    explicit operator bool() const { return fPtr != nullptr; }

    Mark mark(size_t offset=0) const {
        this->validate(offset);
        return Mark(fPtr, offset);
    }

protected:
    BufferWriter() = default;
    BufferWriter(void* ptr, size_t size) : fPtr(ptr) {
        SkDEBUGCODE(fEnd = Mark(ptr, ptr ? size : 0);)
    }
    BufferWriter(void* ptr, Mark end = {}) : fPtr(ptr) {
        SkDEBUGCODE(fEnd = end;)
    }

    BufferWriter& operator=(const BufferWriter&) = delete;
    BufferWriter& operator=(BufferWriter&& that) {
        fPtr = that.fPtr;
        that.fPtr = nullptr;
        SkDEBUGCODE(fEnd = that.fEnd;)
        SkDEBUGCODE(that.fEnd = Mark();)
        return *this;
    }

    // makeOffset effectively splits the current writer from {fPtr, fEnd} into {fPtr, p} and
    // a new writer {p, fEnd}. The same data range is accessible, but each byte can only be
    // set by a single writer. Automatically validates that there is enough bytes remaining in this
    // writer to do such a split.
    //
    // This splitting and validation means that providers of BufferWriters to callers can easily
    // and correctly track everything in a single BufferWriter field and use
    //    return std::exchange(fCurrWriter, fCurrWriter.makeOffset(requestedBytes));
    // This exposes the current writer position to the caller and sets the provider's new current
    // position to be just after the requested bytes.
    //
    // Templated so that it can create subclasses directly.
    template<typename W>
    W makeOffset(size_t offsetInBytes) const {
        this->validate(offsetInBytes);
        void* p = SkTAddOffset<void>(fPtr, offsetInBytes);
        Mark end{SkDEBUGCODE(fEnd)};
        SkDEBUGCODE(fEnd = Mark(p);)
        return W{p, end};
    }

    void validate(size_t bytesToWrite) const {
        // If the buffer writer had an end marked, make sure we're not crossing it.
        // Ideally, all creators of BufferWriters mark the end, but a lot of legacy code is not set
        // up to easily do this.
        SkASSERT(fPtr || bytesToWrite == 0);
        SkASSERT(!fEnd || Mark(fPtr, bytesToWrite) <= fEnd);
    }

protected:
    void* fPtr = nullptr;
    SkDEBUGCODE(mutable Mark fEnd = {};)
};

/**
 * Helper for writing vertex data to a buffer. Usage:
 *  VertexWriter vertices{target->makeVertexSpace(...)};
 *  vertices << A0 << B0 << C0 << ...;
 *  vertices << A1 << B1 << C1 << ...;
 *
 * Each value must be POD (plain old data), or have a specialization of the "<<" operator.
 */
struct VertexWriter : public BufferWriter {
    inline constexpr static uint32_t kIEEE_32_infinity = 0x7f800000;

    VertexWriter() = default;
    // DEPRECATED: Prefer specifying the size of the buffer being written to as well
    explicit VertexWriter(void* ptr) : BufferWriter(ptr, Mark()) {}

    VertexWriter(void* ptr, size_t size) : BufferWriter(ptr, size) {}
    VertexWriter(void* ptr, Mark end) : BufferWriter(ptr, end) {}

    VertexWriter(const VertexWriter&) = delete;
    VertexWriter(VertexWriter&& that) { *this = std::move(that); }

    VertexWriter& operator=(const VertexWriter&) = delete;
    VertexWriter& operator=(VertexWriter&& that) {
        BufferWriter::operator=(std::move(that));
        return *this;
    }

    VertexWriter makeOffset(size_t offsetInBytes) const {
        return this->BufferWriter::makeOffset<VertexWriter>(offsetInBytes);
    }

    template <typename T>
    struct Conditional {
        bool fCondition;
        T fValue;
    };

    template <typename T>
    static Conditional<T> If(bool condition, const T& value) {
        return {condition, value};
    }

    template <typename T>
    struct Skip {};

    template<typename T>
    struct ArrayDesc {
        const T* fArray;
        int fCount;
    };

    template <typename T>
    static ArrayDesc<T> Array(const T* array, int count) {
        return {array, count};
    }

    template<int kCount, typename T>
    struct RepeatDesc {
        const T& fVal;
    };

    template <int kCount, typename T>
    static RepeatDesc<kCount, T> Repeat(const T& val) {
        return {val};
    }

    /**
     * Specialized utilities for writing a four-vertices, with some data being replicated at each
     * vertex, and other data being the appropriate 2-components from an SkRect to construct a
     * triangle strip.
     *
     * - Four sets of data will be written
     *
     * - For any arguments where is_quad<Type>::value is true, a unique point will be written at
     *   each vertex. To make a custom type be emitted as a quad, declare:
     *
     *       template<> struct VertexWriter::is_quad<MyQuadClass> : std::true_type {};
     *
     *   and define:
     *
     *       MyQuadClass::writeVertex(int cornerIdx, VertexWriter&) const { ... }
     *
     * - For any arguments where is_quad<Type>::value is false, its value will be replicated at each
     *   vertex.
     */
    template <typename T>
    struct is_quad : std::false_type {};

    template <typename T>
    struct TriStrip {
        void writeVertex(int cornerIdx, VertexWriter& w) const {
            switch (cornerIdx) {
                case 0: w << l << t; return;
                case 1: w << l << b; return;
                case 2: w << r << t; return;
                case 3: w << r << b; return;
            }
            SkUNREACHABLE;
        }
        T l, t, r, b;
    };

    static TriStrip<float> TriStripFromRect(const SkRect& r) {
        return { r.fLeft, r.fTop, r.fRight, r.fBottom };
    }

    static TriStrip<uint16_t> TriStripFromUVs(const std::array<uint16_t, 4>& rect) {
        return { rect[0], rect[1], rect[2], rect[3] };
    }

    template <typename T>
    struct TriFan {
        void writeVertex(int cornerIdx, VertexWriter& w) const {
            switch (cornerIdx) {
                case 0: w << l << t; return;
                case 1: w << l << b; return;
                case 2: w << r << b; return;
                case 3: w << r << t; return;
            }
            SkUNREACHABLE;
        }
        T l, t, r, b;
    };

    static TriFan<float> TriFanFromRect(const SkRect& r) {
        return { r.fLeft, r.fTop, r.fRight, r.fBottom };
    }

    template <typename... Args>
    void writeQuad(const Args&... remainder) {
        this->writeQuadVertex<0>(remainder...);
        this->writeQuadVertex<1>(remainder...);
        this->writeQuadVertex<2>(remainder...);
        this->writeQuadVertex<3>(remainder...);
    }

private:
    template <int kCornerIdx, typename T, typename... Args>
    std::enable_if_t<!is_quad<T>::value, void> writeQuadVertex(const T& val,
                                                               const Args&... remainder) {
        *this << val;  // Non-quads duplicate their value.
        this->writeQuadVertex<kCornerIdx>(remainder...);
    }

    template <int kCornerIdx, typename Q, typename... Args>
    std::enable_if_t<is_quad<Q>::value, void> writeQuadVertex(const Q& quad,
                                                              const Args&... remainder) {
        quad.writeVertex(kCornerIdx, *this);  // Quads emit a different corner each time.
        this->writeQuadVertex<kCornerIdx>(remainder...);
    }

    template <int kCornerIdx>
    void writeQuadVertex() {}

    template <typename T>
    friend VertexWriter& operator<<(VertexWriter&, const T&);

    template <typename T>
    friend VertexWriter& operator<<(VertexWriter&, const ArrayDesc<T>&);
};

template <typename T>
inline VertexWriter& operator<<(VertexWriter& w, const T& val) {
    static_assert(std::is_trivially_copyable<T>::value, "");
    w.validate(sizeof(T));
    memcpy(w.fPtr, &val, sizeof(T));
    w = w.makeOffset(sizeof(T));
    return w;
}

template <typename T>
inline VertexWriter& operator<<(VertexWriter& w, const VertexWriter::Conditional<T>& val) {
    static_assert(std::is_trivially_copyable<T>::value, "");
    if (val.fCondition) {
        w << val.fValue;
    }
    return w;
}

template <typename T>
inline VertexWriter& operator<<(VertexWriter& w, const VertexWriter::Skip<T>& val) {
    w = w.makeOffset(sizeof(T));
    return w;
}

template <typename T>
inline VertexWriter& operator<<(VertexWriter& w, const VertexWriter::ArrayDesc<T>& array) {
    static_assert(std::is_trivially_copyable<T>::value, "");
    w.validate(array.fCount * sizeof(T));
    memcpy(w.fPtr, array.fArray, array.fCount * sizeof(T));
    w = w.makeOffset(sizeof(T) * array.fCount);
    return w;
}

template <int kCount, typename T>
inline VertexWriter& operator<<(VertexWriter& w, const VertexWriter::RepeatDesc<kCount,T>& repeat) {
    for (int i = 0; i < kCount; ++i) {
        w << repeat.fVal;
    }
    return w;
}

template <>
[[maybe_unused]] inline VertexWriter& operator<<(VertexWriter& w, const skvx::float4& vector) {
    w.validate(sizeof(vector));
    vector.store(w.fPtr);
    w = w.makeOffset(sizeof(vector));
    return w;
}

// Allow r-value/temporary writers to be appended to
template <typename T>
inline VertexWriter& operator<<(VertexWriter&& w, const T& val) { return w << val; }

template <typename T>
struct VertexWriter::is_quad<VertexWriter::TriStrip<T>> : std::true_type {};

template <typename T>
struct VertexWriter::is_quad<VertexWriter::TriFan<T>> : std::true_type {};

/**
 * VertexColor is a helper for writing colors to a vertex buffer. It outputs either four bytes or
 * or four float32 channels, depending on the wideColor parameter. Note that the GP needs to have
 * been constructed with the correct attribute type for colors, to match the usage here.
 */
class VertexColor {
public:
    VertexColor() = default;

    explicit VertexColor(const SkPMColor4f& color, bool wideColor) {
        this->set(color, wideColor);
    }

    void set(const SkPMColor4f& color, bool wideColor) {
        if (wideColor) {
            memcpy(fColor, color.vec(), sizeof(fColor));
        } else {
            fColor[0] = color.toBytes_RGBA();
        }
        fWideColor = wideColor;
    }

    size_t size() const { return fWideColor ? 16 : 4; }

private:
    template <typename T>
    friend VertexWriter& operator<<(VertexWriter&, const T&);

    uint32_t fColor[4];
    bool     fWideColor;
};

template <>
[[maybe_unused]] inline VertexWriter& operator<<(VertexWriter& w, const VertexColor& color) {
    w << color.fColor[0];
    if (color.fWideColor) {
        w << color.fColor[1]
          << color.fColor[2]
          << color.fColor[3];
    }
    return w;
}

///////////////////////////////////////////////////////////////////////////////////////////////////

struct IndexWriter : public BufferWriter {
    IndexWriter() = default;

    IndexWriter(void* ptr, size_t size) : BufferWriter(ptr, size) {}
    IndexWriter(void* ptr, Mark end) : BufferWriter(ptr, end) {}

    IndexWriter(const IndexWriter&) = delete;
    IndexWriter(IndexWriter&& that) { *this = std::move(that); }

    IndexWriter& operator=(const IndexWriter&) = delete;
    IndexWriter& operator=(IndexWriter&& that) {
        BufferWriter::operator=(std::move(that));
        return *this;
    }

    IndexWriter makeOffset(int numIndices) const {
        return this->BufferWriter::makeOffset<IndexWriter>(numIndices * sizeof(uint16_t));
    }

    void writeArray(const uint16_t* array, int count) {
        size_t arraySize = count * sizeof(uint16_t);
        this->validate(arraySize);
        memcpy(fPtr, array, arraySize);
        fPtr = SkTAddOffset<void>(fPtr, arraySize);

    }

    friend IndexWriter& operator<<(IndexWriter& w, uint16_t val);
};

inline IndexWriter& operator<<(IndexWriter& w, uint16_t val) {
    w.validate(sizeof(uint16_t));
    memcpy(w.fPtr, &val, sizeof(uint16_t));
    w = w.makeOffset(1);
    return w;
}

inline IndexWriter& operator<<(IndexWriter& w, int val) { return (w << SkTo<uint16_t>(val)); }

template<typename T>
inline IndexWriter& operator<<(IndexWriter&& w, const T& val) { return w << val; }

///////////////////////////////////////////////////////////////////////////////////////////////////

struct UniformWriter : public BufferWriter {
    UniformWriter() = default;

    UniformWriter(void* ptr, size_t size) : BufferWriter(ptr, size) {}
    UniformWriter(void* ptr, Mark end) : BufferWriter(ptr, end) {}

    UniformWriter(const UniformWriter&) = delete;
    UniformWriter(UniformWriter&& that) { *this = std::move(that); }

    UniformWriter& operator=(const UniformWriter&) = delete;
    UniformWriter& operator=(UniformWriter&& that) {
        BufferWriter::operator=(std::move(that));
        return *this;
    }

    void write(const void* src, size_t bytes) {
        this->validate(bytes);
        memcpy(fPtr, src, bytes);
        fPtr = SkTAddOffset<void>(fPtr, bytes);
    }
    void skipBytes(size_t bytes) {
        this->validate(bytes);
        fPtr = SkTAddOffset<void>(fPtr, bytes);
    }
};

///////////////////////////////////////////////////////////////////////////////////////////////////

struct TextureUploadWriter : public BufferWriter {
    TextureUploadWriter() = default;

    TextureUploadWriter(void* ptr, size_t size) : BufferWriter(ptr, size) {}

    TextureUploadWriter(const TextureUploadWriter&) = delete;
    TextureUploadWriter(TextureUploadWriter&& that) { *this = std::move(that); }

    TextureUploadWriter& operator=(const TextureUploadWriter&) = delete;
    TextureUploadWriter& operator=(TextureUploadWriter&& that) {
        BufferWriter::operator=(std::move(that));
        return *this;
    }

    // Writes a block of image data to the upload buffer, starting at `offset`. The source image is
    // `srcRowBytes` wide, and the written block is `dstRowBytes` wide and `rowCount` bytes tall.
    void write(size_t offset, const void* src, size_t srcRowBytes, size_t dstRowBytes,
               size_t trimRowBytes, int rowCount) {
        this->validate(offset + dstRowBytes * rowCount);
        void* dst = SkTAddOffset<void>(fPtr, offset);
        SkRectMemcpy(dst, dstRowBytes, src, srcRowBytes, trimRowBytes, rowCount);
    }

    void convertAndWrite(size_t offset,
                         const SkImageInfo& srcInfo, const void* src, size_t srcRowBytes,
                         const SkImageInfo& dstInfo, size_t dstRowBytes) {
        SkASSERT(srcInfo.width() == dstInfo.width() && srcInfo.height() == dstInfo.height());
        this->validate(offset + dstRowBytes * dstInfo.height());
        void* dst = SkTAddOffset<void>(fPtr, offset);
        SkAssertResult(SkConvertPixels(dstInfo, dst, dstRowBytes, srcInfo, src, srcRowBytes));
    }

    // Writes a block of image data to the upload buffer. It converts src data of RGB_888x
    // colorType into a 3 channel RGB_888 format.
    void writeRGBFromRGBx(size_t offset, const void* src, size_t srcRowBytes, size_t dstRowBytes,
                          int rowPixels, int rowCount) {
        this->validate(offset + dstRowBytes * rowCount);
        void* dst = SkTAddOffset<void>(fPtr, offset);
        auto* sRow = reinterpret_cast<const char*>(src);
        auto* dRow = reinterpret_cast<char*>(dst);

        for (int y = 0; y < rowCount; ++y) {
            for (int x = 0; x < rowPixels; ++x) {
                memcpy(dRow + 3*x, sRow+4*x, 3);
            }
            sRow += srcRowBytes;
            dRow += dstRowBytes;
        }
    }
};

}  // namespace skgpu

#endif // skgpu_BufferWriter_DEFINED