// Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto.  Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
//

#include <OgnReadPickStateDatabase.h>

#include "PickingNodeCommon.h"

#include <omni/kit/IApp.h>
#include <omni/ui/Workspace.h>

#include <omni/graph/core/PreUsdInclude.h>
#include <pxr/usd/sdf/path.h>
#include <pxr/usd/usd/common.h>
#include <pxr/usd/usd/prim.h>
#include <pxr/usd/usd/relationship.h>
#include <pxr/usd/usdUtils/stageCache.h>
#include <omni/graph/core/PostUsdInclude.h>

namespace omni
{
namespace graph
{
namespace ui
{
class OgnReadPickState
{
public:
    struct InternalState
    {
        carb::events::ISubscriptionPtr pickingSub;
        PickingEventPayloads eventPayloads;
    } m_internalState;

    static void initialize(GraphContextObj const& context, NodeObj const& nodeObj)
    {
        OgnReadPickState& state = OgnReadPickStateDatabase::sm_stateManagerOgnReadPickState
            .getState<OgnReadPickState>(nodeObj.nodeHandle);

        // Subscribe to picking events
        if (omni::kit::IApp* app = carb::getCachedInterface<omni::kit::IApp>())
        {
            state.m_internalState.pickingSub = carb::events::createSubscriptionToPushByType(
                app->getMessageBusEventStream(),
                kPickingEventType,
                [nodeObj](carb::events::IEvent* e)
                {
                    if (e)
                    {
                        OgnReadPickState& state = OgnReadPickStateDatabase::sm_stateManagerOgnReadPickState
                            .getState<OgnReadPickState>(nodeObj.nodeHandle);
                        state.m_internalState.eventPayloads.setPayload(e->payload);
                    }
                }
            );
        }
    }

    static void release(const NodeObj& nodeObj)
    {
        OgnReadPickState& state = OgnReadPickStateDatabase::sm_stateManagerOgnReadPickState
            .getState<OgnReadPickState>(nodeObj.nodeHandle);

        // Unsubscribe from picking events
        if (state.m_internalState.pickingSub.get())
            state.m_internalState.pickingSub.detach()->unsubscribe();
    }

    static bool compute(OgnReadPickStateDatabase& db)
    {
        OgnReadPickState& state = db.internalState<OgnReadPickState>();

        // Get the targeted viewport and gesture
        char const* const viewportWindowName = db.tokenToString(db.inputs.viewport());
        char const* const gestureName = db.tokenToString(db.inputs.gesture());

        if (!omni::ui::Workspace::getWindow(viewportWindowName))
        {
            db.logWarning("Viewport window '%s' not found", viewportWindowName);
        }

        auto const* eventPayloadValuePtr = state.m_internalState.eventPayloads.getPayloadValue(viewportWindowName, gestureName);
        if (eventPayloadValuePtr)
        {
            // Get the picked path and pos for the targeted viewport and gesture
            char const* const pickedPrimPath = eventPayloadValuePtr->pickedPrimPath;
            pxr::GfVec3d const& pickedWorldPos = eventPayloadValuePtr->pickedWorldPos;
        
            // Determine if a tracked prim is picked
            bool isTrackedPrimPicked;

            // First determine if any prim is picked
            bool const isAnyPrimPicked = pickedPrimPath && pickedPrimPath[0] != '\0';
            if (isAnyPrimPicked)
            {
                // If any prim is picked, determine if the picked prim is tracked
                if (db.inputs.usePaths())
                {
                    // Get the list of tracked prims from the path[] input
                    auto& inputPaths = db.inputs.trackedPrimPaths();

                    // If no tracked prims are specified then we consider all prims to be tracked
                    // Else search the list of tracked prims for the picked prim
                    if (inputPaths.empty())
                        isTrackedPrimPicked = true;
                    else
                        isTrackedPrimPicked = std::any_of(inputPaths.begin(), inputPaths.end(),
                            [&db, pickedPrimPath](NameToken const& path) {
                                return (std::strcmp(db.tokenToString(path), pickedPrimPath) == 0);
                            });
                }
                else
                {
                    // Get the list of tracked prims from the bundle input
                    pxr::SdfPathVector inputPaths;

                    long int stageId = db.abi_context().iContext->getStageId(db.abi_context());
                    pxr::UsdStageRefPtr stage = pxr::UsdUtilsStageCache::Get().Find(pxr::UsdStageCache::Id::FromLongInt(stageId));
                    auto nodeObj = db.abi_node();
                    pxr::UsdPrim nodePrim = stage->GetPrimAtPath(pxr::SdfPath(nodeObj.iNode->getPrimPath(nodeObj)));
                    nodePrim.GetRelationship(kTrackedPrimsRelToken).GetTargets(&inputPaths);

                    // If no tracked prims are specified then we consider all prims to be tracked
                    // Else search the list of tracked prims for the picked prim
                    if (inputPaths.empty())
                        isTrackedPrimPicked = true;
                    else
                        isTrackedPrimPicked = std::any_of(inputPaths.begin(), inputPaths.end(),
                            [pickedPrimPath](pxr::SdfPath const& path) {
                                return (std::strcmp(path.GetText(), pickedPrimPath) == 0);
                            });
                }
            }
            else
            {
                // No prim is picked at all, so a tracked prim certainly isn't picked
                isTrackedPrimPicked = false;
            }

            // Set outputs
            db.outputs.pickedPrimPath() = db.stringToToken(pickedPrimPath);
            db.outputs.pickedWorldPos() = pickedWorldPos;
            db.outputs.isTrackedPrimPicked() = isTrackedPrimPicked;
            db.outputs.isValid() = true;
        }
        else
        {
            db.outputs.pickedPrimPath() = db.stringToToken("");
            db.outputs.pickedWorldPos() = {0, 0, 0};
            db.outputs.isTrackedPrimPicked() = false;
            db.outputs.isValid() = false;
        }

        return true;
    }
};

REGISTER_OGN_NODE()
} // ui
} // graph
} // omni
