#include "ViewportCamera.h"

#include <pxr/usd/usdGeom/metrics.h>
#include <pxr/usd/usdGeom/tokens.h>
#include <pxr/usd/usdGeom/camera.h>
#include <pxr/base/gf/rotation.h>
#include <pxr/base/gf/frustum.h>
#include <pxr/imaging/cameraUtil/conformWindow.h>

#include <cmath>
#include <algorithm>

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

namespace UsdLayerManager {

// ---------------------------------------------------------------------------
// Helper: rotation matrix around `axis` by `angleDeg` degrees
// ---------------------------------------------------------------------------
static pxr::GfMatrix4d RotMatrix(const pxr::GfVec3d& axis, double angleDeg)
{
    return pxr::GfMatrix4d(1.0).SetRotate(pxr::GfRotation(axis, angleDeg));
}

// ===========================================================================
// Construction
// ===========================================================================

ViewportCamera::ViewportCamera()
    : m_cameraTransformDirty(true)
    , m_rotTheta(0.0)
    , m_rotPhi(0.0)
    , m_rotPsi(0.0)
    , m_center(0.0, 0.0, 0.0)
    , m_dist(100.0)
    , m_selSize(10.0)
    , m_isZUp(false)
    , m_YZUpMatrix(1.0)
    , m_YZUpInvMatrix(1.0)
    , m_hasClosestVisibleDist(false)
    , m_closestVisibleDist(0.0)
    , m_lastFramedDist(100.0)
    , m_lastFramedClosestDist(0.0)
    , m_overrideNear(-1.0)
    , m_overrideFar(-1.0)
    , m_mode(CameraMode::Free)
{
    // Default: perspective camera, vertical FOV = 60°, square aspect ratio.
    m_camera.SetPerspectiveFromAspectRatioAndFieldOfView(
        1.0f, 60.0f, pxr::GfCamera::FOVVertical);
    m_camera.SetFocusDistance(static_cast<float>(m_dist));
    ResetClippingPlanes();
}

// ===========================================================================
// Stage
// ===========================================================================

void ViewportCamera::SetStage(pxr::UsdStageRefPtr stage)
{
    m_stage = stage;
    m_isZUp = stage &&
              (pxr::UsdGeomGetStageUpAxis(stage) == pxr::UsdGeomTokens->z);

    if (m_isZUp) {
        // GfCamera.Y_UP_TO_Z_UP_MATRIX: rotate -90° around X axis
        m_YZUpMatrix    = pxr::GfMatrix4d(1.0).SetRotate(
                              pxr::GfRotation(pxr::GfVec3d::XAxis(), -90.0));
        m_YZUpInvMatrix = m_YZUpMatrix.GetInverse();
    } else {
        m_YZUpMatrix    = pxr::GfMatrix4d(1.0);
        m_YZUpInvMatrix = pxr::GfMatrix4d(1.0);
    }
    m_cameraTransformDirty = true;
}

// ===========================================================================
// Private: rebuild camera transform from orbital parameters
// Mirrors FreeCamera._pushToCameraTransform()
// ===========================================================================

void ViewportCamera::PushToCameraTransform()
{
    if (!m_cameraTransformDirty) return;

    // camera-to-world transform (same as FreeCamera._pushToCameraTransform):
    //   T(dist*Z) * R(-psi,Z) * R(-phi,X) * R(-theta,Y) * YZUpInv * T(center)
    pxr::GfMatrix4d xform =
        pxr::GfMatrix4d(1.0).SetTranslate(pxr::GfVec3d::ZAxis() * m_dist)
        * RotMatrix(pxr::GfVec3d::ZAxis(), -m_rotPsi)
        * RotMatrix(pxr::GfVec3d::XAxis(), -m_rotPhi)
        * RotMatrix(pxr::GfVec3d::YAxis(), -m_rotTheta)
        * m_YZUpInvMatrix
        * pxr::GfMatrix4d(1.0).SetTranslate(m_center);

    m_camera.SetTransform(xform);
    m_camera.SetFocusDistance(static_cast<float>(m_dist));
    m_cameraTransformDirty = false;
}

// Mirrors FreeCamera._pullFromCameraTransform()
void ViewportCamera::PullFromCameraTransform()
{
    pxr::GfFrustum frustum = m_camera.GetFrustum();
    m_dist    = static_cast<double>(m_camera.GetFocusDistance());
    m_selSize = m_dist / 10.0;
    m_center  = frustum.GetPosition() + m_dist * frustum.ComputeViewDirection();

    pxr::GfMatrix4d camTransform = m_camera.GetTransform() * m_YZUpMatrix;
    camTransform.Orthonormalize();
    pxr::GfRotation rotation = camTransform.ExtractRotation();

    // Decompose: Y → theta, X → phi, Z → psi
    pxr::GfVec3d angles = rotation.Decompose(
        pxr::GfVec3d::YAxis(),
        pxr::GfVec3d::XAxis(),
        pxr::GfVec3d::ZAxis());

    m_rotTheta = -angles[0];
    m_rotPhi   = -angles[1];
    m_rotPsi   = -angles[2];
    m_cameraTransformDirty = true;
}

// ===========================================================================
// Free Camera Operations
// ===========================================================================

void ViewportCamera::Tumble(double dTheta, double dPhi)
{
    m_rotTheta += dTheta;
    m_rotPhi   += dPhi;
    m_cameraTransformDirty = true;
}

void ViewportCamera::AdjustDistance(double scaleFactor)
{
    // Mirrors FreeCamera.AdjustDistance(): prevents getting stuck near dist≈0.
    if (scaleFactor > 1.0 && m_dist < 2.0) {
        double selBasedIncr = m_selSize / 25.0;
        scaleFactor -= 1.0;
        m_dist += std::min(selBasedIncr, scaleFactor);
    } else {
        m_dist *= scaleFactor;
    }
    m_dist = std::max(m_dist, 0.001); // never let dist reach zero

    // Keep closest-visible-distance estimate in sync with new dist
    if (m_hasClosestVisibleDist) {
        if (m_dist > m_lastFramedDist) {
            m_closestVisibleDist = m_lastFramedClosestDist;
        } else {
            m_closestVisibleDist = m_lastFramedClosestDist
                                 - m_lastFramedDist
                                 + m_dist;
        }
    }
    m_cameraTransformDirty = true;
}

void ViewportCamera::Truck(double deltaRight, double deltaUp)
{
    PushToCameraTransform();
    pxr::GfFrustum frustum  = m_camera.GetFrustum();
    pxr::GfVec3d   camUp    = frustum.ComputeUpVector();
    pxr::GfVec3d   camRight = pxr::GfCross(frustum.ComputeViewDirection(), camUp);
    m_center += deltaRight * camRight + deltaUp * camUp;
    m_cameraTransformDirty = true;
}

double ViewportCamera::ComputePixelsToWorldFactor(double viewportHeight)
{
    PushToCameraTransform();
    pxr::GfFrustum frustum      = m_camera.GetFrustum();
    double         frustumHeight = frustum.GetWindow().GetSize()[1];
    return frustumHeight * m_dist / std::max(viewportHeight, 1.0);
}

void ViewportCamera::FrameSelection(const pxr::GfBBox3d& selBBox, double frameFit)
{
    m_hasClosestVisibleDist = false;
    m_center  = selBBox.ComputeCentroid();

    pxr::GfRange3d selRange = selBBox.ComputeAlignedRange();
    pxr::GfVec3d   sz       = selRange.GetSize();
    m_selSize = std::max({ sz[0], sz[1], sz[2] });

    // Distance calculation from FreeCamera.frameSelection()
    double fovRad     = GetFOV() * M_PI / 180.0;
    double halfFovRad = std::max(fovRad * 0.5, 0.00872665); // at least ~0.5°
    double lengthToFit = m_selSize * frameFit * 0.5;
    m_dist = lengthToFit / std::tan(halfFovRad);

    // Prevent camera from intersecting the bounding box
    if (m_dist < kDefaultNear + m_selSize * 0.5) {
        m_dist = kDefaultNear + lengthToFit;
    }
    m_cameraTransformDirty = true;
}

// ===========================================================================
// Clipping Planes
// ===========================================================================

std::pair<double,double> ViewportCamera::RangeOfBoxAlongRay(
    const pxr::GfRay& camRay, const pxr::GfBBox3d& bbox) const
{
    double maxDist = -1e38;
    double minDist =  1e38;
    const pxr::GfRange3d&  boxRange = bbox.GetRange();
    const pxr::GfMatrix4d& boxXform = bbox.GetMatrix();

    for (int i = 0; i < 8; ++i) {
        pxr::GfVec3d corner = boxXform.Transform(boxRange.GetCorner(i));
        double t = 0.0;
        camRay.FindClosestPoint(corner, &t);
        maxDist = std::max(maxDist, t);
        minDist = std::min(minDist, t);
    }
    minDist = (minDist < kDefaultNear) ? kDefaultNear : minDist * 0.99;
    maxDist *= 1.01;
    return { minDist, maxDist };
}

void ViewportCamera::SetClippingPlanes(const pxr::GfBBox3d& stageBBox)
{
    double computedNear, computedFar;

    if (stageBBox.GetRange().IsEmpty()) {
        computedNear = kDefaultNear;
        computedFar  = kDefaultFar;
    } else {
        pxr::GfFrustum frustum = m_camera.GetFrustum();
        pxr::GfVec3d   camPos  = frustum.GetPosition();
        pxr::GfRay     camRay(camPos, frustum.ComputeViewDirection());

        auto boxRange = RangeOfBoxAlongRay(camRay, stageBBox);
        computedNear = boxRange.first;
        computedFar  = boxRange.second;

        double precisionNear = computedFar / kMaxGoodZResolution;
        if (m_hasClosestVisibleDist) {
            double halfClose = m_closestVisibleDist / 2.0;
            if (m_closestVisibleDist < m_lastFramedClosestDist) {
                halfClose = std::max({ precisionNear, halfClose, computedNear });
            }
            if (halfClose < computedNear) {
                computedNear = halfClose;
            } else if (precisionNear > computedNear) {
                computedNear = std::min((precisionNear + halfClose) / 2.0, halfClose);
            }
        }
    }

    double nearVal = (m_overrideNear > 0.0) ? m_overrideNear : computedNear;
    double farVal  = (m_overrideFar  > 0.0) ? m_overrideFar  : computedFar;
    farVal = std::max(nearVal + 1.0, farVal);
    m_camera.SetClippingRange(pxr::GfRange1f(
        static_cast<float>(nearVal), static_cast<float>(farVal)));
}

void ViewportCamera::ResetClippingPlanes()
{
    double nearVal = (m_overrideNear > 0.0) ? m_overrideNear : kDefaultNear;
    double farVal  = (m_overrideFar  > 0.0) ? m_overrideFar  : kDefaultFar;
    m_camera.SetClippingRange(pxr::GfRange1f(
        static_cast<float>(nearVal), static_cast<float>(farVal)));
}

// ===========================================================================
// Camera Resolution
// ===========================================================================

pxr::GfCamera ViewportCamera::ComputeGfCamera(
    const pxr::GfBBox3d& stageBBox, bool autoClip)
{
    PushToCameraTransform();
    if (autoClip) {
        SetClippingPlanes(stageBBox);
    } else {
        ResetClippingPlanes();
    }
    return m_camera;
}

void ViewportCamera::SetClosestVisibleDistFromPoint(const pxr::GfVec3d& point)
{
    PushToCameraTransform();
    pxr::GfFrustum frustum = m_camera.GetFrustum();
    pxr::GfVec3d   camPos  = frustum.GetPosition();
    pxr::GfRay     camRay(camPos, frustum.ComputeViewDirection());
    double t = 0.0;
    camRay.FindClosestPoint(point, &t);
    m_closestVisibleDist    = t;
    m_hasClosestVisibleDist = true;
    m_lastFramedDist        = m_dist;
    m_lastFramedClosestDist = m_closestVisibleDist;
}

// ===========================================================================
// Matrix Accessors
// ===========================================================================

pxr::GfMatrix4d ViewportCamera::GetViewMatrix()
{
    PushToCameraTransform();
    return m_camera.GetFrustum().ComputeViewMatrix();
}

pxr::GfMatrix4d ViewportCamera::GetProjectionMatrix()
{
    PushToCameraTransform();
    return m_camera.GetFrustum().ComputeProjectionMatrix();
}

// ===========================================================================
// Compatibility Accessors
// ===========================================================================

pxr::GfVec3d ViewportCamera::GetEye()
{
    PushToCameraTransform();
    return m_camera.GetFrustum().GetPosition();
}

// ===========================================================================
// Camera Mode
// ===========================================================================

void ViewportCamera::SetUsdCamera(const pxr::SdfPath& cameraPath)
{
    m_mode = CameraMode::UsdCamera;
    m_usdCameraPath = cameraPath;
}

void ViewportCamera::InitOrbitFromEyeAndCenter(
    const pxr::GfVec3d& eye, const pxr::GfVec3d& center, double dist)
{
    m_center  = center;
    m_dist    = std::max(dist, 0.001);
    m_selSize = m_dist / 10.0;

    // Transform the eye-to-center offset into the Y-up orbital frame.
    // For Y-up stages m_YZUpMatrix is identity; for Z-up it's -90° around X.
    pxr::GfVec3d offset = m_YZUpMatrix.TransformDir(eye - center);
    double len = offset.GetLength();
    if (len > 1e-6) {
        offset /= len;
        // In the Y-up orbital frame PushToCameraTransform places the eye at:
        //   (-sin(theta)*cos(phi), sin(phi), cos(theta)*cos(phi))
        m_rotPhi   = std::asin(std::max(-1.0, std::min(1.0, offset[1])))
                     * (180.0 / M_PI);
        m_rotTheta = std::atan2(-offset[0], offset[2]) * (180.0 / M_PI);
    }
    m_rotPsi = 0.0;  // zero roll — avoids gimbal weirdness on switch
    m_cameraTransformDirty = true;
}

void ViewportCamera::SwitchToFreeCamera(const pxr::GfCamera* lastGfCamera)
{
    if (m_mode == CameraMode::Free) return;

    if (lastGfCamera) {
        // Initialize free-camera state from the last rendered GfCamera
        // (mirrors FreeCamera.FromGfCamera)
        m_camera = *lastGfCamera;
        PullFromCameraTransform();
    }
    m_mode = CameraMode::Free;
    m_usdCameraPath = pxr::SdfPath();
    m_cameraTransformDirty = true;
}

// ===========================================================================
// Camera Settings
// ===========================================================================

double ViewportCamera::GetFOV() const
{
    return static_cast<double>(
        m_camera.GetFieldOfView(pxr::GfCamera::FOVVertical));
}

void ViewportCamera::SetFOV(double fov)
{
    m_camera.SetPerspectiveFromAspectRatioAndFieldOfView(
        m_camera.GetAspectRatio(),
        static_cast<float>(fov),
        pxr::GfCamera::FOVVertical);
}

double ViewportCamera::GetAspectRatio() const
{
    return static_cast<double>(m_camera.GetAspectRatio());
}

void ViewportCamera::SetAspectRatio(double aspect)
{
    m_camera.SetPerspectiveFromAspectRatioAndFieldOfView(
        static_cast<float>(aspect),
        m_camera.GetFieldOfView(pxr::GfCamera::FOVVertical),
        pxr::GfCamera::FOVVertical);
}

} // namespace UsdLayerManager