// DKR-R optional mip generation. No subresource-state assumptions or per-frame waits.
#pragma once
#include <algorithm>
#include <array>
#include <cstring>
#include <stdexcept>
#include "rt64_texture.h"
#include "rt64_shader_library.h"
#include "rt64_descriptor_sets.h"
#include "rt64_generated_mip_config.h"
#include "native_resource_safety.hpp"
#include "../native/performance_switches.hpp"

namespace RT64::GeneratedMips {
    // Bound work per PNG and scratch retained by an entire native upload batch.
    constexpr uint64_t MaxTexels = 4ULL * 1024 * 1024;
    constexpr uint64_t MaxScratchBytes = 32ULL * 1024 * 1024;
    // Also bound allocation/descriptor overhead for batches of very small N64 textures.
    constexpr size_t MaxScratchTextures = 512;
    inline uint32_t levelCount(uint32_t w, uint32_t h) {
        if (!w || !h || uint64_t(w) * h > MaxTexels) return 1;
        uint32_t count = 1;
        while (w > 1 || h > 1) { w = std::max(w / 2, 1U); h = std::max(h / 2, 1U); ++count; }
        return count;
    }
    inline uint64_t levelBytes(uint32_t w, uint32_t h) { return ((uint64_t(w) * 4 + 255) & ~255ULL) * h; }
    struct Level { uint32_t width, height; std::vector<uint8_t> pixels; };
    // Integer alpha-weighted RGB; average alpha. Deliberately NOT a binary coverage filter.
    // Partition odd dimensions into disjoint boxes, including their final row/column.
    inline void downsample_into(const Level &src,Level &dst) {
        dst.width=std::max(src.width/2,1U);dst.height=std::max(src.height/2,1U);
        dst.pixels.resize(size_t(dst.width) * dst.height * 4);
        for (uint32_t y = 0; y < dst.height; ++y) for (uint32_t x = 0; x < dst.width; ++x) {
            uint32_t sums[4]{}, raw[3]{}, count = 0;
            for (uint32_t sy = y * src.height / dst.height; sy < (y + 1) * src.height / dst.height; ++sy)
                for (uint32_t sx = x * src.width / dst.width; sx < (x + 1) * src.width / dst.width; ++sx) {
                    const auto *p = &src.pixels[(size_t(sy) * src.width + sx) * 4];
                    for (unsigned c = 0; c < 3; ++c) { sums[c] += p[c] * p[3]; raw[c] += p[c]; }
                    sums[3] += p[3]; ++count;
                }
            auto *p = &dst.pixels[(size_t(y) * dst.width + x) * 4];
            for (unsigned c = 0; c < 3; ++c)
                p[c] = uint8_t(sums[3] ? (sums[c] + sums[3] / 2) / sums[3] : (raw[c] + count / 2) / count);
            p[3] = uint8_t((sums[3] + count / 2) / count);
        }
    }
    inline Level downsample(const Level& src) {Level dst;downsample_into(src,dst);return dst;}
// Included after GeneratedMips::Level/downsample; DKR's own extension only.
struct CpuMipScratch {
    std::vector<Level> levels;
    std::vector<std::uint64_t> offsets;
    bool busy=false;
};
struct CpuMipLease {
    CpuMipScratch& scratch;
    explicit CpuMipLease(CpuMipScratch& value):scratch(value){scratch.busy=true;}
    ~CpuMipLease() {
        std::uint64_t retained=0;
        for(const auto& level:scratch.levels)retained+=level.pixels.capacity();
        // Bound retained CPU payload per worker. Larger images still use the
        // unchanged upload path, but their temporary CPU storage is released.
        if(retained>8ULL*1024*1024) {std::vector<Level>().swap(scratch.levels);std::vector<std::uint64_t>().swap(scratch.offsets);}
        scratch.busy=false;
    }
};
inline CpuMipScratch& reusable_mip_scratch() {static thread_local CpuMipScratch value;return value;}

    struct Scratch {
        std::vector<std::unique_ptr<RenderTexture>> textures;
        std::vector<std::unique_ptr<TextureDecodeDescriptorSet>> descriptors;
        uint64_t bytes = 0;
    };
    inline bool uploadRGBA32(Texture &dst, RenderDevice *device, RenderCommandList *cmd,
                             const uint8_t *bytes, size_t byteCount, uint32_t w, uint32_t h, uint32_t rowPitch,
                             std::unique_ptr<RenderBuffer> &upload,bool single_level=false) {
        const uint32_t count = single_level ? 1U : levelCount(w,h);
        if((count<=1 && !single_level) || !bytes || !w || !h || w>16384 || h>16384 || uint64_t(w)*h>64ULL*1024*1024 || uint64_t(rowPitch)<uint64_t(w)*4 || uint64_t(rowPitch)*h>byteCount) return false;
        CpuMipScratch local;
        auto& reusable=reusable_mip_scratch();
        auto& scratch=dkr::runtime::surgical::mip_scratch() && !reusable.busy ? reusable : local;
        CpuMipLease cpu_lease(scratch);
        auto& levels=scratch.levels;auto& offsets=scratch.offsets;
        std::unique_ptr<RenderTexture> chain;
        std::unique_ptr<RenderBuffer> staging;
        uint64_t total=0;
        try {
            levels.resize(count);offsets.clear();offsets.reserve(count);
            levels[0].width=w;levels[0].height=h;levels[0].pixels.resize(size_t(w)*h*4);
            for(uint32_t y=0;y<h;++y) std::memcpy(levels[0].pixels.data()+size_t(y)*w*4,bytes+size_t(y)*rowPitch,size_t(w)*4);
            for(uint32_t m=1;m<count;++m) downsample_into(levels[m-1],levels[m]);
            for(const auto &l:levels) { offsets.push_back(total);total+=(levelBytes(l.width,l.height)+511)&~511ULL; }
            if(!single_level && total>MaxScratchBytes) return false;
            chain=device->createTexture(RenderTextureDesc::Texture2D(w,h,count,RenderFormat::R8G8B8A8_UNORM));
            staging=device->createBuffer(RenderBufferDesc::UploadBuffer(total));
            if(!resource_valid(chain.get()) || !resource_valid(staging.get())) return false;
        } catch(const std::bad_alloc &) { return false; }
        auto *data=static_cast<uint8_t*>(staging->map());
        if(!data) return false;
        for(uint32_t m=0;m<count;++m) { const auto &l=levels[m]; const uint64_t pitch=(uint64_t(l.width)*4+255)&~255ULL;
            for(uint32_t y=0;y<l.height;++y)std::memcpy(data+offsets[m]+y*pitch,l.pixels.data()+size_t(y)*l.width*4,size_t(l.width)*4); }
        staging->unmap();
        cmd->barriers(RenderBarrierStage::COPY,RenderTextureBarrier(chain.get(),RenderTextureLayout::COPY_DEST));
        for(uint32_t m=0;m<count;++m) { const auto &l=levels[m];
            cmd->copyTextureRegion(RenderTextureCopyLocation::Subresource(chain.get(),m),RenderTextureCopyLocation::PlacedFootprint(staging.get(),RenderFormat::R8G8B8A8_UNORM,l.width,l.height,1,((l.width*4+255)&~255U)/4,offsets[m])); }
        dst.texture=std::move(chain);upload=std::move(staging);dst.width=w;dst.height=h;
        dst.format=RenderFormat::R8G8B8A8_UNORM;dst.mipmaps=count;dst.memorySize=total;dst.generatedMipmaps=count>1;
        return true;
    }
    // Record only after native RGBA decode. Keep Scratch alive until the upload fence.
    // Every read/write pair is two different textures; final chain stays COPY_DEST.
    inline bool record(Texture &dst, RenderDevice *device, RenderCommandList *cmd,
                       const ShaderRecord &shader, Scratch &retained) {
        const uint32_t count = levelCount(dst.width, dst.height);
        if (count <= 1 || !shader.pipeline || !dst.texture) return false;
        uint64_t scratchBytes = levelBytes(dst.width, dst.height), finalBytes = scratchBytes;
        uint32_t w = dst.width, h = dst.height;
        for (uint32_t m = 1; m < count; ++m) {
            w = std::max(w / 2, 1U); h = std::max(h / 2, 1U);
            scratchBytes += levelBytes(w, h); finalBytes += levelBytes(w, h);
        }
        if (retained.bytes + scratchBytes > MaxScratchBytes ||
            retained.textures.size() + count > MaxScratchTextures) return false;
        std::vector<std::unique_ptr<RenderTexture>> scratch;
        std::vector<std::unique_ptr<TextureDecodeDescriptorSet>> sets;
        std::unique_ptr<RenderTexture> chain;
        // Allocate everything before recording or changing the original Texture.
        try {
            chain = device->createTexture(RenderTextureDesc::Texture2D(dst.width, dst.height, count, dst.format));
            if (!resource_valid(chain.get())) return false;
            w = dst.width; h = dst.height;
            for (uint32_t m = 1; m < count; ++m) {
                w = std::max(w / 2, 1U); h = std::max(h / 2, 1U);
                auto tex = device->createTexture(RenderTextureDesc::Texture2D(w, h, 1, dst.format,
                    RenderTextureFlag::STORAGE | RenderTextureFlag::UNORDERED_ACCESS));
                if (!resource_valid(tex.get())) return false;
                auto set = std::make_unique<TextureDecodeDescriptorSet>(device);
                set->setTexture(set->TMEM, m == 1 ? dst.texture.get() : scratch.back().get(), RenderTextureLayout::SHADER_READ);
                set->setTexture(set->RGBA32, tex.get(), RenderTextureLayout::GENERAL);
                scratch.emplace_back(std::move(tex)); sets.emplace_back(std::move(set));
            }
            retained.textures.reserve(retained.textures.size() + count);
            retained.descriptors.reserve(retained.descriptors.size() + count - 1);
        } catch (const std::bad_alloc &) { return false; }
        cmd->barriers(RenderBarrierStage::COPY, RenderTextureBarrier(chain.get(), RenderTextureLayout::COPY_DEST));
        cmd->barriers(RenderBarrierStage::COPY, RenderTextureBarrier(dst.texture.get(), RenderTextureLayout::COPY_SOURCE));
        cmd->copyTextureRegion(RenderTextureCopyLocation::Subresource(chain.get(), 0), RenderTextureCopyLocation::Subresource(dst.texture.get()));
        cmd->barriers(RenderBarrierStage::COMPUTE, RenderTextureBarrier(dst.texture.get(), RenderTextureLayout::SHADER_READ));
        cmd->setComputePipelineLayout(shader.pipelineLayout.get()); cmd->setPipeline(shader.pipeline.get());
        w = dst.width; h = dst.height;
        for (uint32_t m = 1; m < count; ++m) {
            const uint32_t nw = std::max(w / 2, 1U), nh = std::max(h / 2, 1U);
            const uint32_t dims[4]{w, h, nw, nh};
            auto *target = scratch[m - 1].get();
            cmd->barriers(RenderBarrierStage::COMPUTE, RenderTextureBarrier(target, RenderTextureLayout::GENERAL));
            cmd->setComputePushConstants(0, dims); cmd->setComputeDescriptorSet(sets[m - 1]->get(), 0);
            cmd->dispatch((nw + 7) / 8, (nh + 7) / 8, 1);
            cmd->barriers(RenderBarrierStage::COPY, RenderTextureBarrier(target, RenderTextureLayout::COPY_SOURCE));
            cmd->copyTextureRegion(RenderTextureCopyLocation::Subresource(chain.get(), m), RenderTextureCopyLocation::Subresource(target));
            if (m + 1 < count) cmd->barriers(RenderBarrierStage::COMPUTE, RenderTextureBarrier(target, RenderTextureLayout::SHADER_READ));
            w = nw; h = nh;
        }
        cmd->barriers(RenderBarrierStage::GRAPHICS_AND_COMPUTE, RenderTextureBarrier(chain.get(), RenderTextureLayout::SHADER_READ));
        retained.textures.emplace_back(std::move(dst.texture));
        for (auto &t : scratch) retained.textures.emplace_back(std::move(t));
        for (auto &s : sets) retained.descriptors.emplace_back(std::move(s));
        retained.bytes += scratchBytes;
        dst.texture = std::move(chain); dst.mipmaps = count; dst.memorySize = finalBytes;
        dst.generatedMipmaps = true;
        return true;
    }
}
