#include "shader.h"
#include <tool.h>
#include <d3d11shader.h>
#include <d3dcompiler.h>
#include <iostream>
#include <render_type_def.h>
#include <span>
#include "states_render.h"

//MARK: compile
namespace {
constexpr bool HasEntry_(Shader::Entry value, Shader::Entry flag) {
    using T = std::underlying_type_t<Shader::Entry>;
    return (static_cast<T>(value & flag) != 0);
}

UINT GetCompileFlags() {
    UINT flags = D3DCOMPILE_WARNINGS_ARE_ERRORS | D3DCOMPILE_ENABLE_STRICTNESS;
#ifdef _DEBUG
    flags |= D3DCOMPILE_SKIP_OPTIMIZATION;
    flags |= D3DCOMPILE_DEBUG;
#else
    flags |= D3DCOMPILE_OPTIMIZATION_LEVEL3;
#endif
    return flags;
}

HRESULT CompileFromFile_(
    const std::filesystem::path &path,
    const std::string &entryPoint,
    const std::string &target,
    Microsoft::WRL::ComPtr<ID3DBlob> &byteCode,
    const D3D_SHADER_MACRO *macros = nullptr,
    ID3DInclude *includeHandler = D3D_COMPILE_STANDARD_FILE_INCLUDE) {

    byteCode.Reset();
    Microsoft::WRL::ComPtr<ID3DBlob> blob_error{};

    HRESULT hr = D3DCompileFromFile(
        path.wstring().c_str(),
        macros,
        includeHandler,
        entryPoint.c_str(),
        target.c_str(),
        GetCompileFlags(),
        0,
        byteCode.GetAddressOf(),
        blob_error.GetAddressOf()
    );
    if (FAILED(hr)) {
        std::wstring msg = Z::HResultToWString(hr);
        if (blob_error) {
            const char *err = static_cast<const char *>(blob_error->GetBufferPointer());

            int len = MultiByteToWideChar(CP_UTF8, 0, err, -1, nullptr, 0);
            if (len > 0) {
                std::wstring werr(len, L'\0');
                MultiByteToWideChar(CP_UTF8, 0, err, -1, werr.data(), len);

                msg += L"\n\n";
                msg += werr;
            } else {
                msg += L"\n\n";
                msg += L"Shader compile error blob exists, but convert failed.";
            }
        }
        MessageBoxW(nullptr, msg.c_str(), L"Compile Shader Failed", MB_OK);
        return hr;
    }
    return S_OK;
}
}

//MARK: 反射CBuffer
namespace {
uint32_t FindCbBindSlot(ID3D11ShaderReflection *reflector, const char *cbufferName) {
    D3D11_SHADER_DESC shaderDesc{};
    (void)reflector->GetDesc(&shaderDesc);

    for (UINT i = 0; i < shaderDesc.BoundResources; ++i) {
        D3D11_SHADER_INPUT_BIND_DESC bindDesc{};
        (void)reflector->GetResourceBindingDesc(i, &bindDesc);

        if (bindDesc.Type == D3D_SIT_CBUFFER && std::strcmp(bindDesc.Name, cbufferName) == 0) {
            return bindDesc.BindPoint;
        }
    }
    return UINT32_MAX;
}

std::vector<Shader::CBufferDesc> ReflectCBuffers_(ID3DBlob *shaderBlob) {
    std::vector<Shader::CBufferDesc> result;

    Microsoft::WRL::ComPtr<ID3D11ShaderReflection> reflector;

    HRESULT hr = D3DReflect(
        shaderBlob->GetBufferPointer(),
        shaderBlob->GetBufferSize(),
        IID_PPV_ARGS(reflector.GetAddressOf())
    );

    if (FAILED(hr)) {
        MessageBoxW(nullptr, L"Reflect Shader Failed", L"ERROR",MB_OK);
        return result;
    }

    D3D11_SHADER_DESC shaderDesc{};
    (void)reflector->GetDesc(&shaderDesc);

    // 反射Cbuffer
    for (UINT i = 0; i < shaderDesc.ConstantBuffers; ++i) {
        ID3D11ShaderReflectionConstantBuffer *cb = reflector->GetConstantBufferByIndex(i);

        D3D11_SHADER_BUFFER_DESC cbDesc{};
        (void)cb->GetDesc(&cbDesc);

        Shader::CBufferDesc outCB;
        outCB.name = cbDesc.Name;
        outCB.byteSize = cbDesc.Size;
        outCB.bindSlot = FindCbBindSlot(reflector.Get(), cbDesc.Name);

        // 反射 props
        for (UINT j = 0; j < cbDesc.Variables; ++j) {
            ID3D11ShaderReflectionVariable *var = cb->GetVariableByIndex(j);

            D3D11_SHADER_VARIABLE_DESC varDesc{};
            (void)var->GetDesc(&varDesc);

            ID3D11ShaderReflectionType *type = var->GetType();

            D3D11_SHADER_TYPE_DESC typeDesc{};
            (void)type->GetDesc(&typeDesc);

            Shader::PropertyDesc prop;
            prop.name = varDesc.Name;
            prop.offset = varDesc.StartOffset;
            prop.size = varDesc.Size;
            prop.rows = typeDesc.Rows;
            prop.columns = typeDesc.Columns;
            prop.varClass = typeDesc.Class;
            prop.varType = typeDesc.Type;
            // 反射DefaultValue
            if (varDesc.DefaultValue != nullptr) {
                const size_t size = varDesc.Size;
                if (size <= sizeof(prop.defaultValue)) {
                    std::memcpy(&prop.defaultValue, varDesc.DefaultValue, size);
                }
            }

            outCB.properties.push_back(std::move(prop));
        }

        result.push_back(std::move(outCB));
    }
    return result;
}

// 合并不同 Shader Stage 反射出的同名 CBuffer。
void MergeCBufferDescs_(std::vector<Shader::CBufferDesc> &dest,
                        std::vector<Shader::CBufferDesc> src) {
    for (Shader::CBufferDesc &s: src) {
        auto it = std::ranges::find_if(
            dest, [&s](const Shader::CBufferDesc &item) { return item.name == s.name; });

        // 未找到同名CBbuffer同时在VS和PS中被调用,直接添加
        if (it == dest.end()) {
            dest.push_back(std::move(s));
            continue;
        }

        // 同名 CBuffer 在 VS/PS 中必须拥有同一槽位和布局,发出警告,不合并
        if (it->bindSlot != s.bindSlot || it->byteSize != s.byteSize) {
            MessageBoxW(nullptr, L"同名 CBuffer 在不同 Shader Stage 的布局不一致",
                        L"Shader Reflection", MB_OK);
        }
    }
}
}

//MARK: 反射InputLayout
namespace {
constexpr D3D11_INPUT_ELEMENT_DESC inputElements_FULL[] = {
    {
        "POSITION", 0, DXGI_FORMAT_R32G32B32_FLOAT, 0,
        offsetof(Z::VertexData, pos), D3D11_INPUT_PER_VERTEX_DATA,
        0
    },
    {
        "NORMAL", 0, DXGI_FORMAT_R32G32B32_FLOAT, 0,
        offsetof(Z::VertexData, normal), D3D11_INPUT_PER_VERTEX_DATA,
        0
    },
    {
        "TANGENT", 0, DXGI_FORMAT_R32G32B32_FLOAT, 0,
        offsetof(Z::VertexData, tangent), D3D11_INPUT_PER_VERTEX_DATA,
        0
    },
    {
        "TEXCOORD", 0, DXGI_FORMAT_R32G32_FLOAT, 0,
        offsetof(Z::VertexData, uv0), D3D11_INPUT_PER_VERTEX_DATA,
        0
    },
    {
        "TEXCOORD", 1, DXGI_FORMAT_R32G32_FLOAT, 0,
        offsetof(Z::VertexData, uv1), D3D11_INPUT_PER_VERTEX_DATA,
        0
    },
    {
        "COLOR", 0, DXGI_FORMAT_R32G32B32A32_FLOAT, 0,
        offsetof(Z::VertexData, color), D3D11_INPUT_PER_VERTEX_DATA,
        0
    },
    {
        "BLENDINDICES", 0, DXGI_FORMAT_R32G32B32A32_SINT, 0,
        static_cast<UINT>(offsetof(Z::VertexData, ar_boneIndex)),
        D3D11_INPUT_PER_VERTEX_DATA,
        0
    },
    {
        "BLENDWEIGHT", 0, DXGI_FORMAT_R32G32B32A32_FLOAT, 0,
        static_cast<UINT>(offsetof(Z::VertexData, ar_boneWeights)),
        D3D11_INPUT_PER_VERTEX_DATA,
        0
    }
};

struct InputElementDesc {
    std::string semanticName;
    UINT semanticIndex = 0;
    DXGI_FORMAT format = DXGI_FORMAT_UNKNOWN;

    UINT inputSlot = 0;
    UINT alignedByteOffset = D3D11_APPEND_ALIGNED_ELEMENT;

    D3D11_INPUT_CLASSIFICATION inputSlotClass = D3D11_INPUT_PER_VERTEX_DATA;

    UINT instanceDataStepRate = 0;
};

[[maybe_unused]] void PrintInputLayoutElements(const std::vector<InputElementDesc> &ied) {
    // ReSharper disable once CppUseStructuredBinding
    for (auto &item: ied) {
        std::cout << "\033[0m============\033[36m" <<
                "\nsementicName:" << item.semanticName <<
                "\nsementicIndex:" << item.semanticIndex <<
                "\nformat:" << item.format <<
                "\nalignedByOffset:" << item.alignedByteOffset <<
                "\ninputSlot:" << item.inputSlot <<
                "\ninputSlotClass:" << item.inputSlotClass <<
                "\ninstanceDataStepRate:" << item.instanceDataStepRate << "\n";
    }
    std::cout << "\033[0m" << std::endl;
}

const D3D11_INPUT_ELEMENT_DESC *FindDescFromTemplates(
    const D3D11_INPUT_ELEMENT_DESC *templates, int c_temps,
    const char *sementicName, UINT sementicIdx = 0) {

    for (const D3D11_INPUT_ELEMENT_DESC &desc: std::span(templates, c_temps)) {
        if (std::strcmp(desc.SemanticName, sementicName) == 0 && desc.SemanticIndex == sementicIdx) {
            return &desc;
        }
    }
    return nullptr;
}

void ProcessInputElements(std::vector<InputElementDesc> &eles,
                          const D3D11_INPUT_ELEMENT_DESC *templates, int c_temps) {
    std::vector<InputElementDesc> filtered;
    filtered.reserve(eles.size());

    for (auto &item: eles) {
        const auto *desc = FindDescFromTemplates(templates, c_temps,
                                                 item.semanticName.c_str(), item.semanticIndex);

        if (!desc) { continue; }

        item.alignedByteOffset = desc->AlignedByteOffset;
        item.format = desc->Format;
        item.inputSlot = desc->InputSlot;
        item.inputSlotClass = desc->InputSlotClass;
        item.instanceDataStepRate = desc->InstanceDataStepRate;

        filtered.push_back(std::move(item));
    }

    eles = std::move(filtered);

    // PrintInputLayoutElements(eles);

}

DXGI_FORMAT GetFormatFromSignature(const D3D11_SIGNATURE_PARAMETER_DESC &desc) {
    UINT mask = desc.Mask;

    if (desc.ComponentType == D3D_REGISTER_COMPONENT_FLOAT32) {
        if (mask == 1) return DXGI_FORMAT_R32_FLOAT;
        if (mask == 3) return DXGI_FORMAT_R32G32_FLOAT;
        if (mask == 7) return DXGI_FORMAT_R32G32B32_FLOAT;
        if (mask == 15) return DXGI_FORMAT_R32G32B32A32_FLOAT;
    }

    if (desc.ComponentType == D3D_REGISTER_COMPONENT_UINT32) {
        if (mask == 1) return DXGI_FORMAT_R32_UINT;
        if (mask == 3) return DXGI_FORMAT_R32G32_UINT;
        if (mask == 7) return DXGI_FORMAT_R32G32B32_UINT;
        if (mask == 15) return DXGI_FORMAT_R32G32B32A32_UINT;
    }

    if (desc.ComponentType == D3D_REGISTER_COMPONENT_SINT32) {
        if (mask == 1) return DXGI_FORMAT_R32_SINT;
        if (mask == 3) return DXGI_FORMAT_R32G32_SINT;
        if (mask == 7) return DXGI_FORMAT_R32G32B32_SINT;
        if (mask == 15) return DXGI_FORMAT_R32G32B32A32_SINT;
    }

    return DXGI_FORMAT_UNKNOWN;
}

std::vector<InputElementDesc> ReflectInputLayout_(ID3DBlob *vsBlob) {
    std::vector<InputElementDesc> result;

    Microsoft::WRL::ComPtr<ID3D11ShaderReflection> reflector;

    HRESULT hr = D3DReflect(
        vsBlob->GetBufferPointer(),
        vsBlob->GetBufferSize(),
        IID_PPV_ARGS(reflector.GetAddressOf())
    );

    if (FAILED(hr)) {
        MessageBoxW(nullptr, L"Reflect VS Failed", L"ERROR", MB_OK);
        return result;
    }

    D3D11_SHADER_DESC shaderDesc{};
    hr = reflector->GetDesc(&shaderDesc);

    if (FAILED(hr)) {
        MessageBoxW(nullptr, L"Reflect VS GetDesc Failed", L"ERROR", MB_OK);
        return result;
    }

    for (UINT i = 0; i < shaderDesc.InputParameters; ++i) {
        D3D11_SIGNATURE_PARAMETER_DESC paramDesc{};
        hr = reflector->GetInputParameterDesc(i, &paramDesc);

        if (FAILED(hr)) {
            continue;
        }

        // 跳过系统语义，比如 SV_VertexID / SV_InstanceID
        if (paramDesc.SystemValueType != D3D_NAME_UNDEFINED) {
            continue;
        }

        InputElementDesc element;
        element.semanticName = paramDesc.SemanticName;
        element.semanticIndex = paramDesc.SemanticIndex;
        element.format = GetFormatFromSignature(paramDesc);

        element.inputSlot = 0;
        element.alignedByteOffset = D3D11_APPEND_ALIGNED_ELEMENT;
        element.inputSlotClass = D3D11_INPUT_PER_VERTEX_DATA;
        element.instanceDataStepRate = 0;

        if (element.format == DXGI_FORMAT_UNKNOWN) {
            continue;
        }

        result.push_back(element);
    }

    return result;
}

Microsoft::WRL::ComPtr<ID3D11InputLayout> CreateInputLayoutFromVS_(ID3D11Device *device,
                                                                   ID3DBlob *vsBlob) {
    // 先按照shader反射了inputlayout,写入 std::vector<InputElementDesc> reflectedElements
    std::vector<InputElementDesc> reflectedElements = ReflectInputLayout_(vsBlob);

    // 检查是否匹配 FULL 模式的模板，将不符合的剔除,按照描述表查找语义,再写偏移
    int c_tempeleFull = sizeof(inputElements_FULL) / sizeof(D3D11_INPUT_ELEMENT_DESC);
    ProcessInputElements(reflectedElements, inputElements_FULL, c_tempeleFull);

    // 重新拼写交给device的DESC
    std::vector<D3D11_INPUT_ELEMENT_DESC> dxElements;
    dxElements.reserve(reflectedElements.size());

    for (const auto &e: reflectedElements) {
        D3D11_INPUT_ELEMENT_DESC desc{};

        desc.SemanticName = e.semanticName.c_str();
        desc.SemanticIndex = e.semanticIndex;
        desc.Format = e.format;
        desc.InputSlot = e.inputSlot;
        desc.AlignedByteOffset = e.alignedByteOffset;
        desc.InputSlotClass = e.inputSlotClass;
        desc.InstanceDataStepRate = e.instanceDataStepRate;

        dxElements.push_back(desc);
    }

    // 创建InputLayout
    Microsoft::WRL::ComPtr<ID3D11InputLayout> inputLayout;
    HRESULT hr = device->CreateInputLayout(
        dxElements.data(),
        static_cast<UINT>(dxElements.size()),
        vsBlob->GetBufferPointer(),
        vsBlob->GetBufferSize(),
        inputLayout.GetAddressOf()
    );
    if (FAILED(hr)) {
        MessageBoxW(nullptr, L"Create InputLayout Failed", L"ERROR", MB_OK);
        return nullptr;
    }

    return inputLayout;
}

}

//MARK: 反射TexBuffer
namespace {
std::vector<Shader::TexBufDesc> ReflectShaderTextures(ID3DBlob *shaderBlob) {
    std::vector<Shader::TexBufDesc> textures;

    Microsoft::WRL::ComPtr<ID3D11ShaderReflection> reflector;

    HRESULT hr = D3DReflect(shaderBlob->GetBufferPointer(),
                            shaderBlob->GetBufferSize(),IID_PPV_ARGS(reflector.GetAddressOf()));

    if (FAILED(hr)) { return textures; }

    D3D11_SHADER_DESC shaderDesc = {};
    (void)reflector->GetDesc(&shaderDesc);

    for (UINT i = 0; i < shaderDesc.BoundResources; ++i) {
        D3D11_SHADER_INPUT_BIND_DESC bindDesc = {};
        (void)reflector->GetResourceBindingDesc(i, &bindDesc);

        if (bindDesc.Type == D3D_SIT_TEXTURE) {
            Shader::TexBufDesc info;
            info.name = bindDesc.Name;
            info.bindPoint = bindDesc.BindPoint;
            info.bindCount = bindDesc.BindCount;
            info.type = bindDesc.Type;
            info.dimension = bindDesc.Dimension;

            textures.push_back(info);
        }
    }

    return textures;
}
}



std::unordered_map<std::string, Shader *> Shader::map_Name_Shader{};
//= public
Shader *Shader::S_GetDefaultShader(ID3D11Device *device) {
    const std::filesystem::path &shaderPath = StatesRender::S_GetDefaultShaderPath();
    return Shader::S_CompileShader(shaderPath, Entry::VertexShader | Entry::PixelShader, device);
}

Shader *Shader::S_CompileShader(const std::filesystem::path &shaderPath, Entry entry,
                                ID3D11Device *device) {
    std::u8string u8fName = shaderPath.filename().u8string();
    std::string fileName = {reinterpret_cast<const char *>(u8fName.c_str()), u8fName.size()};
    auto it = map_Name_Shader.find(fileName);
    if (it != map_Name_Shader.end()) {
        return it->second;
    }

    using namespace Microsoft::WRL;
    HRESULT hr{};
    Shader *shader = new Shader{fileName.c_str(), shaderPath};

    // 编译VS
    shader->m_entries = entry;
    if (HasEntry_(entry, Entry::VertexShader)) {
        hr = CompileFromFile_(shaderPath, "VSMain", "vs_5_0",
                              shader->m_vertex_blob);
        if (FAILED(hr)) {
            delete shader;
            return nullptr;
        }
        hr = device->CreateVertexShader(shader->m_vertex_blob.Get()->GetBufferPointer(),
                                        shader->m_vertex_blob.Get()->GetBufferSize(),
                                        nullptr,
                                        shader->m_vertex_shader.GetAddressOf());
        if (FAILED(hr)) {
            MessageBoxW(nullptr, L"Create Vertex Shader failed", L"Shader::CreateShader",MB_OK);
            delete shader;
            return nullptr;
        }

        // 反射InputLayout
        shader->m_input_layout = CreateInputLayoutFromVS_(device, shader->m_vertex_blob.Get());

        // 反射 VS 的 CBuffer，但暂不缓存 CBPerMat 指针
        // 后续合并 PS 描述时 vector 可能重分配
        MergeCBufferDescs_(shader->m_cb_desc, ReflectCBuffers_(shader->m_vertex_blob.Get()));
    }


    // 编译PS
    if (HasEntry_(entry, Entry::PixelShader)) {
        hr = CompileFromFile_(shaderPath, "PSMain", "ps_5_0",
                              shader->m_pixel_blob);
        if (FAILED(hr)) {
            delete shader;
            return nullptr;
        }
        hr = device->CreatePixelShader(shader->m_pixel_blob.Get()->GetBufferPointer(),
                                       shader->m_pixel_blob.Get()->GetBufferSize(),
                                       nullptr,
                                       shader->m_pixel_shader.GetAddressOf());
        if (FAILED(hr)) {
            MessageBoxW(nullptr, L"Create Pixel Shader failed", L"Shader::CreateShader",MB_OK);
            delete shader;
            return nullptr;
        }
        // CBPerMat 往往只在 PS 使用，因此必须反射并合并 PS CBuffer。
        MergeCBufferDescs_(shader->m_cb_desc, ReflectCBuffers_(shader->m_pixel_blob.Get()));

        // 反射pixel shader的texture input
        shader->m_texBuf_desc = ReflectShaderTextures(shader->m_pixel_blob.Get());
    }
    // 所有 Stage 的反射结果合并完成后再取得稳定的 CBPerMat 的 Desc的指针
    auto perMatIt = std::ranges::find_if(
        shader->m_cb_desc, [](const CBufferDesc &item) { return item.name == "CBPerMat"; });
    shader->m_perMatCbDesc = perMatIt == shader->m_cb_desc.end() ? nullptr : &*perMatIt;

    map_Name_Shader.emplace(std::pair{fileName, shader});
    return shader;
}
Shader *Shader::S_GetShader(const std::string &fileName) {
    auto it = map_Name_Shader.find(fileName);
    if (it == map_Name_Shader.end()) {
        return nullptr;
    }
    return it->second;
}
bool Shader::HasEntry(Entry entry) const {
    return HasEntry_(this->m_entries, entry);
}


//= prvate

Shader::Shader(const char *name, const std::filesystem::path &path) :
m_name(std::string{name}), m_path(std::filesystem::path{path}) {}




















//
