#pragma once
#include "VoxelWorld.h"
#include "VoxelMesher.h"
#include "TopingSystem.h"
#include "DeferredGPUBuffer.h"
#include "VoxelRTManager.h"
#include "WickedEngine.h"

namespace voxel {

// ── CPU Profiling accumulator ────────────────────────────────────
struct ProfileAccum {
    double totalMs = 0.0;
    uint32_t count = 0;
    void add(float ms) { totalMs += ms; count++; }
    float avg() const { return count > 0 ? (float)(totalMs / count) : 0.0f; }
    void reset() { totalMs = 0.0; count = 0; }
};

// ── GPU-visible chunk info (must match HLSL GPUChunkInfo) ────────
struct GPUChunkInfo {
    XMFLOAT4 worldPos;         // xyz = chunk origin, w = debug flag
    uint32_t quadOffset;        // offset into mega quad buffer
    uint32_t quadCount;         // number of quads for this chunk
    uint32_t pad[2];            // align to 32 bytes
    uint32_t faceOffsets[6];    // per-face quad offset within this chunk's quads
    uint32_t faceCounts[6];     // per-face quad count
    uint32_t neighbors[6];     // chunk index of face neighbors (+X,-X,+Y,-Y,+Z,-Z), 0xFFFFFFFF = none
    uint32_t pad2[2];          // pad to 112 bytes (7 × float4)
};

// ── Voxel Renderer (GPU mesh pipeline) ──────────────────────────
class VoxelRenderer {
    friend class VoxelRenderPath;
public:
    VoxelRenderer();
    ~VoxelRenderer();

    void initialize(wi::graphics::GraphicsDevice* device);
    void shutdown();

    // Mesh dirty chunks and repack the mega-buffer
    void updateMeshes(VoxelWorld& world);

    // Render all visible chunks
    void render(
        wi::graphics::CommandList cmd,
        const wi::scene::CameraComponent& camera,
        const wi::graphics::Texture& depthBuffer,
        const wi::graphics::Texture& renderTarget,
        const wi::graphics::Texture& normalTarget
    ) const;

    // Load material textures from PNG files (RGB=albedo, A=heightmap)
    void loadTextures();

    // Stats
    uint32_t getTotalQuads() const { return totalQuads_; }
    uint32_t getVisibleChunks() const { return visibleChunks_; }
    uint32_t getDrawCalls() const { return drawCalls_; }
    uint32_t getChunkCount() const { return chunkCount_; }
    bool isInitialized() const { return initialized_; }

    bool debugFaceColors_ = false;
    bool debugBlend_ = false;
    float windTime_ = 0.0f;  // set by VoxelRenderPath::Update each frame
    float normalStrength_ = 0.7f; // normal map strength (0=off)
    int debugLighting_ = 0;       // 0=all, 1=no nmap, 2=flat, 3=albedo, 4=NdotL
    XMFLOAT4 sunDirection_ = { -0.7f, -0.4f, -0.3f, 0.0f }; // set by VoxelRenderPath::Update

private:
    void createPipeline();

    wi::graphics::GraphicsDevice* device_ = nullptr;

    // Shaders & Pipeline (voxels)
    wi::graphics::Shader vertexShader_;
    wi::graphics::Shader pixelShader_;
    wi::graphics::PipelineState pso_;
    // Shaders & Pipeline (topings, Phase 4)
    wi::graphics::Shader topingVS_;
    wi::graphics::Shader topingPS_;
    wi::graphics::PipelineState topingPso_;
    wi::graphics::GPUBuffer topingVertexBuffer_;   // StructuredBuffer<TopingVertex>, SRV t4
    DeferredGPUBuffer topingInstanceBuf_;            // StructuredBuffer<float3>, SRV t5
    static constexpr uint32_t MAX_TOPING_INSTANCES = 256 * 1024; // 256K instances max
    // Persistent staging buffers for toping upload (avoids per-frame allocations)
    struct TopingSortedInst { float wx, wy, wz; uint16_t type, variant; };
    struct TopingGPUInst { float x, y, z; };
    std::vector<TopingSortedInst> topingSorted_;
    std::vector<TopingGPUInst> topingGpuInsts_;
    mutable uint32_t topingDrawCalls_ = 0;

    // ── Toping draw groups (shared between render + BLAS CS) ─────
    struct TopingDrawGroup {
        uint16_t type, variant;
        uint32_t instanceOffset, instanceCount;
        uint32_t vertexTemplateOffset, vertexCount; // from TopingDef::variants[]
    };
    std::vector<TopingDrawGroup> topingDrawGroups_; // built in uploadTopingData, reused in renderTopings

    // ── Toping BLAS group staging (passed to VoxelRTManager) ──────
    struct TopingBLASGroupGPU {
        uint32_t globalVertexOffset;    // prefix sum of total vertices before this group
        uint32_t vertexTemplateOffset;  // offset into topingVertices (t4)
        uint32_t vertexCount;           // vertices per instance
        uint32_t instanceOffset;        // offset into topingInstances (t5)
        uint32_t instanceCount;         // instances in this group
    };
    std::vector<TopingBLASGroupGPU> topingBLASGroupsGPU_; // CPU staging for group table
    mutable uint32_t topingBLASTotalVertices_ = 0;

    // Shaders & Pipeline (smooth surfaces, Phase 5)
    wi::graphics::Shader smoothVS_;
    wi::graphics::Shader smoothPS_;
    wi::graphics::RasterizerState smoothRasterizer_;
    wi::graphics::PipelineState smoothPso_;
    mutable uint32_t smoothDrawCalls_ = 0;

    // Texture arrays for materials (512x512, 6 layers each)
    wi::graphics::Texture textureArray_;      // RGBA: RGB=albedo, A=heightmap (t1)
    wi::graphics::Texture normalArray_;       // RGB: tangent-space normal map (t7)
    wi::graphics::Sampler sampler_;

    // ── Mega-buffer architecture (Phase 2) ──────────────────────
    static constexpr uint32_t MEGA_BUFFER_CAPACITY = 2 * 1024 * 1024; // 2M quads max (16 MB)
    static constexpr uint32_t MAX_CHUNKS = 2048;

    wi::graphics::GPUBuffer chunkInfoBuffer_;   // StructuredBuffer<GPUChunkInfo>, SRV t2

    // CPU-side tracking
    struct ChunkSlot {
        ChunkPos pos;
        uint32_t quadOffset;    // offset into mega-buffer (in quads)
        uint32_t quadCount;
    };
    std::vector<ChunkSlot> chunkSlots_;
    std::vector<GPUChunkInfo> cpuChunkInfo_;
    uint32_t chunkCount_ = 0;
    bool megaBufferDirty_ = true;

    // Constants buffer (must match HLSL VoxelCB)
    struct VoxelConstants {
        XMFLOAT4X4 viewProjection;
        XMFLOAT4X4 inverseViewProjection; // for depth-to-world reconstruction (RT shadows)
        XMFLOAT4X4 prevViewProjection;    // previous frame's VP for temporal reprojection (RT AO)
        XMFLOAT4 cameraPosition;
        XMFLOAT4 sunDirection;
        XMFLOAT4 sunColor;
        float chunkSize;
        float textureTiling;
        float blendEnabled;
        float debugBlend;
        XMFLOAT4 frustumPlanes[6]; // ax+by+cz+d=0
        uint32_t chunkCount;
        uint32_t bleedMask;
        uint32_t resistBleedMask;
        float windTime;
        // Stylized lighting (Phase 7)
        XMFLOAT4 skyAmbient;
        XMFLOAT4 groundAmbient;
        XMFLOAT4 shadowTint;
        XMFLOAT4 fogColor;
        XMFLOAT4 fogParams;
        XMFLOAT4 rimColor;
        XMFLOAT4 rimParams;
        XMFLOAT4 toneMapParams;
    };
    wi::graphics::GPUBuffer constantBuffer_;

    // ── GPU Compute Mesher ──────────────────────────────────────────
    wi::graphics::Shader meshShader_;         // voxelMeshCS compute shader
    mutable wi::graphics::GPUBuffer voxelDataBuffer_; // chunk voxel data (StructuredBuffer<uint>)
    wi::graphics::GPUBuffer gpuQuadBuffer_;   // GPU mesh output (RWStructuredBuffer<uint2>)
    wi::graphics::GPUBuffer gpuQuadCounter_;  // atomic counter for GPU mesh output
    wi::graphics::GPUBuffer meshCounterReadback_; // READBACK buffer for quad counter
    bool gpuMesherAvailable_ = false;
    mutable uint32_t gpuMeshQuadCount_ = 0;   // Readback from previous frame (1-frame delay)
    mutable uint32_t voxelDataCapacity_ = 0;  // Current capacity of voxelDataBuffer_ (in uint32s)
    mutable std::vector<uint32_t> packedVoxelCache_; // cached packed voxel data for all chunks
    mutable bool voxelCacheDirty_ = true;     // true: packedVoxelCache_ needs repack from chunks
    mutable bool gpuMeshDirty_ = true;        // true: GPU needs upload + re-dispatch
    mutable bool chunkInfoDirty_ = true;      // true: chunkInfoBuffer needs re-upload

    // ── GPU Smooth Mesher (Phase 5.3) ─────────────────────────────
    wi::graphics::Shader smoothCentroidShader_;      // voxelSmoothCentroidCS (pass 1: centroid grid)
    wi::graphics::Shader smoothMeshShader_;          // voxelSmoothCS (pass 2: emit with smooth normals)
    wi::graphics::GPUBuffer centroidGridBuffer_;     // float4[34^3] per-chunk centroid grid (reused)
    wi::graphics::GPUBuffer gpuSmoothVertexBuffer_;  // RWStructuredBuffer<GPUSmoothVertex>, UAV+SRV
    wi::graphics::GPUBuffer gpuSmoothCounter_;       // atomic counter for smooth vertices
    wi::graphics::GPUBuffer smoothCounterReadback_;   // READBACK buffer for vertex counter
    static constexpr uint32_t CENTROID_GRID_SIZE = 34 * 34 * 34; // 39304 entries per chunk
    static constexpr uint32_t MAX_GPU_SMOOTH_VERTICES = 2 * 1024 * 1024; // 2M vertices max
    mutable uint32_t gpuSmoothVertexCount_ = 0;       // readback from previous frame
    mutable bool gpuSmoothMeshDirty_ = true;

    // ── Ray Tracing (Phase 6) ────────────────────────────────────────
    static constexpr uint32_t MAX_BLAS_VERTICES = MEGA_BUFFER_CAPACITY * 6; // 6 verts per quad
    mutable VoxelRTManager rt_;

    void dispatchGpuMesh(wi::graphics::CommandList cmd, const VoxelWorld& world,
        ProfileAccum* profPack = nullptr, ProfileAccum* profUpload = nullptr,
        ProfileAccum* profDispatch = nullptr) const;
    void dispatchGpuSmoothMesh(wi::graphics::CommandList cmd, const VoxelWorld& world) const;
    void rebuildChunkInfoOnly(VoxelWorld& world);

    // ── GPU Timestamp Queries (comprehensive GPU profiling) ────────
    wi::graphics::GPUQueryHeap timestampHeap_;
    wi::graphics::GPUBuffer timestampReadback_;
    // Timestamp slots: pairs of (BEGIN, END) for each GPU phase
    static constexpr uint32_t TS_GPU_MESH_BEGIN = 0;
    static constexpr uint32_t TS_GPU_MESH_END = 1;
    static constexpr uint32_t TS_GPU_SMOOTH_BEGIN = 2;
    static constexpr uint32_t TS_GPU_SMOOTH_END = 3;
    static constexpr uint32_t TS_BLAS_EXTRACT_BEGIN = 4;
    static constexpr uint32_t TS_BLAS_EXTRACT_END = 5;
    static constexpr uint32_t TS_BLAS_BUILD_BEGIN = 6;
    static constexpr uint32_t TS_BLAS_BUILD_END = 7;
    static constexpr uint32_t TS_DRAW_BEGIN = 8;
    static constexpr uint32_t TS_DRAW_END = 9;
    static constexpr uint32_t TS_RT_SHADOWS_BEGIN = 10;
    static constexpr uint32_t TS_RT_SHADOWS_END = 11;
    static constexpr uint32_t TS_COUNT = 12;
    mutable float gpuMeshTimeMs_ = 0.0f;
    mutable float gpuSmoothMeshTimeMs_ = 0.0f;
    mutable float gpuBLASExtractTimeMs_ = 0.0f;
    mutable float gpuBLASBuildTimeMs_ = 0.0f;
    mutable float gpuDrawTimeMs_ = 0.0f;
    mutable float gpuRTShadowsTimeMs_ = 0.0f;

    // Stats (mutable: updated during const Render() call)
    mutable uint32_t totalQuads_ = 0;
    mutable uint32_t visibleChunks_ = 0;
    mutable uint32_t drawCalls_ = 0;

    bool initialized_ = false;

public:
    float getGpuDrawTimeMs() const { return gpuDrawTimeMs_; }
    float getGpuMeshTimeMs() const { return gpuMeshTimeMs_; }
    float getGpuSmoothMeshTimeMs() const { return gpuSmoothMeshTimeMs_; }
    float getGpuBLASExtractTimeMs() const { return gpuBLASExtractTimeMs_; }
    float getGpuBLASBuildTimeMs() const { return gpuBLASBuildTimeMs_; }
    float getGpuRTShadowsTimeMs() const { return gpuRTShadowsTimeMs_; }
    bool isGpuMeshEnabled() const { return gpuMesherAvailable_; }
    uint32_t getGpuMeshQuadCount() const { return gpuMeshQuadCount_; }
    VoxelRTManager& rt() const { return rt_; }

    // Phase 4: Toping rendering
    void uploadTopingData(const TopingSystem& topingSystem);
    void renderTopings(
        wi::graphics::CommandList cmd,
        const TopingSystem& topingSystem,
        const wi::graphics::Texture& depthBuffer,
        const wi::graphics::Texture& renderTarget,
        const wi::graphics::Texture& normalTarget
    ) const;
    uint32_t getTopingDrawCalls() const { return topingDrawCalls_; }

    // Phase 5: Smooth surface rendering (GPU compute only)
    void renderSmooth(
        wi::graphics::CommandList cmd,
        const wi::graphics::Texture& depthBuffer,
        const wi::graphics::Texture& renderTarget,
        const wi::graphics::Texture& normalTarget
    ) const;
    uint32_t getSmoothVertexCount() const { return gpuSmoothVertexCount_; }
    uint32_t getSmoothDrawCalls() const { return smoothDrawCalls_; }

    // Phase 6: Ray Tracing (delegated to VoxelRTManager)
    bool isRTAvailable() const { return rt_.isAvailable(); }
    bool isRTReady() const { return rt_.isReady(); }
    bool isRTShadowsEnabled() const { return rt_.isShadowsEnabled(); }
    uint32_t getRTBlockyTriCount() const { return rt_.getBlockyTriCount(); }
    uint32_t getRTSmoothTriCount() const { return rt_.getSmoothTriCount(); }
    uint32_t getRTTopingTriCount() const { return rt_.getTopingTriCount(); }
    const wi::graphics::RaytracingAccelerationStructure& getTLAS() const { return rt_.getTLAS(); }
};

// ── Camera Controller ────────────────────────────────────────────
struct CameraController {
    float speed = 50.0f;
    float sensitivity = 0.003f;
    XMFLOAT3 pos = { 256.0f, 100.0f, 256.0f };
    float pitch = -0.3f;
    float yaw = 0.0f;
    bool mouseCaptured = false;

    void set(float x, float y, float z, float p, float yw) {
        pos = { x, y, z }; pitch = p; yaw = yw;
    }
    void handleInput(float dt, wi::scene::CameraComponent* camera);
};

// ── Animation State ─────────────────────────────────────────────
struct AnimationState {
    float windTime = 0.0f;       // continuous, always running
    bool terrainAnimated = false; // toggled with F3
    bool sunOrbit = false;       // toggled with F7: sun orbits in ~10s cycle
    bool showCrosshair = true;   // toggled with F8: crosshair + face debug info
    // F9 debug cycle: 0=all ON, 1=normals OFF, 2=flat lighting, 3=albedo only, 4=NdotL only, 5=normal viz
    int debugLighting = 0;
    static constexpr int DEBUG_LIGHTING_MODES = 6;
    float time = 0.0f;           // current animation time offset
    float accum = 0.0f;          // accumulator for 30 Hz timer
    static constexpr float INTERVAL = 1.0f / 30.0f; // ~33.3ms = 30 Hz

    // Returns true when an animation tick should fire (call every frame).
    bool tick(float dt) {
        windTime += dt;
        if (!terrainAnimated) return false;
        accum += dt;
        if (accum < INTERVAL) return false;
        accum -= INTERVAL;
        time += INTERVAL;
        return true;
    }
};

// ── CPU Profiling (averages every INTERVAL seconds) ─────────────
struct VoxelProfiler {
    static constexpr float INTERVAL = 5.0f;

    // Update() phase
    ProfileAccum regenerate;      // regenerateAnimated
    ProfileAccum updateMeshes;    // updateMeshes (rebuildChunkInfoOnly)
    ProfileAccum topingCollect;   // topingSystem.collectInstances
    ProfileAccum topingUpload;    // uploadTopingData
    ProfileAccum smoothMesh;      // (legacy, unused — GPU smooth only)
    ProfileAccum smoothUpload;    // (legacy, unused — GPU smooth only)
    ProfileAccum frame;           // full frame (Update only - legacy)

    // Render() phase
    ProfileAccum voxelPack;       // voxel data packing in dispatchGpuMesh
    ProfileAccum gpuUpload;       // GPU upload in dispatchGpuMesh
    ProfileAccum gpuDispatch;     // compute dispatches in dispatchGpuMesh
    ProfileAccum gpuMeshDispatch; // GPU mesh compute dispatch (in Render)
    ProfileAccum gpuSmoothDispatch; // GPU smooth mesh dispatch (in Render)
    ProfileAccum blasExtract;     // BLAS position extraction compute
    ProfileAccum blasBuild;       // BLAS/TLAS build
    ProfileAccum deferredUpload;  // deferred GPU buffer uploads
    ProfileAccum render;          // render() draw calls
    ProfileAccum rtShadows;       // RT shadows + AO dispatch

    // Totals
    ProfileAccum fullFrame;       // true full frame (Update + Render + Compose)
    ProfileAccum gpuWait;         // GPU sync: time between Compose end and next Update start
    ProfileAccum wickedRender;    // RenderPath3D::Render() (Wicked internal)
    ProfileAccum trueFrame;       // wall-clock frame-to-frame time

    // Timing helpers
    std::chrono::high_resolution_clock::time_point frameStart;
    std::chrono::high_resolution_clock::time_point lastComposeEnd;
    bool lastComposeEndValid = false;
    float timer = 0.0f;

    void log(const VoxelRenderer& renderer) const;
    void resetAll();
};

// ── Custom RenderPath that integrates voxel rendering ───────────
class VoxelRenderPath : public wi::RenderPath3D {
public:
    VoxelWorld world;
    VoxelRenderer renderer;
    TopingSystem topingSystem;

    bool debugMode = false;
    bool debugSmooth = false;
    bool screenshotMode = false;  // CLI "screenshot": auto-position camera, capture, quit
    void setCamera(float x, float y, float z, float pitch, float yaw) {
        camera_.set(x, y, z, pitch, yaw);
    }
    void resetAOHistory();  // invalidate temporal AO after camera jump

    CameraController camera_;
    AnimationState anim_;
    mutable VoxelProfiler prof_;

    const wi::graphics::Texture& getVoxelRT() const { return voxelRT_; }

    void Start() override;
    void Update(float dt) override;
    void Render() const override;
    void Compose(wi::graphics::CommandList cmd) const override;

private:
    void createRenderTargets();
    mutable bool worldGenerated_ = false;
    mutable int frameCount_ = 0;
    mutable float lastDt_ = 0.016f;
    mutable float smoothFps_ = 60.0f;

    wi::graphics::Texture voxelRT_;
    wi::graphics::Texture voxelNormalRT_;  // Phase 6: world-space normals for RT shadows/AO
    wi::graphics::Texture voxelDepth_;
    mutable bool rtCreated_ = false;

    mutable uint32_t rtBuildSkipCounter_ = 0;  // stagger BLAS builds during animation
    mutable bool rtWasEnabled_ = false;       // saved RT state before animation

    // Cached crosshair raycast result (updated each frame in Compose)
    struct CrosshairHit {
        bool valid = false;
        int x = 0, y = 0, z = 0;
        int face = -1;       // 0=+X,1=-X,2=+Y,3=-Y,4=+Z,5=-Z
        uint8_t matID = 0;
        bool smooth = false;
    };
    mutable CrosshairHit crosshairHit_;

    // Build a full debug log string (used by HUD overlay and screenshot .log)
    std::string buildDebugLog() const;
};

} // namespace voxel
