#pragma once

#include "MayaFlux/Buffers/VKBuffer.hpp"
#include "MayaFlux/Kakshya/NDData/DataAccess.hpp"

namespace MayaFlux::Buffers {

class VKBuffer;
class AudioBuffer;

inline constexpr float k_buffer_growth_factor = 1.5F;

/**
 * @brief Opaque handle to an in-flight transfer.
 */
using TransferHandle = std::shared_ptr<void>;

/**
 * @brief Upload data to a host-visible buffer
 * @param target Target VKBuffer to upload data into
 * @param data DataVariant containing the data to upload
 *
 * This function handles uploading data from a Kakshya::DataVariant into a
 * host-visible VKBuffer. It maps the buffer memory, copies the data, and
 * marks the buffer as dirty for synchronization.
 */
MAYAFLUX_API void upload_host_visible(const std::shared_ptr<VKBuffer>& target, const Kakshya::DataVariant& data);

/**
 * @brief Upload data to a device-local buffer using a staging buffer
 * @param target Target VKBuffer to upload data into
 * @param staging_buffer Host-visible staging VKBuffer used for the upload
 * @param data DataVariant containing the data to upload
 *
 * This function handles uploading data from a Kakshya::DataVariant into a
 * device-local VKBuffer by utilizing a staging buffer. It copies the data
 * into the staging buffer, flushes it, and then issues a command to copy
 * from the staging buffer to the target device-local buffer.
 */
MAYAFLUX_API void upload_device_local(const std::shared_ptr<VKBuffer>& target, const std::shared_ptr<VKBuffer>& staging_buffer, const Kakshya::DataVariant& data);

/**
 * @brief Upload host memory into a raw back_buffers slot.
 * @param slot Destination slot, typically from VKBuffer::get_back_buffer().
 * @param data Source pointer, at least @p size bytes.
 * @param size Byte count to copy.
 * @param staging Persistent staging buffer for device-local slots. Ignored,
 *        and may be null, when slot.mapped_ptr is already valid.
 *
 * Writes through slot.mapped_ptr directly when present (host-visible slot,
 * no transfer). Otherwise copies into @p staging and records a fenced
 * device-to-device copy into slot.buffer, mirroring download_back_buffer
 * with the direction reversed.
 */
MAYAFLUX_API void upload_back_buffer(
    const VKBufferResources::GenerationSlot& slot,
    const void* data,
    size_t size,
    std::shared_ptr<VKBuffer>& staging);

/**
 * @brief Upload to a back_buffers slot without waiting for completion.
 * @param slot Destination slot, typically from VKBuffer::get_back_buffer().
 * @param data Source pointer, at least @p size bytes. Copied into staging
 *        before returning, so the caller may free it immediately.
 * @param size Byte count to copy.
 * @param staging Persistent staging buffer, grown if too small.
 * @return Handle to pass to resolve_transfer, or null for host-visible
 *         slots where the write was a direct memcpy.
 *
 * The staging buffer is in use until the returned handle is resolved. Any
 * further transfer through the same staging buffer must resolve this one
 * first, or the copy source is overwritten mid-flight.
 */
MAYAFLUX_API TransferHandle upload_back_buffer_async(
    const VKBufferResources::GenerationSlot& slot,
    const void* data,
    size_t size,
    std::shared_ptr<VKBuffer>& staging);

/**
 * @brief Wait for a transfer and release its resources.
 * @param handle Handle from upload_back_buffer_async. Nulled on return.
 *
 * No-op on a null handle. Blocks the calling thread until the fence
 * signals, which is immediate when a frame or more has elapsed since
 * submission.
 */
MAYAFLUX_API void resolve_transfer(TransferHandle& handle);

/**
 * @brief Download data from a host-visible buffer
 * @param source Source VKBuffer to download data from
 * @param target Target VKBuffer to store the downloaded data
 *
 * This function handles downloading data from a host-visible VKBuffer.
 * It maps the buffer memory, copies the data into a CPU-accessible format,
 * and updates the associated target buffer.
 */
MAYAFLUX_API void download_host_visible(const std::shared_ptr<VKBuffer>& source, const std::shared_ptr<VKBuffer>& target);

/**
 * @brief Download data from a device-local buffer using a staging buffer
 * @param source Source VKBuffer to download data from
 * @param target Target VKBuffer to store the downloaded data
 * @param staging_buffer Host-visible staging VKBuffer used for the download
 *
 * This function handles downloading data from a device-local VKBuffer by
 * utilizing a staging buffer. It issues a command to copy the data from
 * the device-local buffer to the staging buffer, invalidates the staging
 * buffer memory, and then copies the data into a CPU-accessible format
 * to update the target buffer.
 */
MAYAFLUX_API void download_device_local(const std::shared_ptr<VKBuffer>& source, const std::shared_ptr<VKBuffer>& target, const std::shared_ptr<VKBuffer>& staging_buffer);

/**
 * @brief Download a raw back_buffers slot to host memory.
 * @param slot Source slot, typically from VKBuffer::get_back_buffer().
 * @param size Byte count to copy.
 * @param data Destination pointer, at least @p size bytes.
 * @param staging Persistent staging buffer for device-local slots. Ignored,
 *        and may be null, when slot.mapped_ptr is already valid.
 *
 * Reads slot.mapped_ptr directly when present (host-visible slot, no
 * transfer). Otherwise records a fenced copy into @p staging exactly as
 * download_from_gpu_async does, using slot.buffer as the copy source.
 */
MAYAFLUX_API void download_back_buffer(
    const VKBufferResources::GenerationSlot& slot,
    void* data,
    size_t size,
    std::shared_ptr<VKBuffer>& staging);

/**
 * @brief Upload raw data to GPU buffer (auto-detects host-visible vs device-local)
 * @param data Source data pointer
 * @param size Size in bytes
 * @param target Target GPU buffer
 * @param staging Optional staging buffer (created if needed for device-local)
 *
 * Convenience wrapper over StagingUtils that:
 * - Converts raw pointer → DataVariant
 * - Auto-detects if buffer is host-visible or device-local
 * - Handles staging buffer creation if needed
 */
MAYAFLUX_API void upload_to_gpu(
    const void* data,
    size_t size,
    const std::shared_ptr<VKBuffer>& target,
    const std::shared_ptr<VKBuffer>& staging = nullptr);

/**
 * @brief Upload typed data to GPU buffer
 * @tparam T Data type (float, double, int, etc.)
 * @param data Source data span
 * @param target Target GPU buffer
 * @param staging Optional staging buffer
 */
template <typename T>
void upload_to_gpu(
    std::span<const T> data,
    const std::shared_ptr<VKBuffer>& target,
    const std::shared_ptr<VKBuffer>& staging = nullptr)
{
    upload_to_gpu(data.data(), data.size_bytes(), target, staging);
}

/**
 * @brief Upload vector to GPU buffer
 * @tparam T Data type
 * @param data Source data vector
 * @param target Target GPU buffer
 * @param staging Optional staging buffer
 */
template <typename T>
void upload_to_gpu(
    const std::vector<T>& data,
    const std::shared_ptr<VKBuffer>& target,
    const std::shared_ptr<VKBuffer>& staging = nullptr)
{
    upload_to_gpu(std::span<const T>(data), target, staging);
}

/**
 * @brief Upload @p size bytes to @p target, growing both buffers first if needed.
 * @param data         Source pointer.
 * @param size         Byte count to upload.
 * @param target       Destination GPU buffer.
 * @param staging      Paired staging buffer, or nullptr for host-visible targets.
 * @param growth_factor Forwarded to ensure_gpu_capacity.
 *
 * Combines ensure_gpu_capacity() and upload_to_gpu() into a single call.
 * The upload is clamped to the post-resize capacity, which is always >= @p size
 * after ensure_gpu_capacity returns.
 */
MAYAFLUX_API void upload_resizing(
    const void* data,
    size_t size,
    const std::shared_ptr<VKBuffer>& target,
    const std::shared_ptr<VKBuffer>& staging,
    float growth_factor = 1.5F);

/**
 * @brief Download from GPU buffer to raw data (auto-detects host-visible vs device-local)
 * @param source Source GPU buffer
 * @param data Destination data pointer
 * @param size Size in bytes
 * @param staging Optional staging buffer (created if needed for device-local)
 *
 * Convenience wrapper over StagingUtils that:
 * - Auto-detects if buffer is host-visible or device-local
 * - Handles staging buffer creation if needed
 * - Copies data to destination pointer
 */
MAYAFLUX_API void download_from_gpu(
    const std::shared_ptr<VKBuffer>& source,
    void* data,
    size_t size,
    const std::shared_ptr<VKBuffer>& staging = nullptr);

/**
 * @brief Download from GPU buffer to typed span
 * @tparam T Data type
 * @param source Source GPU buffer
 * @param data Destination data span
 * @param staging Optional staging buffer
 */
template <typename T>
void download_from_gpu(
    const std::shared_ptr<VKBuffer>& source,
    std::span<T> data,
    const std::shared_ptr<VKBuffer>& staging = nullptr)
{
    download_from_gpu(source, data.data(), data.size_bytes(), staging);
}

/**
 * @brief Download from GPU buffer to vector
 * @tparam T Data type
 * @param source Source GPU buffer
 * @param data Destination data vector (resized if needed)
 * @param staging Optional staging buffer
 */
template <typename T>
void download_from_gpu(
    const std::shared_ptr<VKBuffer>& source,
    std::vector<T>& data,
    const std::shared_ptr<VKBuffer>& staging = nullptr)
{
    size_t element_count = source->get_size_bytes() / sizeof(T);
    data.resize(element_count);
    download_from_gpu(source, std::span<T>(data), staging);
}

/**
 * @brief Download from a device-local GPU buffer without stalling the graphics queue.
 *
 * Records a buffer copy into a fenced command buffer, waits on the fence from
 * the calling thread, then memcpys from the mapped staging buffer into @p data.
 * Unlike download_from_gpu, this does not call queue.waitIdle, making it safe
 * to call from any thread that is not the graphics thread itself.
 *
 * @p staging is allocated and cached by the caller to avoid per-call Vulkan
 * object churn. Pass the same staging buffer each frame; it is resized if
 * @p size exceeds its current capacity.
 *
 * @param source   Device-local source buffer.
 * @param data     Destination host pointer, at least @p size bytes.
 * @param size     Byte count to copy.
 * @param staging  Persistent host-visible staging buffer. Created via
 *                 create_staging_buffer(). Resized in-place if too small.
 */
MAYAFLUX_API void download_from_gpu_async(
    const std::shared_ptr<VKBuffer>& source,
    void* data,
    size_t size,
    std::shared_ptr<VKBuffer>& staging);

/**
 * @brief Resolve the GPU-resident image from any GpuImageSource buffer.
 *
 * Selects get_texture() or get_gpu_texture() at compile time based on which
 * the concrete buffer type exposes. Returns nullptr if the buffer has not yet
 * produced a GPU texture.
 *
 * @tparam T A type satisfying GpuImageSource.
 * @param buffer The pixel-bearing buffer to query.
 * @return GPU-resident VKImage, or nullptr.
 */
template <GpuImageSource T>
[[nodiscard]] std::shared_ptr<Core::VKImage> resolve_gpu_image(const T& buffer)
{
    if constexpr (requires(const T& b) {
                      { b.get_texture() } -> std::convertible_to<std::shared_ptr<Core::VKImage>>;
                  }) {
        return buffer.get_texture();
    } else {
        return buffer.get_gpu_texture();
    }
}

/**
 * @brief Download a VKImage to CPU and return a normalised float span.
 *
 * Downloads @p image into @p raw_staging via TextureLoom::download_data,
 * then normalises into @p work via Kakshya::as_normalised_float.
 *
 * @p raw_staging is resized to match the image byte footprint when the image
 * dimensions or format change. @p work is the caller-owned scratch buffer
 * reused across calls to avoid per-frame allocation.
 *
 * Returns an empty span if @p image is null, not initialised, or its format
 * has no ImageFormat mapping.
 *
 * @param image       GPU-resident source image.
 * @param raw_staging Persistent byte buffer for the GPU download. Reuse across calls.
 * @param work        Persistent float buffer for normalisation output. Reuse across calls.
 * @param gpu_staging Optional persistent staging buffer for device-local images. Reuse across calls.
 * @return Normalised float span pointing into @p work (or directly into the
 *         variant storage for float-format images).
 */
[[nodiscard]] MAYAFLUX_API std::span<const float> download_and_normalise(
    const std::shared_ptr<Core::VKImage>& image,
    std::vector<uint8_t>& raw_staging,
    std::vector<float>& work,
    const std::shared_ptr<VKBuffer>& gpu_staging);

/**
 * @brief Create staging buffer for transfers
 * @param size Size in bytes
 * @return Host-visible staging buffer ready for transfers
 */
MAYAFLUX_API std::shared_ptr<VKBuffer> create_staging_buffer(size_t size);

/**
 * @brief Allocate a persistent host-visible staging buffer sized for
 *        repeated streaming uploads to an image of @p size bytes.
 *        Call once; pass the result to upload_image_streaming() every frame.
 * @param size Byte footprint of the target image (use VKImage::get_size_bytes()).
 * @return Initialised host-visible VKBuffer, or nullptr on failure.
 */
[[nodiscard]] MAYAFLUX_API std::shared_ptr<VKBuffer> create_image_staging_buffer(size_t size);

/**
 * @brief Check if buffer is device-local (staging needed)
 * @param buffer Buffer to check
 * @return True if buffer is device-local
 */
MAYAFLUX_API bool is_device_local(const std::shared_ptr<VKBuffer>& buffer);

/**
 * @brief Grow a GPU buffer (and its paired staging buffer) to fit @p required bytes.
 * @param target       Device-local or host-visible GPU buffer to resize.
 * @param staging      Paired staging buffer, or nullptr if the target is host-visible.
 * @param required     Bytes needed for the next upload.
 * @param growth_factor Multiplier applied when a resize is triggered (default 1.5).
 *
 * No-op when @p target already has sufficient capacity. When a resize is necessary
 * both buffers are grown to `required * growth_factor` so subsequent small
 * increments do not trigger further allocations. Existing GPU data is not preserved
 * (preserve_data = false) because the caller is about to overwrite it.
 */
MAYAFLUX_API void ensure_gpu_capacity(
    const std::shared_ptr<VKBuffer>& target,
    const std::shared_ptr<VKBuffer>& staging,
    size_t required,
    float growth_factor = k_buffer_growth_factor);

/**
 * @brief Upload data from DataAccess view to GPU buffer (precision-preserving)
 * @tparam T View type (double, glm::dvec2, glm::dvec3, glm::vec3, float, etc.)
 * @param accessor DataAccess instance providing the view
 * @param target Target GPU buffer
 * @param staging Optional staging buffer (auto-created if needed)
 *
 * Zero-copy when types match, automatic conversion cache when they don't.
 * For AUDIO modalities, defaults to DOUBLE precision to preserve accuracy.
 */
template <typename T>
void upload_from_view(
    const Kakshya::DataAccess& accessor,
    const std::shared_ptr<VKBuffer>& target,
    const std::shared_ptr<VKBuffer>& staging = nullptr)
{
    auto view = accessor.view<T>();

    const void* data_ptr = view.data();
    size_t data_bytes = view.size() * sizeof(T);

    if constexpr (std::is_same_v<T, double>) {
        const auto modality = target->get_modality();
        if (modality != Kakshya::DataModality::AUDIO_1D
            && modality != Kakshya::DataModality::AUDIO_MULTICHANNEL) {
            MF_WARN(Journal::Component::Buffers, Journal::Context::BufferProcessing,
                "Uploading double precision to buffer with modality {}. Consider using AUDIO_1D or AUDIO_MULTICHANNEL.",
                Kakshya::modality_to_string(modality));
        }
    }

    upload_to_gpu(data_ptr, data_bytes, target, staging);
}

/**
 * @brief Upload structured data with GLM double-precision types
 * @tparam T GLM type (glm::dvec2, glm::dvec3, glm::dvec4 for double precision)
 * @param accessor DataAccess with structured dimensions
 * @param target Target GPU buffer
 * @param staging Optional staging buffer
 *
 * Use this for high-precision structured data like audio samples interpreted
 * as multi-dimensional vectors. Supports both single and double precision GLM types.
 */
template <typename T>
    requires GlmType<T>
void upload_structured_view(
    const Kakshya::DataAccess& accessor,
    const std::shared_ptr<VKBuffer>& target,
    const std::shared_ptr<VKBuffer>& staging = nullptr)
{
    if (!accessor.is_structured()) {
        error<std::invalid_argument>(
            Journal::Component::Buffers,
            Journal::Context::BufferProcessing,
            std::source_location::current(),
            "Cannot upload structured view from non-structured data");
    }

    auto structured_view = accessor.view<T>();
    upload_to_gpu(structured_view.data(), structured_view.size_bytes(), target, staging);
}

/**
 * @brief Download GPU buffer to DataAccess-compatible format (precision-preserving)
 */
template <typename T>
Kakshya::DataAccess download_to_view(
    const std::shared_ptr<VKBuffer>& source,
    Kakshya::DataVariant& target_variant,
    const std::vector<Kakshya::DataDimension>& dimensions,
    Kakshya::DataModality modality,
    const std::shared_ptr<VKBuffer>& staging = nullptr)
{
    size_t element_count = source->get_size_bytes() / sizeof(T);

    std::vector<T> temp_buffer(element_count);
    download_from_gpu(source, temp_buffer, staging);

    target_variant = std::move(temp_buffer);

    return Kakshya::DataAccess(target_variant, dimensions, modality);
}

/**
 * @brief Upload AudioBuffer to GPU (always double precision)
 * @param audio_buffer Source audio buffer (double[])
 * @param gpu_buffer Target GPU buffer (must support R64Sfloat format)
 * @param staging Optional staging buffer (auto-created if needed)
 *
 * AudioBuffer is always double precision. This function ensures the GPU buffer
 * is configured for double precision and performs a direct upload with no conversion.
 *
 * @throws std::runtime_error if gpu_buffer doesn't support double precision
 */
void upload_audio_to_gpu(
    const std::shared_ptr<AudioBuffer>& audio_buffer,
    const std::shared_ptr<VKBuffer>& gpu_buffer,
    const std::shared_ptr<VKBuffer>& staging = nullptr);

/**
 * @brief Download GPU buffer to AudioBuffer (expects double precision)
 * @param gpu_buffer Source GPU buffer (should contain double precision data)
 * @param audio_buffer Target audio buffer (always double[])
 * @param staging Optional staging buffer (auto-created if needed)
 *
 * Downloads GPU data and copies to AudioBuffer. If the GPU buffer contains
 * float data instead of double, DataAccess will handle the upconversion
 * (though this is not recommended for audio precision).
 */
void download_audio_from_gpu(
    const std::shared_ptr<VKBuffer>& gpu_buffer,
    const std::shared_ptr<AudioBuffer>& audio_buffer,
    const std::shared_ptr<VKBuffer>& staging = nullptr);

} // namespace MayaFlux::Buffers
