#include "ViewportWidget.h"
#include "Snapping.h"
#include "Camera.h"
#include "ViewCube.h"
#include "SketchGrid.h"
#include "Document.h"
#include "FeatureBrowser.h"
#include "SketchFeature.h"
#include "SketchLine.h"
#include "SketchRectangle.h"
#include "SketchCircle.h"
#include "SketchObject.h"
#include "ApplicationController.h"
#include "RectangleTool.h"
#include "LineTool.h"
#include "CircleTool.h"
#include <QMouseEvent>
#include <QKeyEvent>
#include <QWheelEvent>
#include <QApplication>
#include <QPainter>
#include <QPixmap>
#include <QImage>
#include <QSvgRenderer>
#include <QWheelEvent>
#include <QApplication>
#include <cmath>
#include <QtMath>
#include <QOpenGLShaderProgram>
#include <QVector>
#include <map>

struct PntComparator {
    bool operator()(const gp_Pnt& a, const gp_Pnt& b) const {
        // A tolerance is needed for floating point comparisons
        double tol = 1e-9;
        if (std::abs(a.X() - b.X()) > tol) return a.X() < b.X();
        if (std::abs(a.Y() - b.Y()) > tol) return a.Y() < b.Y();
        if (std::abs(a.Z() - b.Z()) > tol) return a.Z() < b.Z();
        return false; // They are considered equal
    }
};

ViewportWidget::ViewportWidget(QWidget *parent)
    : QOpenGLWidget(parent)
{
    m_camera = new Camera(this);
    connect(m_camera, &Camera::cameraChanged, this, QOverload<>::of(&QWidget::update));
    connect(m_camera, &Camera::restoreStateAnimationFinished, this, &ViewportWidget::onRestoreStateAnimationFinished);
    m_viewCube = new ViewCube();
    m_sketchGrid = new SketchGrid(this);
    m_featureBrowser = new FeatureBrowser();
    setMouseTracking(true);
    setFocusPolicy(Qt::StrongFocus);

    m_currentPlane = SketchPlane::XY;

    m_toolIcons.insert(static_cast<int>(ApplicationController::ToolType::Line), new QSvgRenderer(QString(":/icons/line.svg"), this));
    m_toolIcons.insert(static_cast<int>(ApplicationController::ToolType::Rectangle), new QSvgRenderer(QString(":/icons/rectangle.svg"), this));
    m_toolIcons.insert(static_cast<int>(ApplicationController::ToolType::Circle), new QSvgRenderer(QString(":/icons/circle.svg"), this));

    m_cursorRenderer = new QSvgRenderer(QString(":/icons/cursor.svg"), this);

    m_sketchTools.insert(static_cast<int>(ApplicationController::ToolType::Line), new LineTool(this));
    m_sketchTools.insert(static_cast<int>(ApplicationController::ToolType::Rectangle), new RectangleTool(this));
    m_sketchTools.insert(static_cast<int>(ApplicationController::ToolType::Circle), new CircleTool(this));

    m_snapping = new Snapping(this);
}

ViewportWidget::~ViewportWidget()
{
    makeCurrent();
    delete m_shaderProgram;
    delete m_viewCube;
    delete m_sketchGrid;
    m_vbo.destroy();
    m_vao.destroy();
    doneCurrent();

    delete m_featureBrowser;
    delete m_snapping;
}

void ViewportWidget::setDocument(Document* document)
{
    m_document = document;
    m_featureBrowser->setDocument(document);
}

void ViewportWidget::initializeGL()
{
    initializeOpenGLFunctions();
    glClearColor(0.2f, 0.3f, 0.3f, 1.0f);

    initShaders();

    m_vao.create();
    QOpenGLVertexArrayObject::Binder vaoBinder(&m_vao);
    m_vbo.create();
    m_vbo.bind();
    m_vbo.allocate(nullptr, 0); // Allocate when drawing

    // Position attribute
    m_shaderProgram->enableAttributeArray(0);
    m_shaderProgram->setAttributeBuffer(0, GL_FLOAT, 0, 3, 0);

    m_vbo.release();

    glEnable(GL_DEPTH_TEST);
    m_viewCube->initializeGL();
    m_sketchGrid->initializeGL();
}

void ViewportWidget::paintGL()
{
    const qreal retinaScale = devicePixelRatio();

    // Main scene rendering
    glViewport(0, 0, width() * retinaScale, height() * retinaScale);
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

    if (!m_shaderProgram || !m_shaderProgram->isLinked()) {
        return;
    }

    glEnable(GL_DEPTH_TEST);
    glEnable(GL_CULL_FACE);

    QMatrix4x4 model = m_camera->modelViewMatrix();
    m_shaderProgram->bind();
    m_shaderProgram->setUniformValue(m_projMatrixLoc, projection);
    m_shaderProgram->setUniformValue(m_mvMatrixLoc, model);

    QOpenGLVertexArrayObject::Binder vaoBinder(&m_vao);

    // Sketch grid rendering
    if (m_isSelectingPlane) {
        if (m_highlightedPlane != SketchPlane::NONE) {
            m_sketchGrid->paintGL(static_cast<SketchGrid::SketchPlane>(m_highlightedPlane), m_shaderProgram, m_colorLoc);
        }
    }
    if (m_currentPlane != SketchPlane::NONE) {
        m_sketchGrid->paintGL(static_cast<SketchGrid::SketchPlane>(m_currentPlane), m_shaderProgram, m_colorLoc);
    }

    if (m_isSelectingPlane) {
        drawSelectionPlanes();
    }

    if (m_document) {
        for (Feature* feature : m_document->features()) {
            if (auto sketch = dynamic_cast<SketchFeature*>(feature)) {
                drawSketch(sketch);
            }
        }
    }

    m_snapping->paintGL();

    if (m_camera->isRotating()) {
        const float radius = 0.004f * -m_camera->uiCameraDistance();
        const int numSegments = 16;
        QMatrix4x4 invModelView = m_camera->modelViewMatrix().inverted();
        QVector3D rightVec = invModelView.column(0).toVector3D();
        QVector3D upVec = invModelView.column(1).toVector3D();

        QVector<GLfloat> circleFillVertices;
        QVector3D center = m_camera->rotationPivot();

        circleFillVertices << center.x() << center.y() << center.z(); // Center vertex for fan
        for (int i = 0; i <= numSegments; ++i) { // <= to close the circle
            float angle = (2.0f * M_PI * float(i)) / float(numSegments);
            QVector3D p = center + radius * (cos(angle) * rightVec + sin(angle) * upVec);
            circleFillVertices << p.x() << p.y() << p.z();
        }

        m_shaderProgram->setUniformValue(m_colorLoc, QVector4D(0.0f, 0.0f, 1.0f, 1.0f)); // Blue
        m_vbo.bind();
        m_vbo.allocate(circleFillVertices.constData(), circleFillVertices.size() * sizeof(GLfloat));
        glDrawArrays(GL_TRIANGLE_FAN, 0, circleFillVertices.size() / 3);
    }

    if (m_activeSketchTool) {
        m_activeSketchTool->paintGL();
    }

    m_shaderProgram->release();

    // View cube rendering
    QMatrix4x4 viewCubeModel;
    viewCubeModel.rotate(m_camera->xRotation() / 16.0f, 1, 0, 0);
    viewCubeModel.rotate(m_camera->yRotation() / 16.0f, 0, 1, 0);
    m_viewCube->paintGL(m_shaderProgram, m_colorLoc, viewCubeModel, width() * retinaScale, height() * retinaScale, m_currentMousePos);

    glViewport(0, 0, width() * retinaScale, height() * retinaScale);
    glDisable(GL_DEPTH_TEST);
    glDisable(GL_CULL_FACE);

    QPainter painter(this);
    if (m_currentPlane != SketchPlane::NONE) {
        m_sketchGrid->paintAxisLabels(painter, static_cast<SketchGrid::SketchPlane>(m_currentPlane), model, projection);
    }
    m_featureBrowser->paint(painter, width(), height());
    m_viewCube->paint2D(painter, width(), height());


    if (m_activeSketchTool) {
        m_activeSketchTool->paint2D(painter, model, projection);
    }

    painter.setPen(Qt::white);
    painter.drawText(width() - 350, height() - 10, QString("Camera Distance: %1").arg(-m_camera->uiCameraDistance()));

    painter.end();
}

void ViewportWidget::resizeGL(int w, int h)
{
    projection.setToIdentity();
    projection.perspective(45.0f, w / float(h), 0.01f, 10000.0f);
}

void ViewportWidget::mousePressEvent(QMouseEvent *event)
{
    m_camera->mousePressEvent(event);

    if (event->button() == Qt::MiddleButton && !(QApplication::keyboardModifiers() & Qt::ShiftModifier)) {
        m_camera->startRotation(unproject(event->pos(), m_currentPlane));
        update();
    }

    if (event->button() == Qt::LeftButton) {
        if (m_viewCube->handleMousePress(event->pos(), width(), height())) {
            m_camera->animateToHomeView();
            update();
            return;
        }

        if (m_isSelectingPlane) {
            if (m_highlightedPlane != SketchPlane::NONE) {
                emit planeSelected(m_highlightedPlane);
                m_isSelectingPlane = false;
                m_highlightedPlane = SketchPlane::NONE;
                update();
            }
            return;
        }

        if (m_activeSketchTool) {
            m_activeSketchTool->mousePressEvent(event);
            update();
            return;
        }

    }
}

void ViewportWidget::mouseMoveEvent(QMouseEvent *event)
{
    m_currentMousePos = event->pos();

    if (m_isSelectingPlane) {
        SketchPlane newHighlight = checkPlaneSelection(m_currentMousePos);
        if (newHighlight != m_highlightedPlane) {
            m_highlightedPlane = newHighlight;
            update();
        }
    }

    if (m_snapping->update(m_currentMousePos)) {
        update();
    }

    if (m_activeSketchTool) {
        m_activeSketchTool->mouseMoveEvent(event);
    }

    m_camera->mouseMoveEvent(event, height());
    update();
}

void ViewportWidget::wheelEvent(QWheelEvent *event)
{
    QVector3D worldPos = unproject(event->position().toPoint(), m_currentPlane);
    m_camera->wheelEvent(event, worldPos);
}

void ViewportWidget::keyPressEvent(QKeyEvent *event)
{
    if (m_activeSketchTool) {
        m_activeSketchTool->keyPressEvent(event);
        return;
    }

    if (event->key() == Qt::Key_Escape) {
		if (m_isSelectingPlane) {
            m_isSelectingPlane = false;
            m_highlightedPlane = SketchPlane::NONE;
            m_currentPlane = SketchPlane::XY;
            update();
            return;
        }
		else if (m_activeTool) {
            emit toolDeactivated();
            update();
            return;
        }
    }
    QOpenGLWidget::keyPressEvent(event);
}

void ViewportWidget::mouseReleaseEvent(QMouseEvent *event)
{
    if (event->button() == Qt::MiddleButton) {
        m_camera->stopRotation();
        update();
    }
    QOpenGLWidget::mouseReleaseEvent(event);
}

void ViewportWidget::addLine(const gp_Pnt& start, const gp_Pnt& end)
{
    emit lineAdded(start, end);
}

void ViewportWidget::deactivateActiveTool()
{
    emit toolDeactivated();
}

bool ViewportWidget::isSnappingOrigin() const
{
    return m_snapping->isSnappingOrigin();
}

bool ViewportWidget::isSnappingVertex() const
{
    return m_snapping->isSnappingVertex();
}

const gp_Pnt& ViewportWidget::snapVertex() const
{
    return m_snapping->snapVertex();
}

void ViewportWidget::setSnappingHorizontal(bool snapping)
{
    m_isSnappingHorizontal = snapping;
}

void ViewportWidget::setSnappingVertical(bool snapping)
{
    m_isSnappingVertical = snapping;
}

bool ViewportWidget::focusNextPrevChild(bool next)
{
    if (m_activeTool != static_cast<int>(ApplicationController::ToolType::None)) {
        return false;
    }
    return QOpenGLWidget::focusNextPrevChild(next);
}

void ViewportWidget::onSketchModeStarted(SketchPlane plane)
{
    m_currentPlane = plane;

    m_camera->saveState();
    m_camera->animateToPlaneView(static_cast<int>(plane));
}

void ViewportWidget::onPlaneSelectionModeStarted()
{
    m_isSelectingPlane = true;
    m_highlightedPlane = SketchPlane::NONE;
    m_currentPlane = SketchPlane::NONE;
    update();
}

void ViewportWidget::onSketchModeEnded()
{
    m_camera->animateRestoreState();
}

void ViewportWidget::onRestoreStateAnimationFinished()
{
    // Return to showing the base XY grid when not in a sketch
    m_currentPlane = SketchPlane::XY;
    update();
}

QVector3D ViewportWidget::project(const QVector3D& worldCoord, const QMatrix4x4& modelView, const QMatrix4x4& projection, const QRect& viewport)
{
    QVector4D clipCoord = projection * modelView * QVector4D(worldCoord, 1.0);
    if (qFuzzyCompare(clipCoord.w(), 0.0f)) {
        return QVector3D(-1, -1, -1);
    }

    QVector3D ndc(clipCoord.x() / clipCoord.w(), clipCoord.y() / clipCoord.w(), clipCoord.z() / clipCoord.w());

    float winX = viewport.x() + viewport.width() * (ndc.x() + 1.0) / 2.0;
    float winY = viewport.y() + viewport.height() * (1.0 - (ndc.y() + 1.0) / 2.0); // Y is inverted

    return QVector3D(winX, winY, (ndc.z() + 1.0) / 2.0);
}

void ViewportWidget::onActiveToolChanged(int tool)
{
    if (m_activeSketchTool) {
        m_activeSketchTool->deactivate();
        m_activeSketchTool = nullptr;
    }

    m_activeTool = tool;

    if (m_sketchTools.contains(tool)) {
        m_activeSketchTool = m_sketchTools.value(tool);
        m_activeSketchTool->activate();
    }

    if (tool == static_cast<int>(ApplicationController::ToolType::None)) {
        unsetCursor();
    } else {
        if (m_toolIcons.contains(tool)) {
            QSvgRenderer* renderer = m_toolIcons.value(tool);
            if (renderer && renderer->isValid()) {
                const QSize cursorSize(48, 48);
                QPixmap cursorPixmap(cursorSize);
                cursorPixmap.fill(Qt::transparent);
                QPainter painter(&cursorPixmap);
                painter.setRenderHint(QPainter::Antialiasing);

                // Render arrow cursor from SVG
                if (m_cursorRenderer && m_cursorRenderer->isValid()) {
                    m_cursorRenderer->render(&painter, QRectF(0, 0, 14, 22));
                }

                // Render tool icon
                const QSize iconSize(32, 32);
                QPixmap iconPixmap(iconSize);
                iconPixmap.fill(Qt::transparent);
                QPainter iconPainter(&iconPixmap);
                renderer->render(&iconPainter);
                iconPainter.end();

                // Invert and draw icon onto cursor pixmap
                QImage iconImage = iconPixmap.toImage();
                iconImage.invertPixels(QImage::InvertRgb);
                painter.drawImage(QRect(12, 12, iconSize.width(), iconSize.height()), iconImage);
                painter.end();

                setCursor(QCursor(cursorPixmap, 0, 0));
            }
        } else {
            unsetCursor();
        }
    }
}

QVector3D ViewportWidget::unproject(const QPoint& screenPos, SketchPlane plane)
{
    QMatrix4x4 model = m_camera->modelViewMatrix();

    bool invertible;
    QMatrix4x4 inv = (projection * model).inverted(&invertible);
    if (!invertible) {
        return QVector3D();
    }

    float ndcX = (2.0f * screenPos.x()) / width() - 1.0f;
    float ndcY = 1.0f - (2.0f * screenPos.y()) / height();

    QVector4D nearPoint_ndc(ndcX, ndcY, -1.0f, 1.0f);
    QVector4D farPoint_ndc(ndcX, ndcY, 1.0f, 1.0f);

    QVector4D nearPoint_world = inv * nearPoint_ndc;
    QVector4D farPoint_world = inv * farPoint_ndc;

    if (qFuzzyCompare(nearPoint_world.w(), 0.0f) || qFuzzyCompare(farPoint_world.w(), 0.0f)) {
        return QVector3D();
    }

    nearPoint_world /= nearPoint_world.w();
    farPoint_world /= farPoint_world.w();

    QVector3D rayOrigin(nearPoint_world);
    QVector3D rayDir = (QVector3D(farPoint_world) - rayOrigin).normalized();

    QVector3D planeNormal;
    switch (plane) {
        case SketchPlane::XY: planeNormal = QVector3D(0, 0, 1); break;
        case SketchPlane::XZ: planeNormal = QVector3D(0, 1, 0); break;
        case SketchPlane::YZ: planeNormal = QVector3D(1, 0, 0); break;
        case SketchPlane::NONE: return QVector3D();
    }

    float denom = QVector3D::dotProduct(planeNormal, rayDir);
    if (qAbs(denom) > 1e-6) {
        QVector3D p0(0,0,0);
        float t = QVector3D::dotProduct(p0 - rayOrigin, planeNormal) / denom;
        return rayOrigin + t * rayDir;
    }

    return QVector3D();
}

void ViewportWidget::drawSketch(const SketchFeature* sketch)
{
    glDisable(GL_DEPTH_TEST);
    glLineWidth(2.0f);

    QVector<GLfloat> lineVertices;
    std::map<gp_Pnt, int, PntComparator> vertexCounts;

    for (const auto& obj : sketch->objects()) {
        if (obj->type() == SketchObject::ObjectType::Line) {
            auto line = static_cast<const SketchLine*>(obj);
            const auto& start = line->startPoint();
            const auto& end = line->endPoint();
            lineVertices << start.X() << start.Y() << start.Z();
            lineVertices << end.X() << end.Y() << end.Z();

            vertexCounts[start]++;
            vertexCounts[end]++;
        } else if (obj->type() == SketchObject::ObjectType::Circle) {
            auto circle = static_cast<const SketchCircle*>(obj);
            const auto& center = circle->center();
            double radius = circle->radius();
            QVector3D centerPos(center.X(), center.Y(), center.Z());

            const int numSegments = 64;
            QVector3D u_axis, v_axis;
            switch (sketch->plane()) {
                case SketchFeature::SketchPlane::XY: // Top
                    u_axis = QVector3D(1, 0, 0);
                    v_axis = QVector3D(0, 1, 0);
                    break;
                case SketchFeature::SketchPlane::XZ: // Front
                    u_axis = QVector3D(1, 0, 0);
                    v_axis = QVector3D(0, 0, 1);
                    break;
                case SketchFeature::SketchPlane::YZ: // Right
                    u_axis = QVector3D(0, 1, 0);
                    v_axis = QVector3D(0, 0, 1);
                    break;
            }

            for (int i = 0; i < numSegments; ++i) {
                double angle1 = 2.0 * M_PI * double(i) / double(numSegments);
                QVector3D p1 = centerPos + radius * (cos(angle1) * u_axis + sin(angle1) * v_axis);

                double angle2 = 2.0 * M_PI * double(i + 1) / double(numSegments);
                QVector3D p2 = centerPos + radius * (cos(angle2) * u_axis + sin(angle2) * v_axis);

                lineVertices << p1.x() << p1.y() << p1.z();
                lineVertices << p2.x() << p2.y() << p2.z();
            }
        } else if (obj->type() == SketchObject::ObjectType::Rectangle) {
            auto rect = static_cast<const SketchRectangle*>(obj);
            const auto& p1 = rect->corner1();
            const auto& p3 = rect->corner2();

            gp_Pnt p2, p4;
            if (sketch->plane() == SketchFeature::SketchPlane::XY) {
                p2.SetCoord(p3.X(), p1.Y(), p1.Z());
                p4.SetCoord(p1.X(), p3.Y(), p1.Z());
            } else if (sketch->plane() == SketchFeature::SketchPlane::XZ) {
                p2.SetCoord(p3.X(), p1.Y(), p1.Z());
                p4.SetCoord(p1.X(), p1.Y(), p3.Z());
            } else if (sketch->plane() == SketchFeature::SketchPlane::YZ) {
                p2.SetCoord(p1.X(), p3.Y(), p1.Z());
                p4.SetCoord(p1.X(), p1.Y(), p3.Z());
            }

            lineVertices << p1.X() << p1.Y() << p1.Z();
            lineVertices << p2.X() << p2.Y() << p2.Z();
            lineVertices << p2.X() << p2.Y() << p2.Z();
            lineVertices << p3.X() << p3.Y() << p3.Z();
            lineVertices << p3.X() << p3.Y() << p3.Z();
            lineVertices << p4.X() << p4.Y() << p4.Z();
            lineVertices << p4.X() << p4.Y() << p4.Z();
            lineVertices << p1.X() << p1.Y() << p1.Z();

            vertexCounts[p1] += 2;
            vertexCounts[p2] += 2;
            vertexCounts[p3] += 2;
            vertexCounts[p4] += 2;
        }
    }

    m_shaderProgram->setUniformValue(m_colorLoc, QVector4D(1.0f, 1.0f, 1.0f, 1.0f));
    m_vbo.bind();

    if (!lineVertices.isEmpty()) {
        m_vbo.allocate(lineVertices.constData(), lineVertices.size() * sizeof(GLfloat));
        glDrawArrays(GL_LINES, 0, lineVertices.size() / 3);
    }

    QVector<gp_Pnt> unattachedVertices;
    for (const auto& pair : vertexCounts) {
        if (pair.second == 1) {
            unattachedVertices.append(pair.first);
        }
    }

    if (!unattachedVertices.isEmpty()) {
        const float radius = 0.004f * -m_camera->uiCameraDistance();
        const int numSegments = 16;
        QMatrix4x4 invModelView = m_camera->modelViewMatrix().inverted();
        QVector3D rightVec = invModelView.column(0).toVector3D();
        QVector3D upVec = invModelView.column(1).toVector3D();

        QVector<GLfloat> circleFillVertices;
        QVector<GLfloat> circleOutlineVertices;

        for (const auto& centerPnt : unattachedVertices) {
            QVector3D center(centerPnt.X(), centerPnt.Y(), centerPnt.Z());

            circleFillVertices << center.x() << center.y() << center.z(); // Center vertex for fan

            QVector3D firstCircumferencePoint;
            for (int i = 0; i < numSegments; ++i) {
                float angle = (2.0f * M_PI * float(i)) / float(numSegments);
                QVector3D p = center + radius * (cos(angle) * rightVec + sin(angle) * upVec);
                if (i == 0) {
                    firstCircumferencePoint = p;
                }
                circleFillVertices << p.x() << p.y() << p.z();
                circleOutlineVertices << p.x() << p.y() << p.z();
            }
            circleFillVertices << firstCircumferencePoint.x() << firstCircumferencePoint.y() << firstCircumferencePoint.z();
        }

        // Draw fills
        m_shaderProgram->setUniformValue(m_colorLoc, QVector4D(0.2f, 0.3f, 0.3f, 1.0f));
        m_vbo.allocate(circleFillVertices.constData(), circleFillVertices.size() * sizeof(GLfloat));
        const int vertsPerFan = numSegments + 2;
        for (size_t i = 0; i < unattachedVertices.size(); ++i) {
            glDrawArrays(GL_TRIANGLE_FAN, i * vertsPerFan, vertsPerFan);
        }

        // Draw outlines
        m_shaderProgram->setUniformValue(m_colorLoc, QVector4D(1.0f, 1.0f, 1.0f, 1.0f));
        m_vbo.allocate(circleOutlineVertices.constData(), circleOutlineVertices.size() * sizeof(GLfloat));
        const int vertsPerLoop = numSegments;
        for (size_t i = 0; i < unattachedVertices.size(); ++i) {
            glDrawArrays(GL_LINE_LOOP, i * vertsPerLoop, vertsPerLoop);
        }
    }

    glEnable(GL_DEPTH_TEST);
}

void ViewportWidget::initShaders()
{
    m_shaderProgram = new QOpenGLShaderProgram(this);

    if (!m_shaderProgram->addShaderFromSourceFile(QOpenGLShader::Vertex, ":/shaders/simple.vert")) {
        qCritical() << "Vertex shader compilation failed:" << m_shaderProgram->log();
        return;
    }
    if (!m_shaderProgram->addShaderFromSourceFile(QOpenGLShader::Fragment, ":/shaders/simple.frag")) {
        qCritical() << "Fragment shader compilation failed:" << m_shaderProgram->log();
        return;
    }
    if (!m_shaderProgram->link()) {
        qCritical() << "Shader program linking failed:" << m_shaderProgram->log();
        return;
    }

    m_projMatrixLoc = m_shaderProgram->uniformLocation("projectionMatrix");
    m_mvMatrixLoc = m_shaderProgram->uniformLocation("modelViewMatrix");
    m_colorLoc = m_shaderProgram->uniformLocation("objectColor");
}

void ViewportWidget::drawSelectionPlanes()
{
    const float planeSize = 5.0f;
    const float planeOffset = 1.0f;

    glDisable(GL_CULL_FACE);
    glEnable(GL_BLEND);
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);
    glLineWidth(2.0f);
    glDisable(GL_DEPTH_TEST);

    m_vbo.bind();
    QVector<GLfloat> vertices;

    auto drawPlane = [&](const QVector<GLfloat>& quadVerts, bool highlighted) {
        // Draw fill
        if (highlighted) {
            m_shaderProgram->setUniformValue(m_colorLoc, QVector4D(1.0f, 1.0f, 0.0f, 0.5f));
            m_vbo.allocate(quadVerts.constData(), quadVerts.size() * sizeof(GLfloat));
            glDrawArrays(GL_TRIANGLE_FAN, 0, 4);
        }
        // Draw outline
        m_shaderProgram->setUniformValue(m_colorLoc, QVector4D(1.0f, 1.0f, 0.0f, 1.0f));
        m_vbo.allocate(quadVerts.constData(), quadVerts.size() * sizeof(GLfloat));
        glDrawArrays(GL_LINE_LOOP, 0, 4);
    };

    // XY Plane (Top)
    QVector<GLfloat> xyQuad = { planeOffset, planeOffset, 0, planeOffset + planeSize, planeOffset, 0,
                                planeOffset + planeSize, planeOffset + planeSize, 0, planeOffset, planeOffset + planeSize, 0 };
    drawPlane(xyQuad, m_highlightedPlane == SketchPlane::XY);

    // XZ Plane (Front)
    QVector<GLfloat> xzQuad = { planeOffset, 0, planeOffset, planeOffset + planeSize, 0, planeOffset,
                                planeOffset + planeSize, 0, planeOffset + planeSize, planeOffset, 0, planeOffset + planeSize };
    drawPlane(xzQuad, m_highlightedPlane == SketchPlane::XZ);

    // YZ Plane (Right)
    QVector<GLfloat> yzQuad = { 0, planeOffset, planeOffset, 0, planeOffset + planeSize, planeOffset,
                                0, planeOffset + planeSize, planeOffset + planeSize, 0, planeOffset, planeOffset + planeSize };
    drawPlane(yzQuad, m_highlightedPlane == SketchPlane::YZ);

    glEnable(GL_CULL_FACE);
    glEnable(GL_DEPTH_TEST);
    glDisable(GL_BLEND);
}

ViewportWidget::SketchPlane ViewportWidget::checkPlaneSelection(const QPoint& screenPos)
{
    const float planeSize = 5.0f;
    const float planeOffset = 1.0f;

    QVector3D intersection;

    // Check front to back to handle overlaps
    // YZ plane (Right)
    intersection = unproject(screenPos, SketchPlane::YZ);
    if (intersection.y() >= planeOffset && intersection.y() <= planeOffset + planeSize &&
        intersection.z() >= planeOffset && intersection.z() <= planeOffset + planeSize) {
        return SketchPlane::YZ;
    }

    // XZ plane (Front)
    intersection = unproject(screenPos, SketchPlane::XZ);
    if (intersection.x() >= planeOffset && intersection.x() <= planeOffset + planeSize &&
        intersection.z() >= planeOffset && intersection.z() <= planeOffset + planeSize) {
        return SketchPlane::XZ;
    }

    // XY plane (Top)
    intersection = unproject(screenPos, SketchPlane::XY);
    if (intersection.x() >= planeOffset && intersection.x() <= planeOffset + planeSize &&
        intersection.y() >= planeOffset && intersection.y() <= planeOffset + planeSize) {
        return SketchPlane::XY;
    }

    return SketchPlane::NONE;
}