#pragma once

#include <glad/gl.h>

#include <pxr/usd/usd/stage.h>
#include <pxr/usd/usd/timeCode.h>
#include <pxr/imaging/glf/drawTarget.h>
#include <pxr/usdImaging/usdImagingGL/engine.h>
#include <pxr/imaging/glf/simpleLight.h>
#include <pxr/imaging/glf/simpleMaterial.h>
#include <pxr/base/gf/vec3f.h>
#include <pxr/base/gf/vec4f.h>
#include <pxr/base/gf/vec4d.h>
#include <pxr/base/gf/matrix4d.h>
#include <pxr/base/gf/matrix4f.h>
#include <pxr/base/gf/range3d.h>
#include <pxr/base/gf/camera.h>
#include <pxr/base/gf/bbox3d.h>
#include <pxr/base/gf/frustum.h>
#include <pxr/usd/sdf/path.h>

#include <vector>
#include <string>
#include <memory>

namespace UsdLayerManager {

/// Bounding box display mode for selected prims.
enum class BBoxMode {
    None,           ///< No bounding boxes drawn
    PerObject,      ///< One box per selected prim
    AllSelection    ///< One combined box for the entire selection
};

/// Surface shading mode — maps to UsdImagingGLDrawMode + enableLighting.
enum class ShadingMode {
    SmoothShaded,        ///< Smooth shading with materials (default)
    FlatShaded,          ///< Flat shading with materials
    WireframeOnSurface,  ///< Wireframe overlay on solid surface
    Wireframe,           ///< Wireframe only, no fill
    Unlit,               ///< Smooth geometry, lighting disabled
};

/// Viewport color correction mode — maps to UsdImagingGLRenderParams::colorCorrectionMode.
enum class ColorCorrectionMode {
    Disabled,    ///< No correction (raw linear output)
    sRGB,        ///< Linear → sRGB gamma (default)
    OpenColorIO, ///< Full OCIO pipeline via bundled ACES 1.2 config
};

class UsdSceneRenderer {
public:
    UsdSceneRenderer();
    ~UsdSceneRenderer();

    void SetStage(pxr::UsdStageRefPtr stage);

    /// Render the current stage into the internal draw target.
    void Render(int width, int height);

    /// Compute the world-space bounding box of the whole stage.
    pxr::GfRange3d ComputeStageBounds();

    // -----------------------------------------------------------------------
    // Camera state
    // -----------------------------------------------------------------------

    /// Set free-camera view/proj matrices explicitly.
    void SetCameraState(const pxr::GfMatrix4d& viewMatrix,
                        const pxr::GfMatrix4d& projMatrix);

    /// Set free-camera state from a GfCamera (also extracts clip planes).
    void SetCameraStateFromGfCamera(const pxr::GfCamera& gfCamera);

    /// Set USD prim camera path; renderer resolves it internally.
    void SetCameraPath(const pxr::SdfPath& cameraPath);

    // -----------------------------------------------------------------------
    // Selection
    // -----------------------------------------------------------------------
    void ClearSelected();
    void AddSelected(const pxr::SdfPath& path, int instanceIndex = -1);
    /// Replace the entire selection with a set of paths (for multi-select).
    void SetSelectedPaths(const pxr::SdfPathVector& paths);

    // -----------------------------------------------------------------------
    // Picking
    // -----------------------------------------------------------------------

    /// Single-prim pick: returns true if a hit was found.
    bool PickObject(
        int mouseX, int mouseY,
        int viewWidth, int viewHeight,
        pxr::GfVec3d* outHitPoint,
        pxr::SdfPath* outHitPrimPath);

    /// Rect pick: finds all unique prims whose geometry overlaps the screen rect
    /// defined by (x0,y0)-(x1,y1) in viewport pixel coordinates.
    /// Returns true if at least one prim was hit.
    bool PickObjectsInRect(
        int x0, int y0, int x1, int y1,
        int viewWidth, int viewHeight,
        pxr::SdfPathVector* outHitPaths);

    // -----------------------------------------------------------------------
    // Overlay drawing (called while draw-target FBO is still bound)
    // -----------------------------------------------------------------------

    /// Draw XYZ axis gizmo matching stageView.DrawAxis().
    void DrawAxis(const pxr::GfMatrix4d& viewProjMatrix, double cameraDist);

    /// Draw bounding boxes for the given selected prim paths according to m_bboxMode.
    void DrawBoundingBoxes(
        const pxr::SdfPathVector& selectedPaths,
        const pxr::GfMatrix4d&   viewProjMatrix);

    /// Draw frustum wireframes for all UsdGeomCamera prims in the stage.
    /// selectedPaths — current primary selection (selected camera gets accent colour).
    /// activeCameraPath — camera currently driving the viewport (gets cyan).
    /// viewportCameraDist — used to scale wireframe size (from ViewportCamera::GetDist()).
    void DrawCameraWireframes(
        pxr::UsdStageRefPtr        stage,
        const pxr::SdfPathVector&  selectedPaths,
        const pxr::SdfPath&        activeCameraPath,
        const pxr::GfMatrix4d&     viewProjMatrix,
        double                     viewportCameraDist);

    /// Draw wireframe gizmos for all UsdLux light prims in the stage.
    /// The shape reflects each light's type (sphere/rect/disk/distant/dome/
    /// cylinder); selected lights get an accent colour.  viewportCameraDist
    /// scales direction lines and point/distant/dome markers.
    void DrawLightWireframes(
        pxr::UsdStageRefPtr        stage,
        const pxr::SdfPathVector&  selectedPaths,
        const pxr::GfMatrix4d&     viewProjMatrix,
        double                     viewportCameraDist);

    /// Test a screen-space mouse position against all camera wireframe segments.
    /// Returns true and sets *outCameraPath if a camera wireframe is within
    /// kCameraPickRadius (10 px) of mousePos.  Prioritises the closest hit.
    /// mouseX/Y and imagePosX/Y are absolute screen coordinates.
    bool PickCameraAtPoint(
        pxr::UsdStageRefPtr    stage,
        float mouseX, float mouseY,
        const pxr::GfMatrix4d& viewProjMatrix,
        float imagePosX, float imagePosY,
        int viewW, int viewH,
        double                 viewportCameraDist,
        pxr::SdfPath*          outCameraPath);

    // -----------------------------------------------------------------------
    // Draw-target FBO helpers (for external overlays, e.g. TransformManipulator)
    // -----------------------------------------------------------------------
    /// Bind the internal offscreen FBO and set the GL viewport.
    /// Must be paired with UnbindDrawTarget() after drawing.
    void BindDrawTarget();
    /// Unbind the internal offscreen FBO.
    void UnbindDrawTarget();

    // -----------------------------------------------------------------------
    // Render delegate
    // -----------------------------------------------------------------------

    /// Return all available renderer plugin IDs (e.g. HdStormRendererPlugin).
    static std::vector<pxr::TfToken> GetRendererPlugins();

    /// Return the ID of the currently active renderer plugin.
    pxr::TfToken GetCurrentRendererId() const;

    /// Return the human-readable display name for a given plugin ID.
    static std::string GetRendererDisplayName(const pxr::TfToken& pluginId);

    /// Switch to a different render delegate.  Returns false if the plugin is
    /// unknown or the switch fails.  The renderer is re-initialised if needed.
    bool SetRendererPlugin(const pxr::TfToken& pluginId);

    // -----------------------------------------------------------------------
    // View settings
    // -----------------------------------------------------------------------
    bool              ShowGrid()                const { return m_showGrid; }
    void              SetShowGrid(bool show)          { m_showGrid = show; }

    bool              ShowCameraGuide()         const { return m_showCameraGuide; }
    void              SetShowCameraGuide(bool show)   { m_showCameraGuide = show; }

    const pxr::TfToken& GetCurrentAov()        const { return m_currentAov; }
    pxr::TfTokenVector  GetRendererAovs()       const;
    bool                SetCurrentAov(const pxr::TfToken& aov);

    bool              GetAAEnabled()            const { return m_aaEnabled; }
    void              SetAAEnabled(bool enabled)      { m_aaEnabled = enabled; }

    const pxr::GfVec3f& GetBackgroundColor()               const { return m_backgroundColor; }
    void                SetBackgroundColor(const pxr::GfVec3f& c) { m_backgroundColor = c; }

    void SetForceRefresh(bool val) {
        m_forceRefresh = m_forceRefresh || val;
        if (val) m_cameraCacheDirty = true;
    }

    /// Current evaluation time (timeline frame).  Drives Hydra rendering,
    /// USD camera resolution, and bounding-box computation.
    pxr::UsdTimeCode GetCurrentTimeCode() const  { return m_currentTime; }
    void SetCurrentTimeCode(pxr::UsdTimeCode t)  { m_currentTime = t; }

    // -----------------------------------------------------------------------
    // Bounding box display
    // -----------------------------------------------------------------------
    BBoxMode GetBBoxMode()             const { return m_bboxMode; }
    void     SetBBoxMode(BBoxMode mode)      { m_bboxMode = mode; }

    const pxr::GfVec4f& GetBBoxColor()                  const { return m_bboxColor; }
    void                 SetBBoxColor(const pxr::GfVec4f& c)   { m_bboxColor = c; }

    // -----------------------------------------------------------------------
    // Lighting settings (mirrors stageView.py viewSettings)
    // -----------------------------------------------------------------------
    /// Camera headlight: single point light at the camera position (default ON).
    bool GetAmbientLightOnly()              const { return m_ambientLightOnly; }
    void SetAmbientLightOnly(bool val)            { m_ambientLightOnly = val; }

    /// Dome environment light (default OFF).
    bool GetDomeLightEnabled()              const { return m_domeLightEnabled; }
    void SetDomeLightEnabled(bool val)            { m_domeLightEnabled = val; }

    ShadingMode GetShadingMode()             const { return m_shadingMode; }
    void        SetShadingMode(ShadingMode m)     { m_shadingMode = m; m_forceRefresh = true; }

    ColorCorrectionMode GetColorCorrectionMode()              const { return m_colorCorrectionMode; }
    void                SetColorCorrectionMode(ColorCorrectionMode m) { m_colorCorrectionMode = m; m_forceRefresh = true; }

    const std::string& GetOcioDisplay()          const { return m_ocioDisplay; }
    void               SetOcioDisplay(std::string v)   { m_ocioDisplay     = std::move(v); m_forceRefresh = true; }
    const std::string& GetOcioView()             const { return m_ocioView; }
    void               SetOcioView(std::string v)      { m_ocioView        = std::move(v); m_forceRefresh = true; }
    const std::string& GetOcioColorSpace()       const { return m_ocioColorSpace; }
    void               SetOcioColorSpace(std::string v){ m_ocioColorSpace  = std::move(v); m_forceRefresh = true; }
    const std::string& GetOcioLook()             const { return m_ocioLook; }
    void               SetOcioLook(std::string v)      { m_ocioLook        = std::move(v); m_forceRefresh = true; }

    /// Default material ambient (kA, default 0.2 — matches viewSettingsDataModel.py).
    float GetDefaultMaterialAmbient()       const { return m_defaultMaterialAmbient; }
    void  SetDefaultMaterialAmbient(float v)      { m_defaultMaterialAmbient = v; }

    /// Default material specular (kS, default 0.1 — matches viewSettingsDataModel.py).
    float GetDefaultMaterialSpecular()      const { return m_defaultMaterialSpecular; }
    void  SetDefaultMaterialSpecular(float v)     { m_defaultMaterialSpecular = v; }

    // -----------------------------------------------------------------------
    // Output
    // -----------------------------------------------------------------------
    uint32_t GetColorTextureID();

    /// Dimensions of the most recent Render() call.
    int GetLastRenderWidth()  const { return m_lastRenderWidth;  }
    int GetLastRenderHeight() const { return m_lastRenderHeight; }

    /// Read the last rendered frame into top-down RGBA8 pixels.
    /// Resolves MSAA before reading. Returns false if the draw target is not ready.
    bool CaptureFrame(std::vector<uint8_t>& outRGBA);

    // For unit tests
    pxr::GlfDrawTargetRefPtr GetDrawTargetForTest() const { return m_drawTarget; }

private:
    int m_lastRenderWidth  = 0;
    int m_lastRenderHeight = 0;
    void InitRenderer();

    void InitGridResources();
    void RebuildGridVBO();          ///< (Re)build grid line geometry after up-axis or size change.
    void DestroyGridResources();
    void RenderGrid(int width, int height);
    void InitAxisResources();
    void DestroyAxisResources();
    void InitBBoxResources();
    void DestroyBBoxResources();
    void InitCamWireResources();
    void DestroyCamWireResources();
    void InitLightWireResources();
    void DestroyLightWireResources();
    /// Draw a single axis-aligned box from a GfRange3d.
    void DrawBox(const pxr::GfRange3d& range, const pxr::GfMatrix4f& mvp);

    /// Build camera wireframe line segments (GL_LINES vertex pairs) in world space
    /// for a single camera given its resolved GfCamera and display scale.
    /// outVerts is appended with interleaved XYZ floats (2 verts per segment).
    void BuildCameraWireframeLines(const pxr::GfCamera& gfCam,
                                   double                scale,
                                   std::vector<float>&   outVerts);

    /// Build light-gizmo line segments (GL_LINES vertex pairs) in world space for
    /// a single light prim.  localToWorld is the light's transform at m_currentTime;
    /// scale is the display-size factor (for direction lines / point markers).
    void BuildLightWireframeLines(const pxr::UsdPrim&    lightPrim,
                                  const pxr::GfMatrix4d& localToWorld,
                                  double                 scale,
                                  std::vector<float>&    outVerts);

    /// True if the stage contains at least one authored UsdLux light.
    /// Used to suppress the default camera headlight when real lights exist.
    bool StageHasAuthoredLights() const;

    /// Project a world-space point to absolute screen coordinates (x=imagePosX+pixelX, etc.).
    /// Returns false when the point is behind the camera.
    static bool WorldToScreen(const pxr::GfVec3d&    world,
                               const pxr::GfMatrix4d& viewProj,
                               int viewW, int viewH,
                               float imagePosX, float imagePosY,
                               float& outX, float& outY);

    /// 2-D point-to-segment distance (all values in screen pixels).
    static float PointToSegmentDist(float px, float py,
                                     float ax, float ay,
                                     float bx, float by);

    pxr::UsdStageRefPtr m_stage;
    std::shared_ptr<pxr::UsdImagingGLEngine> m_renderer;
    pxr::TfToken    m_currentRendererPlugin; ///< empty = use Storm/GL heuristic on init
    pxr::GlfDrawTargetRefPtr                  m_drawTarget;

    pxr::UsdImagingGLRenderParams m_renderParams;

    pxr::GfMatrix4d m_viewMatrix;
    pxr::GfMatrix4d m_projMatrix;
    pxr::SdfPath    m_cameraPath;
    bool            m_useCameraPath;

    // Camera frustum used for picking (set when camera state is provided via GfCamera)
    pxr::GfFrustum  m_cameraFrustum;
    bool            m_hasCameraFrustum = false;

    // Clip planes extracted from GfCamera (passed to renderParams)
    std::vector<pxr::GfVec4d> m_clipPlanes;

    bool            m_showGrid;
    bool            m_showCameraGuide;
    pxr::TfToken    m_currentAov;
    bool            m_aaEnabled;        ///< GL_LINE_SMOOTH anti-aliasing for all line overlays
    pxr::GfVec3f    m_backgroundColor;

    // Lighting settings (mirrors stageView.py viewSettings)
    ShadingMode          m_shadingMode;            // draw mode + lighting flags
    ColorCorrectionMode  m_colorCorrectionMode;    // viewport color correction
    std::string          m_ocioDisplay;
    std::string          m_ocioView;
    std::string          m_ocioColorSpace;
    std::string          m_ocioLook;
    bool            m_ambientLightOnly;       // camera headlight
    bool            m_domeLightEnabled;       // dome/IBL light
    bool            m_stageIsZup;             // used for dome light rotation
    float           m_defaultMaterialAmbient; // kA
    float           m_defaultMaterialSpecular;// kS

    bool m_rendererInitialized;
    bool m_forceRefresh;

    pxr::UsdTimeCode m_currentTime = pxr::UsdTimeCode::Default();

    int m_diagFrameCount;

    // --- Grid GL resources (line geometry, reuses m_axisProgram / m_axisUniformMVP) ---
    GLuint   m_gridVAO;
    GLuint   m_gridVBO;
    // Draw-call ranges inside the VBO (vertex index, count pairs for GL_LINES)
    GLint    m_gridMinorFirst,  m_gridMinorCount;   ///< 1-unit minor lines
    GLint    m_gridMajorFirst,  m_gridMajorCount;   ///< 10-unit major lines
    GLint    m_gridAxisAFirst,  m_gridAxisACount;   ///< X-axis  (red,   at b=0)
    GLint    m_gridAxisBFirst,  m_gridAxisBCount;   ///< Z/Y-axis(blue/green, at a=0)
    float    m_gridHalfSize;   ///< Half-extent (default 50 → 100×100 grid)

    // --- Axis GLSL resources (stageView.DrawAxis port) ---
    GLuint m_axisVAO;
    GLuint m_axisVBO;
    GLuint m_axisProgram;
    GLint  m_axisUniformMVP;
    GLint  m_axisUniformColor;

    // --- Bounding box GLSL resources ---
    BBoxMode     m_bboxMode;
    pxr::GfVec4f m_bboxColor;   // default: white
    GLuint m_bboxVAO;
    GLuint m_bboxVBO;           // dynamic: updated per box
    GLuint m_bboxProgram;
    GLint  m_bboxUniformMVP;
    GLint  m_bboxUniformColor;

    // --- Camera wireframe GL resources (dedicated, separate from bbox) ---
    GLuint m_camWireVAO  = 0;
    GLuint m_camWireVBO  = 0;   ///< dynamic VBO; reallocated per frame as needed

    // --- Light wireframe GL resources (mirrors camera wireframe) ---
    GLuint m_lightWireVAO = 0;
    GLuint m_lightWireVBO = 0;  ///< dynamic VBO; reallocated per frame as needed

    // --- Camera wireframe cache ---
    pxr::SdfPathVector m_cachedCameraPaths;
    bool               m_cameraCacheDirty = true;

};

} // namespace UsdLayerManager