UsdLayerManager/src/ui/Application.cpp

#include "Application.h"
#include "../utils/Logger.h"
#include "../utils/FileDialog.h"
#include "../utils/PathUtils.h"
#include <pxr/usd/usd/prim.h>
#include <pxr/usd/usd/references.h>
#include <pxr/usd/sdf/path.h>
#include <pxr/base/tf/token.h>
#include <imgui.h>
#include <imgui_internal.h>
#include <vector>
#include <string>
#include <filesystem>
#include <algorithm>
#include <cctype>

namespace UsdLayerManager {

// ---------------------------------------------------------------------------
// Internal helper: convert a raw string into a valid USD identifier.
// ---------------------------------------------------------------------------
static std::string SanitizeUsdNameApp(const std::string& raw) {
    std::string result;
    result.reserve(raw.size());
    for (char c : raw) {
        if (std::isalnum(static_cast<unsigned char>(c)) || c == '_')
            result += c;
        else
            result += '_';
    }
    if (result.empty() || std::isdigit(static_cast<unsigned char>(result[0])))
        result = "_" + result;
    return result;
}

Application::Application()
    : m_showDemoWindow(false)
    , m_showStageInfo(true)
    , m_showStageEditor(true)
    , m_showSceneHierarchy(true)
    , m_showViewport(true)
    , m_showPropertyPanel(true)
    , m_showTimeline(true)
    , m_running(false) {
}

Application::~Application() {
    Shutdown();
}

bool Application::Initialize(const std::string& windowTitle, int width, int height) {
    LOG_INFO("Initializing USD Layer Manager Application...");
    
    // Create and initialize ImGui context
    m_imguiContext = std::make_unique<ImGuiContext>();
    if (!m_imguiContext->Initialize(windowTitle, width, height)) {
        LOG_ERROR("Failed to initialize ImGui context");
        return false;
    }
    
    // Create managers
    m_stageManager = std::make_unique<UsdStageManager>();
    m_layerManager = std::make_unique<LayerManager>();
    m_propertyManager = std::make_unique<PropertyManager>();
    m_propertyManager->SetCommandHistory(&m_commandHistory);
    m_stageEditorPanel = std::make_unique<StageEditorPanel>();
    m_stageEditorPanel->SetLayerManager(m_layerManager.get());
    m_stageEditorPanel->SetCommandHistory(&m_commandHistory);
    m_sceneHierarchyPanel = std::make_unique<SceneHierarchyPanel>();
    m_sceneHierarchyPanel->SetPropertyManager(m_propertyManager.get());
    m_sceneHierarchyPanel->SetCommandHistory(&m_commandHistory);
    m_sceneHierarchyPanel->SetLayerManager(m_layerManager.get());
    m_viewportPanel = std::make_unique<ViewportPanel>();
    m_viewportPanel->SetCommandHistory(&m_commandHistory);
    m_propertyPanel = std::make_unique<PropertyPanel>();
    m_propertyPanel->SetPropertyManager(m_propertyManager.get());
    m_propertyPanel->SetCommandHistory(&m_commandHistory);

    m_timelinePanel = std::make_unique<TimelinePanel>();
    m_timelinePanel->OnTimeChanged = [this](pxr::UsdTimeCode displayTime,
                                            pxr::UsdTimeCode editTime) {
        m_viewportPanel->SetTimeCodes(displayTime, editTime);
        m_propertyPanel->SetTimeCodes(displayTime, editTime);
    };

    // Initialize IconManager — must happen after OpenGL context is ready (ImGui init above).
    m_iconManager = std::make_unique<IconManager>();
    m_iconManager->Initialize(ResourcePath("resources/icons"), 24);
    m_sceneHierarchyPanel->SetIconManager(m_iconManager.get());
    m_viewportPanel->SetIconManager(m_iconManager.get());
    m_timelinePanel->SetIconManager(m_iconManager.get());
    m_stageEditorPanel->SetIconManager(m_iconManager.get());

    m_sceneHierarchyPanel->SetOnPrimSelected(
        [this](const std::string& path) {
            m_viewportPanel->SetSelectedPrimPath(path);
            m_propertyPanel->SetSelectedPrimPath(path);
        });

    m_sceneHierarchyPanel->SetOnStageMetadataChanged(
        [this]() {
            RefreshManagers();
        });

    // Single click in viewport → sync hierarchy + property panel
    m_viewportPanel->OnPrimPicked = [this](const std::string& path) {
        m_sceneHierarchyPanel->SetSelectedPath(path);
        m_propertyPanel->SetSelectedPrimPath(path);
    };

    // Rect drag in viewport → sync hierarchy + property panel (primary path)
    m_viewportPanel->OnPrimsPickedRect = [this](const std::vector<std::string>& paths) {
        m_sceneHierarchyPanel->SetSelectedPaths(paths);
        m_propertyPanel->SetSelectedPrimPath(paths.empty() ? "" : paths.front());
    };

    if (!m_stageManager->CreateInMemoryStage()) {
        LOG_ERROR("Failed to create default in-memory stage");
    } else {
        RefreshManagers();
    }
    
    LOG_INFO("Application initialized successfully");
    return true;
}

void Application::Run() {
    LOG_INFO("Starting application main loop...");
    m_running = true;
    
    while (m_running && m_imguiContext->ProcessEvents()) {
        Update();
        RenderUI();
    }
    
    LOG_INFO("Application main loop ended");
}

void Application::Shutdown() {
    m_viewportPanel.reset();
    m_sceneHierarchyPanel.reset();
    m_propertyPanel.reset();
    m_stageEditorPanel.reset();
    m_propertyManager.reset();
    m_layerManager.reset();

    if (m_iconManager) {
        m_iconManager->Shutdown();
        m_iconManager.reset();
    }
    
    if (m_stageManager) {
        m_stageManager->CloseStage();
        m_stageManager.reset();
    }
    
    if (m_imguiContext) {
        LOG_INFO("Shutting down application...");
        m_imguiContext->Shutdown();
        m_imguiContext.reset();
    }
}

void Application::RefreshManagers() {
    m_commandHistory.Clear();
    if (m_stageManager->HasStage()) {
        auto stage = m_stageManager->GetStage();
        m_layerManager->SetStage(stage);
        m_propertyManager->SetStage(stage);
        m_sceneHierarchyPanel->SetStage(stage);
        m_viewportPanel->SetStage(stage);
        m_viewportPanel->FrameScene();
        m_propertyPanel->SetStage(stage);
        m_timelinePanel->SetStage(stage);
    } else {
        m_layerManager->SetStage(nullptr);
        m_propertyManager->SetStage(nullptr);
        m_sceneHierarchyPanel->SetStage(nullptr);
        m_viewportPanel->SetStage(nullptr);
        m_propertyPanel->SetStage(nullptr);
        m_timelinePanel->SetStage(nullptr);
    }
}

void Application::Update() {
    ImGuiIO& io = ImGui::GetIO();
    m_timelinePanel->Update(io.DeltaTime);

    // Process undo/redo hotkeys (Ctrl+Z / Ctrl+Y / Ctrl+Shift+Z).
    // Only fire when no ImGui text-input widget has keyboard focus.
    if (!io.WantTextInput) {
        if (io.KeyCtrl && !io.KeyShift && ImGui::IsKeyPressed(ImGuiKey_Z, false)) {
            m_commandHistory.Undo();
        }
        if (io.KeyCtrl && (ImGui::IsKeyPressed(ImGuiKey_Y, false) ||
                           (io.KeyShift && ImGui::IsKeyPressed(ImGuiKey_Z, false)))) {
            m_commandHistory.Redo();
        }
    }
}

void Application::RenderUI() {
    m_imguiContext->NewFrame();

    // Status bar must be rendered before DockSpaceOverViewport so it
    // reserves space at the bottom before the dockspace claims the rest.
    RenderStatusBar();

    ImGui::DockSpaceOverViewport(0, ImGui::GetMainViewport());

    if (m_showDemoWindow) {
        ImGui::ShowDemoWindow(&m_showDemoWindow);
    }

    RenderMenuBar();

    if (m_showStageInfo && m_stageManager->HasStage()) {
        RenderStageInfo();
    }

    if (m_showStageEditor) {
        ImGui::Begin("Stage Editor", &m_showStageEditor, ImGuiWindowFlags_NoCollapse);
        m_stageEditorPanel->Render();
        ImGui::End();
    }

    if (m_showViewport) {
        m_viewportPanel->Render(&m_showViewport);
    }

    // Scene Hierarchy is rendered AFTER the viewport so that viewport picks
    // (OnPrimPicked / OnPrimsPickedRect) are visible to the hierarchy in the
    // same frame — eliminating the one-frame-late scroll/highlight lag.
    if (m_showSceneHierarchy) {
        ImGui::Begin("Scene Hierarchy", &m_showSceneHierarchy, ImGuiWindowFlags_NoCollapse);
        m_sceneHierarchyPanel->Render();
        ImGui::End();
    }

    if (m_showPropertyPanel) {
        ImGui::Begin("Property Panel", &m_showPropertyPanel, ImGuiWindowFlags_NoCollapse);
        m_propertyPanel->Render();
        ImGui::End();
    }

    if (m_showTimeline) {
        ImGui::Begin("Timeline", &m_showTimeline, ImGuiWindowFlags_NoCollapse);
        m_timelinePanel->Render();
        ImGui::End();
    }

    m_imguiContext->Render();
}

void Application::RenderStatusBar() {
    ImGuiViewport* vp     = ImGui::GetMainViewport();
    float          height = ImGui::GetFrameHeight();

    ImGuiWindowFlags flags =
        ImGuiWindowFlags_NoScrollbar     |
        ImGuiWindowFlags_NoSavedSettings |
        ImGuiWindowFlags_MenuBar;

    ImGui::PushStyleVar(ImGuiStyleVar_WindowPadding, ImVec2(8.f, 2.f));
    bool open = ImGui::BeginViewportSideBar("##statusbar", vp,
                                            ImGuiDir_Down, height, flags);
    ImGui::PopStyleVar();

    if (open) {
        ImGui::BeginMenuBar();

        // FPS
        ImGuiIO& io = ImGui::GetIO();
        char buf[32];
        snprintf(buf, sizeof(buf), "FPS: %.1f", io.Framerate);
        // Right-align: measure text, then position cursor.
        float textW = ImGui::CalcTextSize(buf).x;
        float avail = ImGui::GetContentRegionAvail().x;
        if (avail > textW)
            ImGui::SetCursorPosX(ImGui::GetCursorPosX() + avail - textW);
        ImGui::TextDisabled("%s", buf);

        ImGui::EndMenuBar();
    }
    ImGui::End();
}

void Application::RenderMenuBar() {
    if (ImGui::BeginMainMenuBar()) {
        if (ImGui::BeginMenu("File")) {
            if (ImGui::MenuItem("Open...", "Ctrl+O")) {
                OpenUsdFile();
            }
            if (ImGui::MenuItem("New", "Ctrl+N")) {
                CreateNewUsdFile();
            }
            
            ImGui::Separator();
            
            bool hasStage = m_stageManager->HasStage();
            if (ImGui::MenuItem("Save", "Ctrl+S", false, hasStage)) {
                SaveUsdFile();
            }
            if (ImGui::MenuItem("Save As...", "Ctrl+Shift+S", false, hasStage)) {
                SaveUsdFileAs();
            }
            
            ImGui::Separator();
            
            if (ImGui::MenuItem("Close", nullptr, false, hasStage)) {
                CloseUsdFile();
            }
            
            ImGui::Separator();
            
            if (ImGui::MenuItem("Exit", "Alt+F4")) {
                m_running = false;
            }
            ImGui::EndMenu();
        }

        // Edit menu — Undo / Redo
        {
            bool canUndo = m_commandHistory.CanUndo();
            bool canRedo = m_commandHistory.CanRedo();

            std::string undoLabel = canUndo
                ? ("Undo: " + m_commandHistory.GetUndoDescription())
                : "Undo";
            std::string redoLabel = canRedo
                ? ("Redo: " + m_commandHistory.GetRedoDescription())
                : "Redo";

            if (ImGui::BeginMenu("Edit")) {
                ImGui::BeginDisabled(!canUndo);
                if (ImGui::MenuItem(undoLabel.c_str(), "Ctrl+Z"))
                    m_commandHistory.Undo();
                ImGui::EndDisabled();

                ImGui::BeginDisabled(!canRedo);
                if (ImGui::MenuItem(redoLabel.c_str(), "Ctrl+Y"))
                    m_commandHistory.Redo();
                ImGui::EndDisabled();

                ImGui::EndMenu();
            }
        }

        // Stage editing menu — always available (default stage is always present).
        bool hasStage = m_stageManager->HasStage();
        if (ImGui::BeginMenu("Stage", hasStage)) {
            if (ImGui::MenuItem("Add Reference...")) {
                AddReferenceToStage();
            }

            ImGui::Separator();

            if (ImGui::BeginMenu("Create Prim")) {
                static const char* primTypes[] = {
                    "Xform", "Scope",
                    "Mesh", "Sphere", "Cube", "Cylinder", "Cone", "Capsule",
                    "Camera",
                    "SphereLight", "DomeLight", "RectLight", "DiskLight", "CylinderLight", "DistantLight"
                };
                for (const char* t : primTypes) {
                    if (ImGui::MenuItem(t)) {
                        CreatePrimOnStage(t);
                    }
                }
                ImGui::EndMenu();
            }

            ImGui::EndMenu();
        }
        
        if (ImGui::BeginMenu("View")) {
            ImGui::MenuItem("Stage Editor",     nullptr, &m_showStageEditor);
            ImGui::MenuItem("Scene Hierarchy",  nullptr, &m_showSceneHierarchy);
            ImGui::MenuItem("Viewport",         nullptr, &m_showViewport);
            ImGui::MenuItem("Property Panel",   nullptr, &m_showPropertyPanel);
            ImGui::MenuItem("Timeline",         nullptr, &m_showTimeline);
            ImGui::Separator();
            ImGui::MenuItem("Stage Info",       nullptr, &m_showStageInfo);
            ImGui::MenuItem("Demo Window",      nullptr, &m_showDemoWindow);
            ImGui::EndMenu();
        }
        
        if (ImGui::BeginMenu("Help")) {
            if (ImGui::MenuItem("About")) {
                // Future: Show about dialog
            }
            ImGui::EndMenu();
        }
        
        ImGui::EndMainMenuBar();
    }
}

void Application::RenderStageInfo() {
    ImGui::Begin("Stage Info", &m_showStageInfo);
    
    if (m_stageManager->HasStage()) {
        ImGui::Text("Root Layer:");
        ImGui::SameLine();
        ImGui::TextColored(ImVec4(0.5f, 1.0f, 0.5f, 1.0f), "%s", 
            m_stageManager->GetRootLayerIdentifier().c_str());
        
        std::string realPath = m_stageManager->GetRootLayerPath();
        if (!realPath.empty()) {
            ImGui::Text("Real Path:");
            ImGui::SameLine();
            ImGui::TextWrapped("%s", realPath.c_str());
        }
        
        auto stage = m_stageManager->GetStage();
        if (stage) {
            ImGui::Separator();
            ImGui::Text("Pseudo Root: %s", stage->GetPseudoRoot().GetPath().GetText());
            ImGui::Text("Default Prim: %s", 
                stage->HasDefaultPrim() ? stage->GetDefaultPrim().GetPath().GetText() : "(none)");
        }
    } else {
        ImGui::TextDisabled("No stage loaded");
    }
    
    ImGui::End();
}

void Application::OpenUsdFile() {
    std::string filePath = FileDialog::OpenFile(
        "USD Files (*.usd;*.usda;*.usdc)\0*.usd;*.usda;*.usdc\0All Files (*.*)\0*.*\0",
        "Open USD File",
        m_imguiContext->GetWindowHandle()
    );
    
    if (!filePath.empty()) {
        if (m_stageManager->OpenStage(filePath)) {
            RefreshManagers();
        } else {
            LOG_ERROR("Failed to open USD file: " + m_stageManager->GetLastError());
        }
    }
}

void Application::CreateNewUsdFile() {
    // Create a fresh anonymous in-memory stage — no file path required.
    // The user can save via File > Save As... when they are ready.
    if (m_stageManager->CreateInMemoryStage()) {
        RefreshManagers();
        LOG_INFO("Created new default in-memory stage");
    } else {
        LOG_ERROR("Failed to create new in-memory stage");
    }
}

void Application::SaveUsdFile() {
    if (!m_stageManager->HasStage()) {
        return;
    }

    // If the current stage is in-memory (anonymous root layer), fall through to
    // Save As — UsdStage::Save() cannot write anonymous layers to disk.
    auto rootLayer = m_stageManager->GetRootLayer();
    if (!rootLayer || rootLayer->IsAnonymous()) {
        SaveUsdFileAs();
        return;
    }

    if (!m_stageManager->SaveStage()) {
        LOG_ERROR("Failed to save USD file: " + m_stageManager->GetLastError());
    }
}

void Application::SaveUsdFileAs() {
    if (!m_stageManager->HasStage()) {
        return;
    }

    std::string filePath = FileDialog::SaveFile(
        "USD Files (*.usd;*.usda;*.usdc)\0*.usd;*.usda;*.usdc\0All Files (*.*)\0*.*\0",
        "Save USD File As",
        "usd",
        m_imguiContext->GetWindowHandle()
    );

    if (!filePath.empty()) {
        if (!m_stageManager->SaveStageAs(filePath)) {
            LOG_ERROR("Failed to save USD file: " + m_stageManager->GetLastError());
        } else {
            // SaveStageAs reopens m_stageManager's stage from the new file path.
            // RefreshManagers syncs m_layerManager (and others) to that new stage;
            // without this, subsequent sublayer edits go to the old (now stale) stage
            // and are silently lost on the next save.
            RefreshManagers();
        }
    }
}

void Application::CloseUsdFile() {
    m_stageManager->CloseStage();
    // Re-create a fresh default stage so the app is always in an editable state.
    if (m_stageManager->CreateInMemoryStage()) {
        RefreshManagers();
        LOG_INFO("Closed stage — reset to new default in-memory stage");
    } else {
        RefreshManagers();
        LOG_ERROR("Failed to re-create default stage after close");
    }
}

void Application::AddReferenceToStage() {
    if (!m_stageManager->HasStage()) return;

    std::string filePath = FileDialog::OpenFile(
        "USD Files (*.usd;*.usda;*.usdc;*.usdz)\0*.usd;*.usda;*.usdc;*.usdz\0All Files (*.*)\0*.*\0",
        "Add Reference File",
        m_imguiContext->GetWindowHandle()
    );
    if (filePath.empty()) return;

    auto stage = m_stageManager->GetStage();

    // Derive a valid USD prim name from the file stem.
    std::string stem = std::filesystem::path(filePath).stem().string();
    std::string xformName = SanitizeUsdNameApp(stem);
    if (xformName.empty()) xformName = "Reference";

    // Avoid name collision — append _N if the path already exists.
    std::string finalName = xformName;
    int suffix = 1;
    while (stage->GetPrimAtPath(pxr::SdfPath("/" + finalName)).IsValid()) {
        finalName = xformName + "_" + std::to_string(suffix++);
    }

    try {
        pxr::SdfPath xformPath("/" + finalName);
        pxr::UsdPrim xformPrim = stage->DefinePrim(xformPath, pxr::TfToken("Xform"));
        if (xformPrim.IsValid()) {
            bool ok = xformPrim.GetReferences().AddReference(filePath);
            if (ok) {
                LOG_INFO("Added reference '" + filePath + "' under prim: " + xformPath.GetString());
            } else {
                LOG_ERROR("Failed to add reference '" + filePath + "' to: " + xformPath.GetString());
            }
        } else {
            LOG_ERROR("Failed to define Xform prim: " + xformPath.GetString());
        }
    } catch (const std::exception& e) {
        LOG_ERROR(std::string("Add reference error: ") + e.what());
    }
}

void Application::CreatePrimOnStage(const std::string& typeName) {
    if (!m_stageManager->HasStage()) return;

    auto stage = m_stageManager->GetStage();

    // Auto-generate a unique prim name from the type (e.g. Sphere → /Sphere, /Sphere_1, …).
    std::string baseName = typeName;
    std::string finalName = baseName;
    int suffix = 1;
    while (stage->GetPrimAtPath(pxr::SdfPath("/" + finalName)).IsValid()) {
        finalName = baseName + "_" + std::to_string(suffix++);
    }

    try {
        pxr::SdfPath primPath("/" + finalName);
        pxr::UsdPrim prim = stage->DefinePrim(primPath, pxr::TfToken(typeName));
        if (prim.IsValid()) {
            LOG_INFO("Created prim '" + primPath.GetString() + "' of type " + typeName);
            // Sync selection to the new prim.
            m_sceneHierarchyPanel->SetSelectedPath(primPath.GetString());
            m_propertyPanel->SetSelectedPrimPath(primPath.GetString());
        } else {
            LOG_ERROR("Failed to create prim of type: " + typeName);
        }
    } catch (const std::exception& e) {
        LOG_ERROR(std::string("Create prim error: ") + e.what());
    }
}

} // namespace UsdLayerManager