// Copyright 2019 The Dawn & Tint Authors
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
//    list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its
//    contributors may be used to endorse or promote products derived from
//    this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include <limits>
#include <memory>
#include <span>
#include <utility>

#include "dawn/wire/WireCmd_autogen.h"
#include "src/dawn/common/StringViewUtils.h"
#include "src/dawn/wire/BufferConsumer.h"
#include "src/dawn/wire/WireResult.h"
#include "src/dawn/wire/server/Server.h"
#include "src/utils/assert.h"

namespace dawn::wire::server {

WireResult Server::PreHandleDeviceCreateErrorBuffer(const DeviceCreateErrorBufferCmd& cmd) {
    // mappedAtCreation isn't implemented in CreateErrorBuffer.
    // The client blocks this, so we can treat it as a fatal wire error (for fuzzers).
    if (cmd.descriptor->mappedAtCreation) {
        return WireResult::FatalError;
    }

    return WireResult::Success;
}

WireResult Server::PreHandleBufferUnmap(const BufferUnmapCmd& cmd) {
    Known<WGPUBuffer> buffer;
    WIRE_TRY(Get(cmd.selfId, &buffer));

    buffer->mapState.Use([&](auto mapState) {
        if (buffer->mappedAtCreation &&
            !(buffer->usage & (WGPUBufferUsage_MapWrite | WGPUBufferUsage_MapRead))) {
            // This indicates the memoryHandle is for mappedAtCreation only. Destroy on unmap
            // memoryHandle could have possibly been deleted if buffer is already destroyed so we
            // don't assert it's non-null
            mapState->memoryHandle = nullptr;
        }
    });

    return WireResult::Success;
}

WireResult Server::PreHandleBufferDestroy(const BufferDestroyCmd& cmd) {
    // Destroying a buffer does an implicit unmapping.
    Known<WGPUBuffer> buffer;
    WIRE_TRY(Get(cmd.selfId, &buffer));

    // The buffer was destroyed. Clear the MemoryHandle.
    buffer->mapState.Use([](auto mapState) { mapState->memoryHandle = nullptr; });

    return WireResult::Success;
}

WireResult Server::DoBufferMapAsync(Known<WGPUBuffer> buffer,
                                    Known<WGPUInstance> instance,
                                    Future future,
                                    wgpu::MapMode mode,
                                    size_t offset,
                                    size_t size) {
    // These requests are just forwarded to the buffer, with userdata containing what the
    // client will require in the return command.
    std::unique_ptr<MapUserdata> userdata = MakeUserdata<MapUserdata>();
    userdata->buffer = buffer.AsHandle();
    userdata->instanceId = instance.id;
    userdata->future = future;
    userdata->mode = mode;

    // Make sure that the size is not WGPU_WHOLE_MAP_SIZE because we want the client to give us an
    // explicit size (so we don't have to handle the defaulting here). The client needs to track
    // the size of the buffer and can do the default size computation.
    if (size >= WGPU_WHOLE_MAP_SIZE) {
        return WireResult::FatalError;
    }

    userdata->offset = offset;
    userdata->size = size;

    mProcs->bufferMapAsync(
        buffer->handle, ToAPI(mode), offset, size,
        MakeCallbackInfo<WGPUBufferMapCallbackInfo, &Server::OnBufferMapAsyncCallback>(
            userdata.release()));

    return WireResult::Success;
}

WireResult Server::DoDeviceCreateBuffer(Known<WGPUDevice> device,
                                        const BufferDescriptor* descriptor,
                                        ObjectHandle bufferHandle,
                                        Span<const std::byte> memoryHandleCreateInfo) {
    // Create and register the buffer object.
    Reserved<WGPUBuffer> buffer;
    WIRE_TRY(Allocate(&buffer, bufferHandle));
    buffer->usage = ToAPI(descriptor->usage);
    buffer->mappedAtCreation = descriptor->mappedAtCreation;

    bool isMappable =
        descriptor->mappedAtCreation != 0u ||
        ((descriptor->usage & (wgpu::BufferUsage::MapRead | wgpu::BufferUsage::MapWrite)) !=
         wgpu::BufferUsage::None);

    // If the buffer is not mappable, create it normally without shared memory.
    if (!isMappable) {
        buffer->handle = mProcs->deviceCreateBuffer(device->handle, ToAPI(descriptor));
        return WireResult::Success;
    }

    // Return `FatalError` if we fail to deserialize the memory handle.
    std::unique_ptr<MemoryTransferService::MemoryHandle> memoryHandle =
        mMemoryTransferService->DeserializeMemoryHandle(memoryHandleCreateInfo);
    if (memoryHandle == nullptr) {
        buffer->handle = mProcs->deviceCreateBuffer(device->handle, ToAPI(descriptor));
        return WireResult::FatalError;
    }

    // Try to wrap the shared memory into a buffer. Note that:
    // - `TryWrapInBuffer()` may return a valid buffer, an error buffer or nullptr to align with the
    //   behavior of `DeviceBase::APICreateBuffer()`.
    // - When a valid buffer is returned, `beginAccess()` must have been called on it.
    // - Index/indirect buffers are not supported to prevent the TOCTOU issue with races between
    //   GPU-side validation and the client modifying the buffer.
    if (!(descriptor->usage & (wgpu::BufferUsage::Indirect | wgpu::BufferUsage::Index))) {
        buffer->handle =
            memoryHandle->TryWrapInBuffer(mProcs.get(), device->handle, ToAPI(descriptor));
    }
    if (buffer->handle != nullptr) {
        buffer->backedWithSharedMemory = true;
        return buffer->mapState.Use([&](auto mapState) {
            mapState->memoryHandle = std::move(memoryHandle);
            return WireResult::Success;
        });
    } else {
        // Returning `nullptr` from `TryWrapInBuffer()` indicates that we failed to wrap the shared
        // memory into a buffer. Try to create it normally without shared memory.
        buffer->handle = mProcs->deviceCreateBuffer(device->handle, ToAPI(descriptor));
    }

    // A null buffer indicates that mapping-at-creation failed inside createBuffer. Unmark the
    // buffer as allocated so we will skip freeing it.
    if (buffer->handle == nullptr) {
        DAWN_ASSERT(descriptor->mappedAtCreation);
        buffer->state = AllocationState::Reserved;
        return WireResult::Success;
    }

    return buffer->mapState.Use([&](auto mapState) {
        mapState->memoryHandle = std::move(memoryHandle);
        return WireResult::Success;
    });
}

WireResult Server::DoBufferUpdateMappedData(Known<WGPUBuffer> buffer,
                                            Span<const std::byte> writeDataUpdateInfo,
                                            size_t offset,
                                            size_t size) {
    if (size == WGPU_WHOLE_MAP_SIZE) {
        return WireResult::FatalError;
    }

    // Buffers backed with shared memory are directly shared with client; no deserialization needed.
    if (buffer->backedWithSharedMemory) {
        return WireResult::Success;
    }

    return buffer->mapState.Use([&](auto mapState) {
        std::byte* mappedData =
            static_cast<std::byte*>(mProcs->bufferGetMappedRange(buffer->handle, offset, size));

        // There are a few valid reasons why getting the mapped range would fail here:
        //  - The buffer was implicitly unmapped because of a device.Destroy() call.
        //  - The buffer was an error buffer created just to replace an OOM mappedAtCreation buffer.
        // Unfortunately validating exactly that the failure is due to a valid reason and not
        // another is difficult, so we return WireResult::Success even for misuses of the wire
        // protocol (like a size being larger than the buffer's size, etc).
        if (mappedData == nullptr) {
            return WireResult::Success;
        }

        // SAFETY: If GetMappedRange returns non-null, it points to at least `size` valid bytes.
        Span<std::byte> DAWN_UNSAFE_BUFFERS(mappedRange{mappedData, size});

        // However it is easy to check for misuses of the wire protocol to UpdateMappedData without
        // a MemoryHandle.
        if (!mapState->memoryHandle) {
            return WireResult::FatalError;
        }

        // Deserialize the flush info and flush updated data from the handle into mappedRange.
        if (!mapState->memoryHandle->DeserializeDataUpdate(writeDataUpdateInfo, offset, size,
                                                           mappedRange)) {
            return WireResult::FatalError;
        }
        return WireResult::Success;
    });
}

void Server::OnBufferMapAsyncCallback(MapUserdata* data,
                                      wgpu::MapAsyncStatus status,
                                      StringView message) {
    // Skip sending the callback if the buffer has already been destroyed.
    Known<WGPUBuffer> buffer;
    if (Get(data->buffer.id, &buffer) != WireResult::Success ||
        buffer->generation != data->buffer.generation) {
        return;
    }

    bool isSuccess = status == wgpu::MapAsyncStatus::Success;

    ReturnBufferMapAsyncCallbackCmd cmd = {};
    cmd.instanceId = data->instanceId;
    cmd.future = data->future;
    cmd.status = status;
    cmd.message = message;

    if (!isSuccess) {
        SerializeCommand(std::move(cmd));
        return;
    }

    switch (data->mode) {
        case wgpu::MapMode::Read: {
            DAWN_ASSERT(data->size != WGPU_WHOLE_MAP_SIZE);  // Validated in DoBufferMapAsync.

            // Buffers backed with shared memory are directly shared with client; no deserialization
            // needed.
            if (buffer->backedWithSharedMemory) {
                SerializeCommand(std::move(cmd));
                break;
            }

            buffer->mapState.Use([&](auto mapState) {
                const std::byte* mappedData = static_cast<const std::byte*>(
                    mProcs->bufferGetConstMappedRange(buffer->handle, data->offset, data->size));

                // SAFETY: If GetConstMappedRange with size != WGPU_WHOLE_MAP_SIZE returns non-null,
                // it points to at least `size` valid bytes.
                Span<const std::byte> DAWN_UNSAFE_BUFFERS(mappedRange{mappedData, data->size});

                auto* memoryHandle = mapState->memoryHandle.get();
                if (!memoryHandle) {
                    cmd.status = wgpu::MapAsyncStatus::Aborted;
                    SerializeCommand(std::move(cmd));
                    return;
                }

                size_t dataUpdateInfoLength =
                    memoryHandle->GetSerializeDataUpdateSize(data->offset, data->size);
                SerializeCommand(
                    std::move(cmd),
                    // Extensions to replace fields skipped by skip_serialize.
                    CommandExtension<&ReturnBufferMapAsyncCallbackCmd::readDataUpdateInfo>{
                        dataUpdateInfoLength, [memoryHandle, data, mappedRange](
                                                  Span<volatile std::byte> serializeBuffer) {
                            // The in-flight map request returned successfully.
                            memoryHandle->SerializeDataUpdate(serializeBuffer, data->offset,
                                                              data->size, mappedRange);
                        }});
            });
            break;
        }
        case wgpu::MapMode::Write: {
            SerializeCommand(std::move(cmd));
            break;
        }
        default:
            // If we are not one of the two possible modes, we should never succeed.
            DAWN_UNREACHABLE();
            break;
    }
}

}  // namespace dawn::wire::server