#include "Application.h"
#include "../utils/Logger.h"
#include "../utils/FileDialog.h"
#include "../utils/PathUtils.h"
#include "../utils/OcioConfigParser.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 <fstream>
#include <filesystem>
#include <algorithm>
#include <cctype>
#include <cstring>
#include "../utils/MovieEncoder.h"

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_showRenderLayerPanel(true)
    , m_showSceneHierarchy(true)
    , m_showViewport(true)
    , m_showPropertyPanel(true)
    , m_showTimeline(true)
    , m_showCurveEditor(false)
    , m_showMaterialEditor(false)
    , 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_renderLayerManager = std::make_unique<RenderLayerManager>();
    m_renderLayerManager->SetLayerManager(m_layerManager.get());
    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_stageEditorPanel->SetRenderLayerManager(m_renderLayerManager.get());
    m_sceneHierarchyPanel = std::make_unique<SceneHierarchyPanel>();
    m_sceneHierarchyPanel->SetPropertyManager(m_propertyManager.get());
    m_sceneHierarchyPanel->SetCommandHistory(&m_commandHistory);
    m_sceneHierarchyPanel->SetLayerManager(m_layerManager.get());
    m_sceneHierarchyPanel->SetRenderLayerManager(m_renderLayerManager.get());
    m_renderLayerPanel = std::make_unique<RenderLayerPanel>();
    m_renderLayerPanel->SetRenderLayerManager(m_renderLayerManager.get());
    m_renderLayerPanel->SetCommandHistory(&m_commandHistory);
    m_renderLayerPanel->SetSceneHierarchyPanel(m_sceneHierarchyPanel.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_curveEditorPanel = std::make_unique<CurveEditorPanel>();
    m_curveEditorPanel->SetCommandHistory(&m_commandHistory);

    m_materialManager = std::make_unique<MaterialManager>();

    m_materialEditorPanel = std::make_unique<MaterialEditorPanel>();
    m_materialEditorPanel->SetMaterialManager(m_materialManager.get());
    m_materialEditorPanel->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);
        m_curveEditorPanel->SetTimeCodes(displayTime, editTime);
    };

    m_curveEditorPanel->OnScrub = [this](double frame) {
        m_timelinePanel->SetCurrentFrameExternal(frame);
    };
    m_timelinePanel->OnBrowseFolder = [this]() -> std::string {
        HWND hwnd = m_imguiContext ? m_imguiContext->GetWindowHandle() : nullptr;
        return FileDialog::BrowseFolder("Select Playblast Output Folder", hwnd);
    };

    // 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_propertyPanel->SetIconManager(m_iconManager.get());
    m_viewportPanel->SetIconManager(m_iconManager.get());
    m_timelinePanel->SetIconManager(m_iconManager.get());
    m_stageEditorPanel->SetIconManager(m_iconManager.get());
    m_materialEditorPanel->SetIconManager(m_iconManager.get());
    m_renderLayerPanel->SetIconManager(m_iconManager.get());

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

    m_sceneHierarchyPanel->SetOnPrimsSelected(
        [this](const std::vector<std::string>& paths) {
            m_viewportPanel->SetSelectedPrimPaths(paths);
        });

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

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

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

    // "Edit Material" in the Property Panel's Material Binding section →
    // open the Material Editor on the resolved material.
    m_propertyPanel->OnEditMaterialRequested = [this](const std::string& materialPath) {
        m_showMaterialEditor = true;
        m_materialEditorPanel->OpenOrCreateMaterial(materialPath);
    };


    if (!m_stageManager->CreateInMemoryStage()) {
        LOG_ERROR("Failed to create default in-memory stage");
    } else {
        RefreshManagers();
    }
    
    // Set OCIO to bundled ACES 1.2 config if not already configured externally.
    // Must happen before the first Render() call so Hydra picks up the env var.
    if (!getenv("OCIO")) {
        namespace fs = std::filesystem;
        std::string ocioConfig = ResourcePath(
            "resources/OpenColorIO-Configs/aces_1.2/config.ocio");
        if (fs::exists(ocioConfig)) {
            _putenv_s("OCIO", ocioConfig.c_str());
            LOG_INFO("OCIO config: " + ocioConfig);
        }
    }

    // NOTE: HDARNOLD_osl_includepath / PXR_MTLX_STDLIB_SEARCH_PATHS are set in
    // main.cpp, NOT here.  They must be in the environment before the plugin
    // DLLs are loaded, because TF_DEFINE_ENV_SETTING caches the value when
    // hdArnold.dll registers its settings.  Setting them at this point is too
    // late and is silently ignored.
    if (const char* oslInc = getenv("HDARNOLD_osl_includepath")) {
        LOG_INFO(std::string("Arnold OSL include path: ") + oslInc);
    } else {
        LOG_WARNING("HDARNOLD_osl_includepath not set — Arnold MaterialX "
                    "shaders will fail to compile ('mx_funcs.h' not found)");
    }

    // Load per-viewport settings and global preferences from AppData.
    if (const char* appData = getenv("APPDATA")) {
        namespace fs = std::filesystem;
        fs::path dir = fs::path(appData) / "UsdLayerManager";
        fs::create_directories(dir);
        m_viewportSettingsPath = (dir / "viewport_settings.ini").string();
        m_prefsPath            = (dir / "preferences.ini").string();
        m_viewportPanel->LoadSettings(m_viewportSettingsPath);
        LoadPreferences();
        ValidateOcioPreferences();
        // Apply pref delegate only to tiles that have no saved delegate.
        m_viewportPanel->ApplyDefaultDelegate(m_prefs.renderDelegate);
        // Color correction is always global — apply to all tiles.
        m_viewportPanel->ApplyGlobalColorCorrection(
            m_prefs.colorCorrectionMode, m_prefs.ocioDisplay,
            m_prefs.ocioView, m_prefs.ocioColorSpace, m_prefs.ocioLook);
        m_materialEditorPanel->SetColorCorrectionFromPrefs(
            m_prefs.colorCorrectionMode, m_prefs.ocioDisplay,
            m_prefs.ocioView, m_prefs.ocioColorSpace, m_prefs.ocioLook);
    }

    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() {
    if (m_viewportPanel && !m_viewportSettingsPath.empty())
        m_viewportPanel->SaveSettings(m_viewportSettingsPath);
    if (!m_prefsPath.empty())
        SavePreferences();

    m_viewportPanel.reset();
    m_sceneHierarchyPanel.reset();
    m_propertyPanel.reset();
    m_stageEditorPanel.reset();
    m_renderLayerPanel.reset();
    // Owns the shader-ball preview's Hydra engine + GL draw target — must be
    // destroyed while the GL context still exists, like m_viewportPanel.
    m_materialEditorPanel.reset();
    m_propertyManager.reset();
    m_renderLayerManager.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);
        // After LayerManager (ActivateRenderLayer routes edit-target changes
        // through it) but before RestorePersistedActiveLayer needs it.
        m_renderLayerManager->SetStage(stage);
        m_propertyManager->SetStage(stage);
        m_sceneHierarchyPanel->SetStage(stage);
        m_renderLayerPanel->SetStage(stage);
        m_viewportPanel->SetStage(stage);
        m_viewportPanel->FrameScene();
        m_propertyPanel->SetStage(stage);
        m_timelinePanel->SetStage(stage);
        m_curveEditorPanel->SetStage(stage);
        m_materialManager->SetStage(stage);
        m_materialEditorPanel->SetStage(stage);
        // Re-applies whichever render layer was active last time this stage
        // was saved (mute state itself isn't serialized — see
        // RenderLayerManager::RestorePersistedActiveLayer).
        m_renderLayerManager->RestorePersistedActiveLayer();
    } else {
        m_layerManager->SetStage(nullptr);
        m_renderLayerManager->SetStage(nullptr);
        m_propertyManager->SetStage(nullptr);
        m_sceneHierarchyPanel->SetStage(nullptr);
        m_renderLayerPanel->SetStage(nullptr);
        m_viewportPanel->SetStage(nullptr);
        m_propertyPanel->SetStage(nullptr);
        m_timelinePanel->SetStage(nullptr);
        m_curveEditorPanel->SetStage(nullptr);
        m_materialManager->SetStage(nullptr);
        m_materialEditorPanel->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_showRenderLayerPanel) {
        ImGui::Begin("Render Layers", &m_showRenderLayerPanel, ImGuiWindowFlags_NoCollapse);
        m_renderLayerPanel->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();
    }

    if (m_showCurveEditor) {
        ImGui::SetNextWindowSize({960, 320}, ImGuiCond_FirstUseEver);
        ImGui::Begin("Curve Editor", &m_showCurveEditor, ImGuiWindowFlags_NoCollapse);
        m_curveEditorPanel->Render();
        ImGui::End();
    }

    if (m_showMaterialEditor) {
        ImGui::SetNextWindowSize({960, 320}, ImGuiCond_FirstUseEver);
        ImGui::Begin("Material Editor", &m_showMaterialEditor, ImGuiWindowFlags_NoCollapse);
        m_materialEditorPanel->Render();
        ImGui::End();
    }

    if (m_showPreferences)
        RenderPreferencesDialog();

    // Capture happens after the timeline so that the "Start" button click in
    // this frame is detected, and after the viewport has rendered the scene.
    if (m_timelinePanel->IsPlayblasting()) {
        CapturePlayblastFrame();
    }

    // If user aborted playblast mid-sequence, close any open encoder.
    if (m_movieEncoder.IsOpen() && !m_timelinePanel->IsPlayblasting()) {
        m_movieEncoder.Close();
    }

    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 — right-aligned
        ImGuiIO& io = ImGui::GetIO();
        char buf[32];
        snprintf(buf, sizeof(buf), "FPS: %.1f", io.Framerate);
        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();
}

// ---------------------------------------------------------------------------
// Global preferences persistence
// ---------------------------------------------------------------------------
void Application::LoadPreferences()
{
    std::ifstream f(m_prefsPath);
    if (!f) return;
    std::string line;
    while (std::getline(f, line)) {
        if (!line.empty() && line.back() == '\r') line.pop_back();
        if (line.empty() || line.front() == '#' || line.front() == '[') continue;
        size_t eq = line.find('=');
        if (eq == std::string::npos) continue;
        std::string key = line.substr(0, eq);
        std::string val = line.substr(eq + 1);
        if      (key == "RenderDelegate")     m_prefs.renderDelegate      = val;
        else if (key == "ColorCorrectionMode")
            m_prefs.colorCorrectionMode = [&]{ try { return std::stoi(val); } catch(...){ return 1; } }();
        else if (key == "OcioDisplay")        m_prefs.ocioDisplay         = val;
        else if (key == "OcioView")           m_prefs.ocioView            = val;
        else if (key == "OcioColorSpace")     m_prefs.ocioColorSpace      = val;
        else if (key == "OcioLook")           m_prefs.ocioLook            = val;
        else if (key == "MaterialBrowserWidth")
            m_prefs.materialBrowserWidth = [&]{ try { return std::stof(val); } catch(...){ return 220.0f; } }();
        else if (key == "MaterialPreviewWidth")
            m_prefs.materialPreviewWidth = [&]{ try { return std::stof(val); } catch(...){ return 320.0f; } }();
        else if (key == "KeepGraphNodeViewSettingsInUsd")
            m_prefs.keepGraphNodeViewSettingsInUsd = (val == "1");
    }

    if (m_materialEditorPanel) {
        m_materialEditorPanel->SetColumnWidths(m_prefs.materialBrowserWidth,
                                               m_prefs.materialPreviewWidth);
        m_materialEditorPanel->SetKeepGraphNodeViewSettingsInUsd(m_prefs.keepGraphNodeViewSettingsInUsd);
    }
}

void Application::ValidateOcioPreferences()
{
    // A persisted display/view that doesn't exist in the active OCIO config
    // makes HdxColorCorrectionTask throw ("Display 'x' not found") every frame
    // and silently skip correction, which reads as a washed-out viewport --
    // worst on delegates like hdEmbree whose output depends on the transform.
    // Note display and view are easy to invert: in the bundled ACES 1.2 config
    // the only display is "ACES" and "sRGB" is one of its views.
    const OcioConfig& cfg = GetCurrentOcioConfig();
    if (!cfg.valid) {
        // No $OCIO / unreadable config -- nothing to validate against. Leave
        // the prefs alone rather than clobbering values that may be correct
        // for a config supplied later.
        return;
    }

    auto contains = [](const std::vector<std::string>& v, const std::string& s) {
        return std::find(v.begin(), v.end(), s) != v.end();
    };

    // --- display ---
    if (m_prefs.ocioDisplay.empty() || !contains(cfg.displays, m_prefs.ocioDisplay)) {
        if (!m_prefs.ocioDisplay.empty()) {
            LOG_WARNING("OCIO display '" + m_prefs.ocioDisplay
                        + "' not in config; falling back to '"
                        + cfg.defaultDisplay + "'");
        }
        m_prefs.ocioDisplay = cfg.defaultDisplay;
    }

    // --- view (must belong to the display resolved above) ---
    auto viewsIt = cfg.views.find(m_prefs.ocioDisplay);
    const std::vector<std::string>* views =
        (viewsIt != cfg.views.end()) ? &viewsIt->second : nullptr;

    if (views && !views->empty()) {
        if (m_prefs.ocioView.empty() || !contains(*views, m_prefs.ocioView)) {
            // Prefer the config default when it's valid for this display,
            // otherwise take the display's first view.
            const std::string fallback =
                contains(*views, cfg.defaultView) ? cfg.defaultView : views->front();
            if (!m_prefs.ocioView.empty()) {
                LOG_WARNING("OCIO view '" + m_prefs.ocioView
                            + "' not valid for display '" + m_prefs.ocioDisplay
                            + "'; falling back to '" + fallback + "'");
            }
            m_prefs.ocioView = fallback;
        }
    }

    // --- color space (same failure mode if it names a missing space) ---
    if (!m_prefs.ocioColorSpace.empty()
        && !contains(cfg.colorSpaces, m_prefs.ocioColorSpace)) {
        LOG_WARNING("OCIO color space '" + m_prefs.ocioColorSpace
                    + "' not in config; clearing it");
        m_prefs.ocioColorSpace.clear();
    }

    // --- look ---
    if (!m_prefs.ocioLook.empty()
        && !contains(cfg.looks, m_prefs.ocioLook)) {
        LOG_WARNING("OCIO look '" + m_prefs.ocioLook
                    + "' not in config; clearing it");
        m_prefs.ocioLook.clear();
    }

    LOG_INFO("OCIO display/view: '" + m_prefs.ocioDisplay + "' / '"
             + m_prefs.ocioView + "'");
}

void Application::SavePreferences()
{
    // Pull the latest splitter positions; called from Shutdown before the
    // panel is reset, and from the preferences dialog while it's alive.
    if (m_materialEditorPanel)
        m_materialEditorPanel->GetColumnWidths(m_prefs.materialBrowserWidth,
                                               m_prefs.materialPreviewWidth);

    std::ofstream f(m_prefsPath);
    if (!f) return;
    f << "[Preferences]\n";
    f << "RenderDelegate="      << m_prefs.renderDelegate      << "\n";
    f << "ColorCorrectionMode=" << m_prefs.colorCorrectionMode << "\n";
    f << "OcioDisplay="         << m_prefs.ocioDisplay         << "\n";
    f << "OcioView="            << m_prefs.ocioView            << "\n";
    f << "OcioColorSpace="      << m_prefs.ocioColorSpace      << "\n";
    f << "OcioLook="            << m_prefs.ocioLook            << "\n";
    f << "MaterialBrowserWidth=" << m_prefs.materialBrowserWidth << "\n";
    f << "MaterialPreviewWidth=" << m_prefs.materialPreviewWidth << "\n";
    f << "KeepGraphNodeViewSettingsInUsd=" << (m_prefs.keepGraphNodeViewSettingsInUsd ? 1 : 0) << "\n";
}

void Application::ApplyPrefsToAllViewports()
{
    if (!m_viewportPanel) return;
    m_viewportPanel->ApplyGlobalColorCorrection(
        m_prefs.colorCorrectionMode,
        m_prefs.ocioDisplay,
        m_prefs.ocioView,
        m_prefs.ocioColorSpace,
        m_prefs.ocioLook);
    if (m_materialEditorPanel)
        m_materialEditorPanel->SetColorCorrectionFromPrefs(
            m_prefs.colorCorrectionMode,
            m_prefs.ocioDisplay,
            m_prefs.ocioView,
            m_prefs.ocioColorSpace,
            m_prefs.ocioLook);
}

void Application::RenderPreferencesDialog()
{
    ImGui::SetNextWindowSize(ImVec2(520, 400), ImGuiCond_FirstUseEver);
    ImGui::SetNextWindowPos(ImGui::GetMainViewport()->GetCenter(),
                            ImGuiCond_FirstUseEver, ImVec2(0.5f, 0.5f));
    if (!ImGui::Begin("Preferences", &m_showPreferences,
                      ImGuiWindowFlags_NoCollapse | ImGuiWindowFlags_NoDocking))
    {
        ImGui::End();
        return;
    }

    bool ccChanged = false;

    if (ImGui::BeginTabBar("PrefTabs")) {

        // ── General ─────────────────────────────────────────────────────────
        if (ImGui::BeginTabItem("General")) {
            ImGui::Spacing();
            ImGui::TextDisabled("No general settings yet.");
            ImGui::EndTabItem();
        }

        // ── Animation ───────────────────────────────────────────────────────
        if (ImGui::BeginTabItem("Animation")) {
            ImGui::Spacing();
            ImGui::TextDisabled("No animation settings yet.");
            ImGui::EndTabItem();
        }

        // ── Material ────────────────────────────────────────────────────────
        if (ImGui::BeginTabItem("Material")) {
            ImGui::Spacing();
            if (ImGui::Checkbox("Keep graph node view settings in USD customData",
                                &m_prefs.keepGraphNodeViewSettingsInUsd)) {
                if (m_materialEditorPanel)
                    m_materialEditorPanel->SetKeepGraphNodeViewSettingsInUsd(
                        m_prefs.keepGraphNodeViewSettingsInUsd);
            }
            ImGui::TextWrapped(
                "When enabled, each node's pin display mode (all pins vs. "
                "connected only) and any manually revealed pins are saved into "
                "the material file and restored when it's reopened. When "
                "disabled, this is a session-only view setting.");
            ImGui::EndTabItem();
        }

        // ── Viewport ────────────────────────────────────────────────────────
        if (ImGui::BeginTabItem("Viewport")) {
            ImGui::Spacing();

            // Render Delegate default — applied to viewports with no saved delegate on startup.
            // Each viewport's toolbar controls its delegate independently after load.
            ImGui::SeparatorText("Default Render Delegate");
            ImGui::TextDisabled("Applied to viewports with no saved setting on startup.");
            ImGui::Spacing();
            auto plugins = UsdSceneRenderer::GetRendererPlugins();
            pxr::TfToken curId(m_prefs.renderDelegate);
            for (const auto& pluginId : plugins) {
                std::string name = UsdSceneRenderer::GetRendererDisplayName(pluginId);
                if (name.empty()) name = pluginId.GetString();
                bool sel = (pluginId == curId);
                if (ImGui::RadioButton(name.c_str(), sel))
                    m_prefs.renderDelegate = pluginId.GetString();
            }

            ImGui::Spacing();

            // Color Correction
            ImGui::SeparatorText("Color Correction");

            auto ccItem = [&](const char* label, int mode) {
                if (ImGui::RadioButton(label, m_prefs.colorCorrectionMode == mode)) {
                    m_prefs.colorCorrectionMode = mode;
                    m_prefOcioFieldsSynced = false;
                    ccChanged = true;
                }
            };
            ccItem("Disabled",    0);
            ccItem("sRGB",        1);
            ccItem("OpenColorIO", 2);

            if (m_prefs.colorCorrectionMode == 2) {
                ImGui::Spacing();
                const OcioConfig& ocfg = GetCurrentOcioConfig();

                if (!m_prefOcioFieldsSynced) {
                    std::string disp = m_prefs.ocioDisplay.empty() ? ocfg.defaultDisplay : m_prefs.ocioDisplay;
                    std::string view = m_prefs.ocioView.empty()    ? ocfg.defaultView    : m_prefs.ocioView;
                    strncpy(m_prefOcioDisplayBuf,    disp.c_str(),                    127);
                    strncpy(m_prefOcioViewBuf,       view.c_str(),                    127);
                    strncpy(m_prefOcioColorSpaceBuf, m_prefs.ocioColorSpace.c_str(), 127);
                    strncpy(m_prefOcioLookBuf,       m_prefs.ocioLook.c_str(),        127);
                    if (m_prefs.ocioDisplay.empty()) m_prefs.ocioDisplay = disp;
                    if (m_prefs.ocioView.empty())    m_prefs.ocioView    = view;
                    m_prefOcioFieldsSynced = true;
                }

                ImGui::PushItemWidth(280.f);

                if (ImGui::BeginCombo("Display##pref", m_prefOcioDisplayBuf)) {
                    for (const auto& d : ocfg.displays) {
                        bool sel = (d == m_prefOcioDisplayBuf);
                        if (ImGui::Selectable(d.c_str(), sel)) {
                            strncpy(m_prefOcioDisplayBuf, d.c_str(), 127);
                            m_prefs.ocioDisplay = d;
                            auto vit = ocfg.views.find(d);
                            if (vit != ocfg.views.end() && !vit->second.empty()) {
                                strncpy(m_prefOcioViewBuf, vit->second[0].c_str(), 127);
                                m_prefs.ocioView = vit->second[0];
                            }
                            ccChanged = true;
                        }
                        if (sel) ImGui::SetItemDefaultFocus();
                    }
                    ImGui::EndCombo();
                }

                {
                    static const std::vector<std::string> kEmpty;
                    auto vit = ocfg.views.find(std::string(m_prefOcioDisplayBuf));
                    const auto& views = (vit != ocfg.views.end()) ? vit->second : kEmpty;
                    if (ImGui::BeginCombo("View##pref", m_prefOcioViewBuf)) {
                        for (const auto& v : views) {
                            bool sel = (v == m_prefOcioViewBuf);
                            if (ImGui::Selectable(v.c_str(), sel)) {
                                strncpy(m_prefOcioViewBuf, v.c_str(), 127);
                                m_prefs.ocioView = v;
                                ccChanged = true;
                            }
                            if (sel) ImGui::SetItemDefaultFocus();
                        }
                        ImGui::EndCombo();
                    }
                }

                if (ImGui::BeginCombo("Color Space##pref",
                        m_prefOcioColorSpaceBuf[0] ? m_prefOcioColorSpaceBuf : "(default)")) {
                    if (ImGui::Selectable("(default)", m_prefOcioColorSpaceBuf[0] == '\0')) {
                        m_prefOcioColorSpaceBuf[0] = '\0';
                        m_prefs.ocioColorSpace.clear();
                        ccChanged = true;
                    }
                    if (m_prefOcioColorSpaceBuf[0] == '\0') ImGui::SetItemDefaultFocus();
                    for (const auto& cs : ocfg.colorSpaces) {
                        bool sel = (cs == m_prefOcioColorSpaceBuf);
                        if (ImGui::Selectable(cs.c_str(), sel)) {
                            strncpy(m_prefOcioColorSpaceBuf, cs.c_str(), 127);
                            m_prefs.ocioColorSpace = cs;
                            ccChanged = true;
                        }
                        if (sel) ImGui::SetItemDefaultFocus();
                    }
                    ImGui::EndCombo();
                }

                if (ImGui::BeginCombo("Look##pref",
                        m_prefOcioLookBuf[0] ? m_prefOcioLookBuf : "(none)")) {
                    if (ImGui::Selectable("(none)", m_prefOcioLookBuf[0] == '\0')) {
                        m_prefOcioLookBuf[0] = '\0';
                        m_prefs.ocioLook.clear();
                        ccChanged = true;
                    }
                    if (m_prefOcioLookBuf[0] == '\0') ImGui::SetItemDefaultFocus();
                    for (const auto& look : ocfg.looks) {
                        bool sel = (look == m_prefOcioLookBuf);
                        if (ImGui::Selectable(look.c_str(), sel)) {
                            strncpy(m_prefOcioLookBuf, look.c_str(), 127);
                            m_prefs.ocioLook = look;
                            ccChanged = true;
                        }
                        if (sel) ImGui::SetItemDefaultFocus();
                    }
                    ImGui::EndCombo();
                }

                ImGui::PopItemWidth();
            }

            ImGui::EndTabItem();
        }

        ImGui::EndTabBar();
    }

    if (ccChanged)
        ApplyPrefsToAllViewports();

    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::Separator();
                if (ImGui::MenuItem("Preferences..."))
                    m_showPreferences = true;

                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("Render Layers",    nullptr, &m_showRenderLayerPanel);
            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::MenuItem("Curve Editor",     nullptr, &m_showCurveEditor);
            ImGui::MenuItem("Material Editor",  nullptr, &m_showMaterialEditor);
            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());
        return;
    }
    SaveDirtyRenderLayers();
}

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 {
            // m_renderLayerManager's layer refs are identifier-keyed SdfLayer
            // objects, independent of which UsdStage currently references
            // them — safe to save before the stage swap RefreshManagers()
            // below performs.
            SaveDirtyRenderLayers();
            // 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::SaveDirtyRenderLayers() {
    if (!m_renderLayerManager) return;
    for (const auto& info : m_renderLayerManager->GetRenderLayers()) {
        if (info.isDefault || !info.layer) continue;
        if (info.layer->IsDirty())
            info.layer->Save();
    }
}

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());
    }
}

// ---------------------------------------------------------------------------
// Writes top-down RGBA8 pixel data as a 24-bit BMP file.
// Negative height in the DIB header tells viewers the rows are top-down.
static bool WriteBMP(const char* path, int w, int h, const uint8_t* rgba) {
    FILE* f = nullptr;
    if (fopen_s(&f, path, "wb") != 0 || !f) return false;

    int rowStride = (w * 3 + 3) & ~3;
    int dataSize  = rowStride * h;
    int fileSize  = 54 + dataSize;

    uint8_t hdr[54] = {};
    hdr[0] = 'B'; hdr[1] = 'M';
    auto put32  = [&](int off, uint32_t v) { memcpy(hdr + off, &v, 4); };
    auto put32s = [&](int off, int32_t  v) { memcpy(hdr + off, &v, 4); };
    auto put16  = [&](int off, uint16_t v) { memcpy(hdr + off, &v, 2); };
    put32(2, uint32_t(fileSize));
    put32(10, 54u);
    put32(14, 40u);
    put32s(18, w);
    put32s(22, -h);   // negative = top-down storage
    put16(26, 1u);
    put16(28, 24u);
    put32(34, uint32_t(dataSize));
    fwrite(hdr, 1, 54, f);

    std::vector<uint8_t> row(rowStride, 0);
    for (int y = 0; y < h; ++y) {
        const uint8_t* src = rgba + y * w * 4;
        for (int x = 0; x < w; ++x) {
            row[x*3+0] = src[x*4+2]; // B
            row[x*3+1] = src[x*4+1]; // G
            row[x*3+2] = src[x*4+0]; // R
        }
        fwrite(row.data(), 1, rowStride, f);
    }
    fclose(f);
    return true;
}

void Application::CapturePlayblastFrame() {
    // The frame BeginPlayblast() is called, the viewport has already rendered
    // at the old time. Skip one frame so the next viewport render picks up the
    // start timecode before we read pixels.
    if (m_timelinePanel->GetPlayblastSkipFirst()) {
        m_timelinePanel->ClearPlayblastSkipFirst();
        return;
    }

    const auto& settings = m_timelinePanel->GetPlayblastSettings();
    auto& renderer = m_viewportPanel->GetRenderer();

    // Re-render at the requested capture resolution. The renderer retains the
    // camera matrices from the viewport render, and SetFraming() adapts the
    // frustum to the new aspect ratio, so the output is correct regardless of
    // whether it matches the display viewport size.
    int cw = settings.captureWidth;
    int ch = settings.captureHeight;
    renderer.Render(cw, ch);

    std::vector<uint8_t> pixels;
    if (!renderer.CaptureFrame(pixels)) {
        LOG_ERROR("Playblast: CaptureFrame failed");
        m_timelinePanel->AbortPlayblast();
        return;
    }

    int idx = m_timelinePanel->GetPlayblastFrameIndex();

    std::error_code ec;
    std::filesystem::create_directories(settings.outputDir, ec);

    // ── Movie encode (direct libav streaming) ─────────────────────────────
    if (settings.exportMovie) {
        if (idx == 0) {
            char moviePath[1024];
            snprintf(moviePath, sizeof(moviePath), "%s\\%s.mp4",
                     settings.outputDir, settings.filePrefix);
            if (!m_movieEncoder.Open(moviePath, cw, ch, settings.fps)) {
                LOG_ERROR("MovieEncoder open failed: " + m_movieEncoder.GetLastError());
                m_timelinePanel->AbortPlayblast();
                return;
            }
        }
        if (!m_movieEncoder.WriteFrame(pixels.data())) {
            LOG_ERROR("MovieEncoder write failed: " + m_movieEncoder.GetLastError());
            m_movieEncoder.Close();
            m_timelinePanel->AbortPlayblast();
            return;
        }
    }

    // ── BMP frames (always written when not exporting movie, or when keepFrames) ─
    if (!settings.exportMovie || settings.keepFrames) {
        char filename[1024];
        snprintf(filename, sizeof(filename), "%s\\%s.%04d.bmp",
                 settings.outputDir, settings.filePrefix, idx);
        if (!WriteBMP(filename, cw, ch, pixels.data())) {
            LOG_ERROR(std::string("Playblast: failed to write ") + filename);
            if (settings.exportMovie) m_movieEncoder.Close();
            m_timelinePanel->AbortPlayblast();
            return;
        }
    }

    if (!m_timelinePanel->AdvancePlayblast()) {
        // Last frame — close the encoder (writes trailer).
        if (settings.exportMovie && m_movieEncoder.IsOpen()) {
            m_movieEncoder.Close();
            LOG_INFO(std::string("Playblast movie: ") +
                     settings.outputDir + "\\" + settings.filePrefix + ".mp4");
        }
        LOG_INFO("Playblast complete: " + std::to_string(idx + 1) +
                 " frames captured to " + settings.outputDir);
    }
}

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