#include "ShaderAuthoringInternal.h"

#include "../ShaderSourceUtils.h"

#include <sstream>
#include <string>
#include <unordered_set>
#include <vector>

namespace XCEngine {
namespace Resources {
namespace Internal {

namespace {

Containers::String BuildShaderKeywordSetSignature(const ShaderKeywordSet& keywordSet) {
    ShaderKeywordSet normalizedKeywords = keywordSet;
    NormalizeShaderKeywordSetInPlace(normalizedKeywords);

    Containers::String signature;
    for (size_t keywordIndex = 0; keywordIndex < normalizedKeywords.enabledKeywords.Size(); ++keywordIndex) {
        if (keywordIndex > 0) {
            signature += ";";
        }

        signature += normalizedKeywords.enabledKeywords[keywordIndex];
    }

    return signature;
}

} // namespace

void AppendAuthoringSourceBlock(
    Containers::String& target,
    const Containers::String& sourceBlock) {
    if (sourceBlock.Empty()) {
        return;
    }

    if (!target.Empty()) {
        target += '\n';
    }
    target += sourceBlock;
}

void CollectQuotedIncludeDependencyPaths(
    const Containers::String& sourcePath,
    const Containers::String& sourceText,
    std::unordered_set<std::string>& seenPaths,
    Containers::Array<Containers::String>& outDependencies) {
    std::istringstream stream(ToStdString(sourceText));
    std::string rawLine;
    while (std::getline(stream, rawLine)) {
        const std::string line = TrimCopy(StripAuthoringLineComment(rawLine));
        if (line.rfind("#include", 0) != 0) {
            continue;
        }

        const size_t firstQuote = line.find('"');
        if (firstQuote == std::string::npos) {
            continue;
        }

        const size_t secondQuote = line.find('"', firstQuote + 1);
        if (secondQuote == std::string::npos || secondQuote <= firstQuote + 1) {
            continue;
        }

        const Containers::String includePath(line.substr(firstQuote + 1, secondQuote - firstQuote - 1).c_str());
        const Containers::String resolvedPath = ResolveShaderDependencyPath(includePath, sourcePath);
        const std::string key = ToStdString(resolvedPath);
        if (!key.empty() && seenPaths.insert(key).second) {
            outDependencies.PushBack(resolvedPath);
        }
    }
}

bool IsShaderAuthoringPragmaDirective(const std::string& line) {
    std::vector<std::string> pragmaTokens;
    if (!TryTokenizeQuotedArguments(line, pragmaTokens) ||
        pragmaTokens.empty() ||
        pragmaTokens[0] != "#pragma" ||
        pragmaTokens.size() < 2u) {
        return false;
    }

    return pragmaTokens[1] == "vertex" ||
           pragmaTokens[1] == "fragment" ||
           pragmaTokens[1] == "compute" ||
           pragmaTokens[1] == "target" ||
           pragmaTokens[1] == "multi_compile" ||
           pragmaTokens[1] == "multi_compile_local" ||
           pragmaTokens[1] == "shader_feature" ||
           pragmaTokens[1] == "shader_feature_local" ||
           pragmaTokens[1] == "backend";
}

Containers::String StripShaderAuthoringPragmas(const Containers::String& sourceText) {
    std::istringstream stream(ToStdString(sourceText));
    std::string rawLine;
    std::string strippedSource;

    while (std::getline(stream, rawLine)) {
        const std::string normalizedLine = TrimCopy(StripAuthoringLineComment(rawLine));
        if (IsShaderAuthoringPragmaDirective(normalizedLine)) {
            continue;
        }

        strippedSource += rawLine;
        strippedSource += '\n';
    }

    return strippedSource.c_str();
}

std::vector<ShaderKeywordSet> BuildShaderKeywordVariantSets(
    const Containers::Array<ShaderKeywordDeclaration>& declarations) {
    std::vector<ShaderKeywordSet> keywordSets(1);

    for (const ShaderKeywordDeclaration& declaration : declarations) {
        if (declaration.options.Empty()) {
            continue;
        }

        std::vector<ShaderKeywordSet> nextKeywordSets;
        std::unordered_set<std::string> seenSignatures;
        nextKeywordSets.reserve(keywordSets.size() * declaration.options.Size());

        for (const ShaderKeywordSet& currentKeywordSet : keywordSets) {
            for (const Containers::String& option : declaration.options) {
                ShaderKeywordSet nextKeywordSet = currentKeywordSet;
                const Containers::String normalizedKeyword = NormalizeShaderKeywordToken(option);
                if (!normalizedKeyword.Empty()) {
                    nextKeywordSet.enabledKeywords.PushBack(normalizedKeyword);
                }

                NormalizeShaderKeywordSetInPlace(nextKeywordSet);
                const std::string signature =
                    ToStdString(BuildShaderKeywordSetSignature(nextKeywordSet));
                if (seenSignatures.insert(signature).second) {
                    nextKeywordSets.push_back(std::move(nextKeywordSet));
                }
            }
        }

        if (!nextKeywordSets.empty()) {
            keywordSets = std::move(nextKeywordSets);
        }
    }

    if (keywordSets.empty()) {
        keywordSets.emplace_back();
    }

    return keywordSets;
}

Containers::String BuildKeywordVariantSource(
    const Containers::String& baseSource,
    const ShaderKeywordSet& requiredKeywords) {
    if (requiredKeywords.enabledKeywords.Empty()) {
        return baseSource;
    }

    std::string defineBlock;
    for (const Containers::String& keyword : requiredKeywords.enabledKeywords) {
        defineBlock += "#define ";
        defineBlock += ToStdString(keyword);
        defineBlock += " 1\n";
    }

    if (baseSource.Empty()) {
        return Containers::String(defineBlock.c_str());
    }

    const std::string sourceText = ToStdString(baseSource);

    size_t lineStart = 0;
    while (lineStart < sourceText.size()) {
        const size_t lineEnd = sourceText.find_first_of("\r\n", lineStart);
        const size_t lineLength =
            lineEnd == std::string::npos ? sourceText.size() - lineStart : lineEnd - lineStart;
        const std::string line = sourceText.substr(lineStart, lineLength);
        const std::string trimmedLine = TrimCopy(line);

        if (trimmedLine.empty() ||
            trimmedLine.rfind("//", 0) == 0u ||
            trimmedLine.rfind("/*", 0) == 0u ||
            trimmedLine.rfind("*", 0) == 0u) {
            if (lineEnd == std::string::npos) {
                break;
            }

            lineStart = lineEnd + 1u;
            if (sourceText[lineEnd] == '\r' &&
                lineStart < sourceText.size() &&
                sourceText[lineStart] == '\n') {
                ++lineStart;
            }
            continue;
        }

        if (trimmedLine.rfind("#version", 0) != 0u) {
            break;
        }

        size_t insertionPos = lineEnd == std::string::npos ? sourceText.size() : lineEnd + 1u;
        if (lineEnd != std::string::npos &&
            sourceText[lineEnd] == '\r' &&
            insertionPos < sourceText.size() &&
            sourceText[insertionPos] == '\n') {
            ++insertionPos;
        }

        std::string variantSource = sourceText.substr(0, insertionPos);
        variantSource += defineBlock;
        variantSource += sourceText.substr(insertionPos);
        return Containers::String(variantSource.c_str());
    }

    std::string variantSource = defineBlock;
    variantSource += '\n';
    variantSource += sourceText;
    return Containers::String(variantSource.c_str());
}

} // namespace Internal
} // namespace Resources
} // namespace XCEngine