#include "dx_texture.h"

#include "dx_renderer.h"
#include "misc/scope_exit.h"

int get_cpu_access_flag(device_memory_type in_memory_type) {
    switch (in_memory_type) {
    case device_memory_type::gpu:
        return 0;
    case device_memory_type::shared:
        return D3D11_CPU_ACCESS_WRITE;
    case device_memory_type::cpu:
        return D3D11_CPU_ACCESS_WRITE;
    default:
        return D3D11_CPU_ACCESS_WRITE;
    }
}

dx_texture::dx_texture(const Eigen::Vector2i& size, const texture_format format, device_memory_type in_memory_type, uint32_t count) : renderer_texture(format) {
    memory_type = in_memory_type;
    create_texture(size, format, in_memory_type, count);
}

dx_texture::~dx_texture() {
    release_texture();
}

void* dx_texture::lock(uint32_t* out_row_pitch) {
    return lock(0, out_row_pitch);
}

void dx_texture::unlock() {
    return unlock(0);
}

void* dx_texture::lock(uint32_t index, uint32_t* out_row_pitch) {
    auto r = aorii::get_renderer<dx_renderer>();
    auto hr = r->get_d3d_device()->CreateTexture2D(&upload_desc, nullptr, &upload_texture);
    if (FAILED(hr)) {
        spdlog::critical("无法创建上传纹理, 错误: 0x{:08X}", static_cast<unsigned int>(hr));
        return nullptr;
    }

    D3D11_MAPPED_SUBRESOURCE mappedResource;
    const auto context = aorii::get_renderer<dx_renderer>()->get_d3d_context();
    hr = context->Map(upload_texture, index, D3D11_MAP_WRITE_DISCARD, 0, &mappedResource);
    if (FAILED(hr)) {
        spdlog::critical("无法映射纹理, 错误: 0x{:08X}", static_cast<unsigned int>(hr));
        return nullptr;
    }
    if (out_row_pitch) *out_row_pitch = mappedResource.RowPitch;
    return mappedResource.pData;
}

void dx_texture::unlock(uint32_t index) const {
    const auto context = aorii::get_renderer<dx_renderer>()->get_d3d_context();
    context->Unmap(upload_texture, index);
    // 复制到目标纹理
    context->CopyResource(m_texture, upload_texture);
}

void dx_texture::update_subresource(const Eigen::AlignedBox2i& in_area, const void* in_buffer, const uint32_t in_row_pitch) {
    update_subresource(0, in_area, in_buffer, in_row_pitch);
}

void dx_texture::update_subresource(uint32_t index, const Eigen::AlignedBox2i& in_area, const void* in_buffer, const uint32_t in_row_pitch) {
    auto r = aorii::get_renderer<dx_renderer>();
    const auto context = r->get_d3d_context();
    auto device = r->get_d3d_device();
    const auto& size = in_area.sizes();

    // 更新上传纹理描述符
    upload_desc.Width = size.x();
    upload_desc.Height = size.y();

    // 定义目标区域
    D3D11_BOX box{};
    box.left = in_area.min().x();
    box.top = in_area.min().y();
    box.right = in_area.max().x();
    box.bottom = in_area.max().y();
    box.front = 0;
    box.back = 1;

    // 定义源区域
    D3D11_BOX src_box{};
    src_box.left = 0;
    src_box.top = 0;
    src_box.right = size.x();
    src_box.bottom = size.y();
    src_box.front = 0;
    src_box.back = 1;

    // 创建临时上传纹理
    auto hr = device->CreateTexture2D(&upload_desc, nullptr, &upload_texture);
    if (FAILED(hr)) {
        spdlog::critical("无法创建上传纹理, 错误: 0x{:08X}", static_cast<unsigned int>(hr));
        return;
    }
    ON_SCOPE_EXIT{
        // 删除临时纹理
        upload_texture->Release();
        upload_texture = nullptr;
    };

    // 映射上传纹理
    D3D11_MAPPED_SUBRESOURCE mapped;
    hr = context->Map(upload_texture, 0, D3D11_MAP_WRITE, 0, &mapped);
    if (FAILED(hr)) {
        spdlog::critical("无法映射上传纹理, 错误: 0x{:08X}", static_cast<unsigned int>(hr));
        return;
    }

    // 复制数据到上传纹理
    const auto src_data = static_cast<const uint8_t*>(in_buffer);
    const auto dst_data = static_cast<uint8_t*>(mapped.pData);

    for (int y = 0; y < size.y(); ++y) {
        memcpy(
            dst_data + mapped.RowPitch * y,
            src_data + in_row_pitch * y,
            size.x()
        );
    }

    // 解除映射
    context->Unmap(upload_texture, 0);

    // 复制到目标纹理
    context->CopySubresourceRegion(
        m_texture,                          // 目标纹理
        index,                              // 目标子资源索引
        box.left, box.top, 0,     // 目标位置
        upload_texture,                     // 源纹理
        0,                                  // 源子资源索引
        &src_box                            // 源区域
    );
}

bool dx_texture::resize(const Eigen::Vector2i& size) {
    D3D11_TEXTURE2D_DESC desc;
    m_texture->GetDesc(&desc);
    if (desc.Width == size.x() && desc.Height == size.y()) {
        return true;
    }
    return create_texture(size, get_format(), memory_type, get_count());
}

Eigen::Vector2i dx_texture::size() {
    D3D11_TEXTURE2D_DESC desc;
    m_texture->GetDesc(&desc);
    return { desc.Width, desc.Height };
}

uint32_t dx_texture::get_count() const {
    D3D11_TEXTURE2D_DESC desc;
    m_texture->GetDesc(&desc);
    return desc.ArraySize;
}

void dx_texture::set_count(const uint32_t in_count) {
    if (in_count == get_count()) {
        return;
    }
    D3D11_TEXTURE2D_DESC desc;
    m_texture->GetDesc(&desc);
    desc.ArraySize = in_count;
    create_texture({ desc.Width, desc.Height }, get_format(), memory_type, in_count);
}

bool dx_texture::create_texture(const Eigen::Vector2i& in_size, texture_format in_format, device_memory_type in_memory_type, uint32_t in_count) {
    // 释放旧资源
    release_texture();

    // 创建新资源
    auto d3d_device = aorii::get_renderer<dx_renderer>()->get_d3d_device();
    D3D11_TEXTURE2D_DESC desc = {};
    desc.Width = in_size.x();
    desc.Height = in_size.y();
    desc.MipLevels = 1;
    desc.ArraySize = in_count;
    desc.Format = get_dxgi_format(in_format);
    desc.SampleDesc.Count = 1;
    desc.SampleDesc.Quality = 0;
    desc.Usage = D3D11_USAGE_DEFAULT;
    desc.BindFlags = D3D11_BIND_SHADER_RESOURCE;
    desc.CPUAccessFlags = get_cpu_access_flag(in_memory_type);
    desc.MiscFlags = 0;
    auto hr = d3d_device->CreateTexture2D(&desc, nullptr, &m_texture);
    if (FAILED(hr)) {
        spdlog::critical("无法创建纹理, 错误: 0x{:08X}", static_cast<unsigned int>(hr));
        return false;
    }

    // 创建临时的上传纹理
    upload_desc = desc;
    upload_desc.ArraySize = 1;
    upload_desc.Usage = D3D11_USAGE_STAGING;  // 用于CPU到GPU的传输
    upload_desc.BindFlags = 0;
    upload_desc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;

    D3D11_SHADER_RESOURCE_VIEW_DESC srv_desc = {};
    srv_desc.Format = desc.Format;
    srv_desc.ViewDimension = D3D11_SRV_DIMENSION_TEXTURE2D;
    srv_desc.Texture2D.MostDetailedMip = 0;
    srv_desc.Texture2D.MipLevels = 1;
    hr = d3d_device->CreateShaderResourceView(m_texture, &srv_desc, &srv);
    if (FAILED(hr)) {
        spdlog::critical("无法创建着色器资源视图, 错误: 0x{:08X}", static_cast<unsigned int>(hr));
        return false;
    }
    return true;
}

void dx_texture::release_texture() {
    if (srv) srv->Release();
    if (m_texture) m_texture->Release();
    if (upload_texture) upload_texture->Release();
    srv = nullptr;
    m_texture = nullptr;
    upload_texture = nullptr;
}
