#include <XCEngine/Rendering/Graph/RenderGraphCompiler.h>

#include <algorithm>
#include <unordered_set>
#include <utility>

namespace XCEngine {
namespace Rendering {
namespace {

RHI::ResourceStates ResolveRequiredState(
    RenderGraphPassType passType,
    RenderGraphAccessMode accessMode,
    RenderGraphTextureAspect aspect) {
    if (accessMode == RenderGraphAccessMode::Write) {
        if (passType == RenderGraphPassType::Raster &&
            aspect == RenderGraphTextureAspect::Depth) {
            return RHI::ResourceStates::DepthWrite;
        }

        return passType == RenderGraphPassType::Compute
            ? RHI::ResourceStates::UnorderedAccess
            : RHI::ResourceStates::RenderTarget;
    }

    if (passType == RenderGraphPassType::Raster &&
        aspect == RenderGraphTextureAspect::Depth) {
        return RHI::ResourceStates::DepthRead;
    }

    return passType == RenderGraphPassType::Compute
        ? RHI::ResourceStates::GenericRead
        : RHI::ResourceStates::PixelShaderResource;
}

void WriteError(
    const Containers::String& message,
    Containers::String* outErrorMessage) {
    if (outErrorMessage != nullptr) {
        *outErrorMessage = message;
    }
}

} // namespace

void CompiledRenderGraph::Reset() {
    m_passes.clear();
    m_textures.clear();
    m_textureLifetimes.clear();
    m_textureTransitionPlans.clear();
}

const Containers::String& CompiledRenderGraph::GetPassName(size_t index) const {
    static const Containers::String kEmptyString;
    return index < m_passes.size() ? m_passes[index].name : kEmptyString;
}

RenderGraphPassType CompiledRenderGraph::GetPassType(size_t index) const {
    return index < m_passes.size()
        ? m_passes[index].type
        : RenderGraphPassType::Raster;
}

bool CompiledRenderGraph::TryGetTextureLifetime(
    RenderGraphTextureHandle handle,
    RenderGraphTextureLifetime& outLifetime) const {
    if (!handle.IsValid() || handle.index >= m_textureLifetimes.size()) {
        return false;
    }

    outLifetime = m_textureLifetimes[handle.index];
    return true;
}

bool CompiledRenderGraph::TryGetImportedTextureOptions(
    RenderGraphTextureHandle handle,
    RenderGraphImportedTextureOptions& outOptions) const {
    if (!handle.IsValid() || handle.index >= m_textures.size()) {
        return false;
    }

    const CompiledRenderGraph::CompiledTexture& texture = m_textures[handle.index];
    if (texture.kind != RenderGraphTextureKind::Imported) {
        return false;
    }

    outOptions = texture.importedOptions;
    return true;
}

bool CompiledRenderGraph::TryGetTextureTransitionPlan(
    RenderGraphTextureHandle handle,
    RenderGraphTextureTransitionPlan& outPlan) const {
    if (!handle.IsValid() || handle.index >= m_textureTransitionPlans.size()) {
        return false;
    }

    outPlan = m_textureTransitionPlans[handle.index];
    return true;
}

bool RenderGraphCompiler::Compile(
    const RenderGraph& graph,
    CompiledRenderGraph& outCompiledGraph,
    Containers::String* outErrorMessage) {
    outCompiledGraph.Reset();

    for (const RenderGraph::TextureResource& texture : graph.m_textures) {
        if (!texture.desc.IsValid()) {
            WriteError(
                Containers::String("RenderGraph texture desc is invalid: ") + texture.name,
                outErrorMessage);
            return false;
        }

        if (texture.kind == RenderGraphTextureKind::Imported &&
            texture.importedOptions.graphOwnsTransitions &&
            texture.importedView == nullptr) {
            WriteError(
                Containers::String("RenderGraph imported texture requires a valid view when graph owns transitions: ") +
                    texture.name,
                outErrorMessage);
            return false;
        }
    }

    const size_t passCount = graph.m_passes.size();
    const size_t textureCount = graph.m_textures.size();

    std::vector<std::vector<Core::uint32>> outgoingEdges(passCount);
    std::vector<Core::uint32> incomingEdgeCount(passCount, 0u);
    std::unordered_set<Core::uint64> edgeKeys;
    std::vector<Core::uint32> lastWriter(textureCount, kInvalidRenderGraphHandle);
    std::vector<std::vector<Core::uint32>> lastReaders(textureCount);

    auto addUniqueReader = [](
        std::vector<Core::uint32>& readers,
        Core::uint32 passIndex) {
        if (std::find(readers.begin(), readers.end(), passIndex) == readers.end()) {
            readers.push_back(passIndex);
        }
    };

    auto addEdge = [&](
        Core::uint32 from,
        Core::uint32 to) {
        if (from == kInvalidRenderGraphHandle || to == kInvalidRenderGraphHandle || from == to) {
            return;
        }

        const Core::uint64 key =
            (static_cast<Core::uint64>(from) << 32u) | static_cast<Core::uint64>(to);
        if (!edgeKeys.insert(key).second) {
            return;
        }

        outgoingEdges[from].push_back(to);
        ++incomingEdgeCount[to];
    };

    for (Core::uint32 passIndex = 0u; passIndex < static_cast<Core::uint32>(passCount); ++passIndex) {
        const RenderGraph::PassNode& pass = graph.m_passes[passIndex];
        for (const RenderGraph::TextureAccess& access : pass.accesses) {
            if (!access.texture.IsValid() || access.texture.index >= textureCount) {
                WriteError(
                    Containers::String("RenderGraph pass '") + pass.name +
                        "' references an invalid texture handle",
                    outErrorMessage);
                return false;
            }

            const RenderGraph::TextureResource& texture = graph.m_textures[access.texture.index];
            std::vector<Core::uint32>& readers = lastReaders[access.texture.index];
            Core::uint32& writer = lastWriter[access.texture.index];

            if (access.mode == RenderGraphAccessMode::Read) {
                if (texture.kind == RenderGraphTextureKind::Transient &&
                    writer == kInvalidRenderGraphHandle) {
                    WriteError(
                        Containers::String("RenderGraph transient texture '") + texture.name +
                            "' is read before any pass writes it",
                        outErrorMessage);
                    return false;
                }

                addEdge(writer, passIndex);
                addUniqueReader(readers, passIndex);
                continue;
            }

            addEdge(writer, passIndex);
            for (Core::uint32 readerPassIndex : readers) {
                addEdge(readerPassIndex, passIndex);
            }
            readers.clear();
            writer = passIndex;
        }
    }

    std::vector<Core::uint32> executionOrder;
    executionOrder.reserve(passCount);
    std::vector<bool> emitted(passCount, false);

    while (executionOrder.size() < passCount) {
        bool progressed = false;
        for (Core::uint32 passIndex = 0u; passIndex < static_cast<Core::uint32>(passCount); ++passIndex) {
            if (emitted[passIndex] || incomingEdgeCount[passIndex] != 0u) {
                continue;
            }

            emitted[passIndex] = true;
            executionOrder.push_back(passIndex);
            for (Core::uint32 dependentPassIndex : outgoingEdges[passIndex]) {
                if (incomingEdgeCount[dependentPassIndex] > 0u) {
                    --incomingEdgeCount[dependentPassIndex];
                }
            }

            progressed = true;
            break;
        }

        if (!progressed) {
            WriteError(
                "RenderGraph failed to compile because pass dependencies contain a cycle",
                outErrorMessage);
            outCompiledGraph.Reset();
            return false;
        }
    }

    outCompiledGraph.m_textureLifetimes.resize(textureCount);
    outCompiledGraph.m_textureTransitionPlans.resize(textureCount);
    for (size_t textureIndex = 0u; textureIndex < textureCount; ++textureIndex) {
        CompiledRenderGraph::CompiledTexture compiledTexture = {};
        compiledTexture.name = graph.m_textures[textureIndex].name;
        compiledTexture.desc = graph.m_textures[textureIndex].desc;
        compiledTexture.kind = graph.m_textures[textureIndex].kind;
        compiledTexture.importedView = graph.m_textures[textureIndex].importedView;
        compiledTexture.importedOptions = graph.m_textures[textureIndex].importedOptions;
        outCompiledGraph.m_textures.push_back(std::move(compiledTexture));

        outCompiledGraph.m_textureLifetimes[textureIndex].kind = graph.m_textures[textureIndex].kind;
        RenderGraphTextureTransitionPlan& transitionPlan =
            outCompiledGraph.m_textureTransitionPlans[textureIndex];
        transitionPlan.graphOwnsTransitions =
            graph.m_textures[textureIndex].kind == RenderGraphTextureKind::Transient ||
            graph.m_textures[textureIndex].importedOptions.graphOwnsTransitions;
        transitionPlan.initialState =
            graph.m_textures[textureIndex].kind == RenderGraphTextureKind::Imported
                ? graph.m_textures[textureIndex].importedOptions.initialState
                : RHI::ResourceStates::Common;
        transitionPlan.finalState =
            graph.m_textures[textureIndex].kind == RenderGraphTextureKind::Imported
                ? graph.m_textures[textureIndex].importedOptions.finalState
                : RHI::ResourceStates::Common;
    }

    outCompiledGraph.m_passes.reserve(passCount);
    for (Core::uint32 compiledPassIndex = 0u;
         compiledPassIndex < static_cast<Core::uint32>(executionOrder.size());
         ++compiledPassIndex) {
        const Core::uint32 originalPassIndex = executionOrder[compiledPassIndex];
        const RenderGraph::PassNode& sourcePass = graph.m_passes[originalPassIndex];

        CompiledRenderGraph::CompiledPass compiledPass = {};
        compiledPass.name = sourcePass.name;
        compiledPass.type = sourcePass.type;
        compiledPass.originalPassIndex = originalPassIndex;
        compiledPass.accesses.reserve(sourcePass.accesses.size());
        for (const RenderGraph::TextureAccess& access : sourcePass.accesses) {
            CompiledRenderGraph::CompiledTextureAccess compiledAccess = {};
            compiledAccess.texture = access.texture;
            compiledAccess.mode = access.mode;
            compiledAccess.aspect = access.aspect;
            compiledAccess.requiredState =
                ResolveRequiredState(sourcePass.type, access.mode, access.aspect);
            compiledPass.accesses.push_back(compiledAccess);
        }
        compiledPass.executeCallback = sourcePass.executeCallback;
        outCompiledGraph.m_passes.push_back(std::move(compiledPass));

        for (const RenderGraph::TextureAccess& access : sourcePass.accesses) {
            RenderGraphTextureLifetime& lifetime =
                outCompiledGraph.m_textureLifetimes[access.texture.index];
            RenderGraphTextureTransitionPlan& transitionPlan =
                outCompiledGraph.m_textureTransitionPlans[access.texture.index];
            if (!lifetime.used) {
                lifetime.used = true;
                lifetime.firstPassIndex = compiledPassIndex;
            }
            lifetime.lastPassIndex = compiledPassIndex;

            const RHI::ResourceStates accessState =
                ResolveRequiredState(sourcePass.type, access.mode, access.aspect);
            if (!transitionPlan.hasFirstAccessState) {
                transitionPlan.hasFirstAccessState = true;
                transitionPlan.firstAccessState = accessState;
            }
            transitionPlan.hasLastAccessState = true;
            transitionPlan.lastAccessState = accessState;
            if (graph.m_textures[access.texture.index].kind == RenderGraphTextureKind::Transient) {
                transitionPlan.finalState = accessState;
            }
        }
    }

    if (outErrorMessage != nullptr) {
        outErrorMessage->Clear();
    }
    return true;
}

} // namespace Rendering
} // namespace XCEngine