#include "UsdStageManager.h"
#include "../utils/Logger.h"
#include <pxr/usd/usd/stage.h>
#include <pxr/usd/sdf/layer.h>
#include <pxr/usd/sdf/primSpec.h>
#include <pxr/usd/sdf/copyUtils.h>

PXR_NAMESPACE_USING_DIRECTIVE

namespace UsdLayerManager {

/// Merge all root-prim specs from the session layer into the root layer so
/// they are preserved when the stage is saved.  The session layer is cleared
/// afterwards so that the opinions are not double-applied when the file is
/// reopened.  No-op if the session layer is empty.
static void MergeSessionLayerIntoRoot(UsdStageRefPtr stage) {
    if (!stage) return;

    SdfLayerHandle sessionLayer = stage->GetSessionLayer();
    SdfLayerHandle rootLayer    = stage->GetRootLayer();
    if (!sessionLayer || !rootLayer) return;
    if (sessionLayer->IsEmpty()) return;

    LOG_INFO("Merging session layer into root layer before save");

    // Iterate the immediate children of the pseudo-root in the session layer
    // (these are all root-level prims that have opinions there).
    SdfPrimSpecHandle pseudoRoot = sessionLayer->GetPseudoRoot();
    if (!pseudoRoot) return;

    for (const SdfPrimSpecHandle& primSpec : pseudoRoot->GetNameChildren()) {
        if (!primSpec) continue;
        const SdfPath& primPath = primSpec->GetPath();
        if (!SdfCopySpec(sessionLayer, primPath, rootLayer, primPath)) {
            LOG_ERROR("MergeSessionLayer: SdfCopySpec failed for " + primPath.GetString());
        }
    }

    sessionLayer->Clear();
    LOG_INFO("Session layer merged and cleared");
}

UsdStageManager::UsdStageManager()
    : m_stage(nullptr) {
}

UsdStageManager::~UsdStageManager() {
    CloseStage();
}

bool UsdStageManager::OpenStage(const std::string& filePath) {
    LOG_INFO("Opening USD stage: " + filePath);
    
    try {
        // Close existing stage if any
        CloseStage();
        
        // Open the stage
        m_stage = UsdStage::Open(filePath);
        
        if (!m_stage) {
            SetError("Failed to open USD stage: " + filePath);
            return false;
        }
        
        LOG_INFO("Successfully opened USD stage: " + filePath);
        return true;
        
    } catch (const std::exception& e) {
        SetError(std::string("Exception while opening stage: ") + e.what());
        return false;
    }
}

bool UsdStageManager::CreateNewStage(const std::string& filePath) {
    LOG_INFO("Creating new USD stage: " + filePath);
    
    try {
        // Close existing stage if any
        CloseStage();
        
        // Create new stage
        m_stage = UsdStage::CreateNew(filePath);
        
        if (!m_stage) {
            SetError("Failed to create new USD stage: " + filePath);
            return false;
        }
        
        LOG_INFO("Successfully created new USD stage: " + filePath);
        return true;
        
    } catch (const std::exception& e) {
        SetError(std::string("Exception while creating stage: ") + e.what());
        return false;
    }
}

bool UsdStageManager::CreateInMemoryStage() {
    LOG_INFO("Creating in-memory USD stage");
    
    try {
        // Close existing stage if any
        CloseStage();
        
        // Create in-memory stage
        m_stage = UsdStage::CreateInMemory();
        
        if (!m_stage) {
            SetError("Failed to create in-memory USD stage");
            return false;
        }
        
        LOG_INFO("Successfully created in-memory USD stage");
        return true;
        
    } catch (const std::exception& e) {
        SetError(std::string("Exception while creating in-memory stage: ") + e.what());
        return false;
    }
}

void UsdStageManager::CloseStage() {
    if (m_stage) {
        LOG_INFO("Closing USD stage");
        m_stage.Reset();
        m_stage = nullptr;
    }
}

bool UsdStageManager::SaveStage() {
    if (!m_stage) {
        SetError("No stage to save");
        return false;
    }

    SdfLayerHandle rootLayer = m_stage->GetRootLayer();
    if (!rootLayer || rootLayer->IsAnonymous()) {
        SetError("Stage has no file path — use Save As to save to disk");
        return false;
    }

    try {
        LOG_INFO("Saving USD stage");

        // Persist any session-layer opinions into the root layer so they are
        // not lost when the file is saved (the session layer is anonymous and
        // cannot be saved by UsdStage::Save() itself).
        MergeSessionLayerIntoRoot(m_stage);

        // Explicitly save the root layer first.  UsdStage::Save() iterates
        // the PCP composition stack and saves dirty layers, but sublayer-path
        // edits (InsertSubLayerPath) are root-layer metadata changes that may
        // not be flushed by stage-level Save() in all USD versions.
        rootLayer->Save();

        // Save any other dirty sublayers that have file paths.
        m_stage->Save();
        LOG_INFO("Successfully saved USD stage");
        return true;

    } catch (const std::exception& e) {
        SetError(std::string("Exception while saving stage: ") + e.what());
        return false;
    }
}

bool UsdStageManager::SaveStageAs(const std::string& filePath) {
    if (!m_stage) {
        SetError("No stage to save");
        return false;
    }
    
    try {
        LOG_INFO("Saving USD stage as: " + filePath);

        // Merge session-layer opinions into the root layer so they survive
        // the copy.  The session layer is anonymous and cannot be exported
        // directly; its content must live in the root layer to be persisted.
        MergeSessionLayerIntoRoot(m_stage);

        // Export the root layer to the new path.  Unlike UsdStage::Export()
        // (which flattens all sublayers into one file), SdfLayer::Export()
        // writes only the root layer's own opinions while preserving its
        // subLayerPaths references.  This keeps the sublayer structure intact.
        SdfLayerHandle rootLayer = m_stage->GetRootLayer();
        if (!rootLayer) {
            SetError("No root layer to export");
            return false;
        }

        if (!rootLayer->Export(filePath)) {
            SetError("Failed to export root layer to: " + filePath);
            return false;
        }

        // Reopen the stage from the newly written file so the application
        // reflects the saved path from this point on.
        m_stage = UsdStage::Open(filePath);
        if (!m_stage) {
            SetError("Failed to reopen stage after save as: " + filePath);
            return false;
        }

        LOG_INFO("Successfully saved USD stage as: " + filePath);
        return true;
        
    } catch (const std::exception& e) {
        SetError(std::string("Exception while saving stage as: ") + e.what());
        return false;
    }
}

std::string UsdStageManager::GetRootLayerPath() const {
    if (!m_stage) {
        return "";
    }
    
    SdfLayerHandle rootLayer = m_stage->GetRootLayer();
    if (!rootLayer) {
        return "";
    }
    
    return rootLayer->GetRealPath();
}

std::string UsdStageManager::GetRootLayerIdentifier() const {
    if (!m_stage) {
        return "";
    }
    
    SdfLayerHandle rootLayer = m_stage->GetRootLayer();
    if (!rootLayer) {
        return "";
    }
    
    return rootLayer->GetIdentifier();
}

SdfLayerHandle UsdStageManager::GetRootLayer() const {
    if (!m_stage) {
        return nullptr;
    }
    
    return m_stage->GetRootLayer();
}

SdfLayerHandle UsdStageManager::GetSessionLayer() const {
    if (!m_stage) {
        return nullptr;
    }
    
    return m_stage->GetSessionLayer();
}

void UsdStageManager::SetError(const std::string& error) {
    m_lastError = error;
    LOG_ERROR(error);
}

} // namespace UsdLayerManager
