#include "ibl_baker.h"

#include <algorithm>
#include <cstring>
#include <d3dcompiler.h>
#include <string>
#include <utility>

#pragma comment(lib, "d3dcompiler.lib")

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

// [新增] IBL 烘焙 Shader 使用多个入口，因此不走当前只支持 VSMain/PSMain 的 Shader 包装类。
bool CompileBakeShader(const std::filesystem::path &path,
                       const char *entryPoint,
                       const char *target,
                       Microsoft::WRL::ComPtr<ID3DBlob> &outBlob) {
    outBlob.Reset();
    Microsoft::WRL::ComPtr<ID3DBlob> errorBlob;

    const HRESULT hr = D3DCompileFromFile(
        path.c_str(), nullptr, D3D_COMPILE_STANDARD_FILE_INCLUDE,
        entryPoint, target, GetShaderCompileFlags(), 0,
        outBlob.GetAddressOf(), errorBlob.GetAddressOf());

    if (SUCCEEDED(hr)) return true;

    if (errorBlob) {
        const char *message = static_cast<const char *>(errorBlob->GetBufferPointer());
        OutputDebugStringA(message);
    }

    std::wstring message = L"编译 IBL 烘焙 Shader 失败。\n文件：";
    message += path.wstring();
    message += L"\n入口：";

    const int required = MultiByteToWideChar(CP_UTF8, 0, entryPoint, -1, nullptr, 0);
    if (required > 0) {
        std::wstring wideEntry(static_cast<std::size_t>(required), L'\0');
        MultiByteToWideChar(CP_UTF8, 0, entryPoint, -1, wideEntry.data(), required);
        message += wideEntry.c_str();
    }

    MessageBoxW(nullptr, message.c_str(), L"IBLBaker", MB_OK);
    return false;
}
}

bool IBLBaker::Initialize(ID3D11Device *device,
                          const std::filesystem::path &bakeShaderPath) {
    if (!device || bakeShaderPath.empty()) return false;
    if (m_device.Get() && m_fullscreenVS.Get() && m_irradiancePS.Get() &&
        m_prefilterPS.Get() && m_brdfLutPS.Get()) {
        return true;
    }

    Microsoft::WRL::ComPtr<ID3DBlob> fullscreenVsBlob;
    Microsoft::WRL::ComPtr<ID3DBlob> irradiancePsBlob;
    Microsoft::WRL::ComPtr<ID3DBlob> prefilterPsBlob;
    Microsoft::WRL::ComPtr<ID3DBlob> brdfLutPsBlob;

    if (!CompileBakeShader(bakeShaderPath, "VSFullscreen", "vs_5_0", fullscreenVsBlob) ||
        !CompileBakeShader(bakeShaderPath, "PSIrradiance", "ps_5_0", irradiancePsBlob) ||
        !CompileBakeShader(bakeShaderPath, "PSPrefilter", "ps_5_0", prefilterPsBlob) ||
        !CompileBakeShader(bakeShaderPath, "PSBrdfLut", "ps_5_0", brdfLutPsBlob)) {
        return false;
    }

    if (FAILED(device->CreateVertexShader(
            fullscreenVsBlob->GetBufferPointer(), fullscreenVsBlob->GetBufferSize(),
            nullptr, m_fullscreenVS.ReleaseAndGetAddressOf()))) {
        return false;
    }
    if (FAILED(device->CreatePixelShader(
            irradiancePsBlob->GetBufferPointer(), irradiancePsBlob->GetBufferSize(),
            nullptr, m_irradiancePS.ReleaseAndGetAddressOf()))) {
        return false;
    }
    if (FAILED(device->CreatePixelShader(
            prefilterPsBlob->GetBufferPointer(), prefilterPsBlob->GetBufferSize(),
            nullptr, m_prefilterPS.ReleaseAndGetAddressOf()))) {
        return false;
    }
    if (FAILED(device->CreatePixelShader(
            brdfLutPsBlob->GetBufferPointer(), brdfLutPsBlob->GetBufferSize(),
            nullptr, m_brdfLutPS.ReleaseAndGetAddressOf()))) {
        return false;
    }

    D3D11_BUFFER_DESC constantBufferDesc{};
    constantBufferDesc.ByteWidth = static_cast<UINT>(sizeof(BakeConstantBufferData));
    constantBufferDesc.Usage = D3D11_USAGE_DYNAMIC;
    constantBufferDesc.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
    constantBufferDesc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
    if (FAILED(device->CreateBuffer(
            &constantBufferDesc, nullptr,
            m_constantBuffer.ReleaseAndGetAddressOf()))) {
        return false;
    }

    D3D11_SAMPLER_DESC samplerDesc{};
    samplerDesc.Filter = D3D11_FILTER_MIN_MAG_MIP_LINEAR;
    samplerDesc.AddressU = D3D11_TEXTURE_ADDRESS_CLAMP;
    samplerDesc.AddressV = D3D11_TEXTURE_ADDRESS_CLAMP;
    samplerDesc.AddressW = D3D11_TEXTURE_ADDRESS_CLAMP;
    samplerDesc.ComparisonFunc = D3D11_COMPARISON_NEVER;
    samplerDesc.MinLOD = 0.0f;
    samplerDesc.MaxLOD = D3D11_FLOAT32_MAX;
    if (FAILED(device->CreateSamplerState(
            &samplerDesc, m_sourceSampler.ReleaseAndGetAddressOf()))) {
        return false;
    }

    // [新增] 全屏三角形不依赖绕序，关闭剔除避免驱动默认状态影响离线烘焙。
    D3D11_RASTERIZER_DESC rasterizerDesc{};
    rasterizerDesc.FillMode = D3D11_FILL_SOLID;
    rasterizerDesc.CullMode = D3D11_CULL_NONE;
    rasterizerDesc.DepthClipEnable = TRUE;
    if (FAILED(device->CreateRasterizerState(
            &rasterizerDesc, m_rasterizerState.ReleaseAndGetAddressOf()))) {
        return false;
    }

    m_device = device;
    return true;
}

bool IBLBaker::Bake(ID3D11DeviceContext *context,
                    const TextureCube &sourceEnvironment,
                    Result &outResult) {
    return Bake(context, sourceEnvironment, outResult, Settings{});
}

bool IBLBaker::Bake(ID3D11DeviceContext *context,
                    const TextureCube &sourceEnvironment,
                    Result &outResult,
                    const Settings &settings) {
    if (!context || !m_device.Get() || !sourceEnvironment.Resource() ||
        !sourceEnvironment.SRV() ||
        settings.irradianceSize == 0 || settings.prefilteredSize == 0 ||
        settings.brdfLutSize == 0 ||
        settings.irradianceSampleCount == 0 ||
        settings.prefilterSampleCount == 0 ||
        settings.brdfSampleCount == 0) {
        return false;
    }

    // [新增] 烘焙可以在初始化阶段或运行时换环境时调用；先清空旧管线绑定，
    // 避免源 Cube 仍被其他槽位绑定导致 D3D11 资源冲突。
    context->ClearState();

    const UINT prefilteredMipCount = CalculateMipCount(settings.prefilteredSize);

    Result result{};
    result.irradiance = TextureCube::CreateRenderTarget(
        m_device.Get(), "IBL_Irradiance", settings.irradianceSize, 1,
        DXGI_FORMAT_R16G16B16A16_FLOAT);
    result.prefiltered = TextureCube::CreateRenderTarget(
        m_device.Get(), "IBL_Prefiltered", settings.prefilteredSize,
        prefilteredMipCount, DXGI_FORMAT_R16G16B16A16_FLOAT);

    Texture2D::RenderTargetCreateDesc brdfDesc{};
    brdfDesc.name = "IBL_BRDF_LUT";
    brdfDesc.width = settings.brdfLutSize;
    brdfDesc.height = settings.brdfLutSize;
    brdfDesc.format = DXGI_FORMAT_R16G16_FLOAT;
    result.brdfLut = Texture2D::CreateRenderTarget(m_device.Get(), brdfDesc);

    if (!result.IsValid()) {
        context->ClearState();
        return false;
    }

    if (!RenderCube(context, sourceEnvironment, *result.irradiance,
                    m_irradiancePS.Get(), settings.irradianceSampleCount, false) ||
        !RenderCube(context, sourceEnvironment, *result.prefiltered,
                    m_prefilterPS.Get(), settings.prefilterSampleCount, true) ||
        !RenderBrdfLut(context, *result.brdfLut,
                       settings.brdfLutSize, settings.brdfSampleCount)) {
        context->ClearState();
        return false;
    }

    // [新增] 烘焙结束后不把临时 b0/t0/RTV 等状态泄漏到正式渲染管线。
    context->ClearState();
    outResult = std::move(result);
    return true;
}

bool IBLBaker::UpdateConstantBuffer(ID3D11DeviceContext *context,
                                    const float roughness,
                                    const float sourceSize,
                                    const float sourceMaxMip,
                                    const UINT faceIndex,
                                    const UINT sampleCount) const {
    if (!context || !m_constantBuffer.Get()) return false;

    const BakeConstantBufferData data{
        DirectX::XMFLOAT4{roughness, sourceSize, sourceMaxMip, 0.0f},
        DirectX::XMUINT4{faceIndex, sampleCount, 0u, 0u}
    };

    D3D11_MAPPED_SUBRESOURCE mapped{};
    if (FAILED(context->Map(
            m_constantBuffer.Get(), 0, D3D11_MAP_WRITE_DISCARD, 0, &mapped))) {
        return false;
    }

    std::memcpy(mapped.pData, &data, sizeof(data));
    context->Unmap(m_constantBuffer.Get(), 0);
    return true;
}

bool IBLBaker::RenderCube(ID3D11DeviceContext *context,
                          const TextureCube &sourceEnvironment,
                          TextureCube &destination,
                          ID3D11PixelShader *pixelShader,
                          const UINT sampleCount,
                          const bool roughnessFromMip) const {
    if (!context || !pixelShader || !destination.Resource() ||
        !sourceEnvironment.Resource() || !sourceEnvironment.SRV()) {
        return false;
    }

    D3D11_TEXTURE2D_DESC sourceDesc{};
    sourceEnvironment.Resource()->GetDesc(&sourceDesc);
    if (sourceDesc.Width == 0 || sourceDesc.MipLevels == 0) return false;

    context->IASetInputLayout(nullptr);
    context->IASetPrimitiveTopology(D3D11_PRIMITIVE_TOPOLOGY_TRIANGLELIST);
    context->VSSetShader(m_fullscreenVS.Get(), nullptr, 0);
    context->PSSetShader(pixelShader, nullptr, 0);
    context->RSSetState(m_rasterizerState.Get());
    context->OMSetDepthStencilState(nullptr, 0);

    ID3D11Buffer *constantBuffer = m_constantBuffer.Get();
    context->PSSetConstantBuffers(0, 1, &constantBuffer);

    ID3D11ShaderResourceView *sourceSrv = sourceEnvironment.SRV();
    context->PSSetShaderResources(0, 1, &sourceSrv);

    ID3D11SamplerState *sourceSampler = m_sourceSampler.Get();
    context->PSSetSamplers(0, 1, &sourceSampler);

    const UINT mipCount = destination.MipLevels();
    const float sourceMaxMip = static_cast<float>(sourceDesc.MipLevels - 1u);

    for (UINT mip = 0; mip < mipCount; ++mip) {
        const UINT mipSize = std::max(1u, destination.Size() >> mip);
        const float roughness = roughnessFromMip && mipCount > 1
            ? static_cast<float>(mip) / static_cast<float>(mipCount - 1u)
            : 0.0f;

        D3D11_VIEWPORT viewport{};
        viewport.Width = static_cast<float>(mipSize);
        viewport.Height = static_cast<float>(mipSize);
        viewport.MinDepth = 0.0f;
        viewport.MaxDepth = 1.0f;
        context->RSSetViewports(1, &viewport);

        for (UINT face = 0; face < 6; ++face) {
            D3D11_RENDER_TARGET_VIEW_DESC rtvDesc{};
            rtvDesc.Format = destination.Format();
            rtvDesc.ViewDimension = D3D11_RTV_DIMENSION_TEXTURE2DARRAY;
            rtvDesc.Texture2DArray.MipSlice = mip;
            rtvDesc.Texture2DArray.FirstArraySlice = face;
            rtvDesc.Texture2DArray.ArraySize = 1;

            Microsoft::WRL::ComPtr<ID3D11RenderTargetView> rtv;
            if (FAILED(m_device->CreateRenderTargetView(
                    destination.Resource(), &rtvDesc, rtv.GetAddressOf()))) {
                return false;
            }

            ID3D11RenderTargetView *rawRtv = rtv.Get();
            context->OMSetRenderTargets(1, &rawRtv, nullptr);

            if (!UpdateConstantBuffer(
                    context, roughness, static_cast<float>(sourceDesc.Width),
                    sourceMaxMip, face, sampleCount)) {
                return false;
            }

            context->Draw(3, 0);
        }
    }

    context->OMSetRenderTargets(0, nullptr, nullptr);
    ID3D11ShaderResourceView *nullSrv = nullptr;
    context->PSSetShaderResources(0, 1, &nullSrv);
    return true;
}

bool IBLBaker::RenderBrdfLut(ID3D11DeviceContext *context,
                             Texture2D &destination,
                             const UINT size,
                             const UINT sampleCount) const {
    if (!context || !destination.Resource() || size == 0) return false;

    Microsoft::WRL::ComPtr<ID3D11RenderTargetView> rtv;
    if (FAILED(m_device->CreateRenderTargetView(
            destination.Resource(), nullptr, rtv.GetAddressOf()))) {
        return false;
    }

    D3D11_VIEWPORT viewport{};
    viewport.Width = static_cast<float>(size);
    viewport.Height = static_cast<float>(size);
    viewport.MinDepth = 0.0f;
    viewport.MaxDepth = 1.0f;

    context->IASetInputLayout(nullptr);
    context->IASetPrimitiveTopology(D3D11_PRIMITIVE_TOPOLOGY_TRIANGLELIST);
    context->VSSetShader(m_fullscreenVS.Get(), nullptr, 0);
    context->PSSetShader(m_brdfLutPS.Get(), nullptr, 0);
    context->RSSetState(m_rasterizerState.Get());
    context->RSSetViewports(1, &viewport);
    context->OMSetDepthStencilState(nullptr, 0);

    ID3D11RenderTargetView *rawRtv = rtv.Get();
    context->OMSetRenderTargets(1, &rawRtv, nullptr);

    if (!UpdateConstantBuffer(context, 0.0f, 0.0f, 0.0f, 0u, sampleCount)) {
        return false;
    }

    ID3D11Buffer *constantBuffer = m_constantBuffer.Get();
    context->PSSetConstantBuffers(0, 1, &constantBuffer);
    context->Draw(3, 0);
    context->OMSetRenderTargets(0, nullptr, nullptr);
    return true;
}

UINT IBLBaker::CalculateMipCount(UINT size) noexcept {
    UINT mipCount = 1;
    while (size > 1) {
        size >>= 1u;
        ++mipCount;
    }
    return mipCount;
}
