#include "RenderLayerManager.h"
#include "LayerManager.h"
#include "../utils/Logger.h"

#include <pxr/usd/usd/editContext.h>
#include <pxr/usd/usd/editTarget.h>
#include <pxr/usd/usd/prim.h>
#include <pxr/usd/sdf/primSpec.h>
#include <pxr/usd/sdf/relationshipSpec.h>
#include <pxr/usd/sdf/proxyTypes.h>
#include <pxr/usd/usdGeom/imageable.h>
#include <pxr/usd/usdGeom/tokens.h>
#include <pxr/base/vt/dictionary.h>
#include <pxr/base/vt/array.h>
#include <pxr/base/tf/token.h>

#include <algorithm>
#include <cctype>
#include <filesystem>
#include <functional>
#include <set>

PXR_NAMESPACE_USING_DIRECTIVE

namespace UsdLayerManager {

namespace {
// Reserved, undefined ("over"-only) admin prim path each render layer's own
// SdfLayer carries its membership relationship on. Never a `def`, so it never
// shows up in UsdStage::Traverse() / the Scene Hierarchy panel / GetChildren().
const SdfPath kAdminPrimPath("/RenderLayerData");
const TfToken kMembersIncludesName("collection:members:includes");

// SdfLayer custom-layer-data keys.
const std::string kMarkerKey       = "usdLayerManager:renderLayer";
const std::string kNameKey         = "usdLayerManager:renderLayerName";
const std::string kAutoHiddenKey   = "usdLayerManager:autoHiddenPaths";
// Root-layer-only key: persists which render layer was active across save/reload.
const std::string kActiveLayerKey  = "usdLayerManager:activeRenderLayer";

std::string SanitizeFileStem(const std::string& raw) {
    std::string result;
    result.reserve(raw.size());
    for (char c : raw) {
        if (std::isalnum(static_cast<unsigned char>(c)) || c == '_' || c == '-')
            result += c;
        else
            result += '_';
    }
    return result;
}

SdfPath MembersIncludesPath() {
    return kAdminPrimPath.AppendProperty(kMembersIncludesName);
}
} // namespace

RenderLayerManager::RenderLayerManager() {}
RenderLayerManager::~RenderLayerManager() {}

void RenderLayerManager::SetStage(UsdStageRefPtr stage) {
    m_stage = stage;
    Refresh();
}

void RenderLayerManager::Refresh() {
    BuildList();
}

void RenderLayerManager::BuildList() {
    m_layers.clear();
    m_layerRefs.clear();

    RenderLayerInfo def;
    def.name      = "Default";
    def.isDefault = true;
    m_layers.push_back(def);

    if (!m_stage) return;
    SdfLayerHandle rootLayer = m_stage->GetRootLayer();
    if (!rootLayer) return;

    bool anyActive = false;
    for (const auto& subPath : rootLayer->GetSubLayerPaths()) {
        SdfLayerRefPtr sub = SdfLayer::FindOrOpenRelativeToLayer(rootLayer, subPath);
        if (!sub) continue;

        VtDictionary cld = sub->GetCustomLayerData();
        auto markerIt = cld.find(kMarkerKey);
        if (markerIt == cld.end() || !markerIt->second.IsHolding<bool>() ||
            !markerIt->second.Get<bool>())
            continue; // not a render layer sublayer — leave for LayerManager/StageEditorPanel

        m_layerRefs.push_back(sub);

        RenderLayerInfo info;
        info.layer           = SdfLayerHandle(sub);
        info.layerIdentifier = sub->GetIdentifier();
        auto nameIt = cld.find(kNameKey);
        info.name = (nameIt != cld.end() && nameIt->second.IsHolding<std::string>())
            ? nameIt->second.Get<std::string>()
            : LayerManager::ExtractDisplayName(info.layerIdentifier);
        info.isActive = !m_stage->IsLayerMuted(info.layerIdentifier);
        if (info.isActive) anyActive = true;

        m_layers.push_back(info);
    }

    m_layers[0].isActive = !anyActive;
}

/*static*/ bool RenderLayerManager::IsRenderLayerSublayer(const SdfLayerHandle& layer) {
    if (!layer) return false;
    VtDictionary cld = layer->GetCustomLayerData();
    auto it = cld.find(kMarkerKey);
    return it != cld.end() && it->second.IsHolding<bool>() && it->second.Get<bool>();
}

void RenderLayerManager::RefreshLayerManager() {
    if (m_layerManager) m_layerManager->Refresh();
}

SdfLayerHandle RenderLayerManager::FindRenderLayer(const std::string& layerId) const {
    for (const auto& info : m_layers)
        if (!info.isDefault && info.layerIdentifier == layerId)
            return info.layer;
    return SdfLayerHandle();
}

// ---------------------------------------------------------------------------
// CRUD
// ---------------------------------------------------------------------------

std::string RenderLayerManager::CreateRenderLayer(const std::string& name) {
    if (!m_stage) return {};
    SdfLayerHandle rootLayer = m_stage->GetRootLayer();
    if (!rootLayer) return {};

    namespace fs = std::filesystem;
    fs::path baseDir;
    std::string realPath = rootLayer->GetRealPath();
    baseDir = !realPath.empty()
        ? fs::path(realPath).parent_path() / ".renderlayers"
        : fs::current_path() / ".renderlayers";
    std::error_code ec;
    fs::create_directories(baseDir, ec);

    std::string stem = SanitizeFileStem(name.empty() ? "RenderLayer" : name);
    if (stem.empty()) stem = "RenderLayer";
    std::string finalStem = stem;
    fs::path filePath = baseDir / (finalStem + ".usda");
    int suffix = 1;
    while (fs::exists(filePath)) {
        finalStem = stem + "_" + std::to_string(suffix++);
        filePath  = baseDir / (finalStem + ".usda");
    }

    SdfLayerRefPtr newLayer = SdfLayer::CreateNew(filePath.string());
    if (!newLayer) {
        LOG_ERROR("RenderLayerManager: failed to create layer file: " + filePath.string());
        return {};
    }

    VtDictionary cld;
    cld[kMarkerKey] = VtValue(true);
    cld[kNameKey]   = VtValue(name.empty() ? finalStem : name);
    newLayer->SetCustomLayerData(cld);

    SdfPrimSpecHandle adminPrim = SdfCreatePrimInLayer(newLayer, kAdminPrimPath);
    if (adminPrim)
        SdfRelationshipSpec::New(adminPrim, kMembersIncludesName.GetString(),
                                 /*custom=*/false, SdfVariabilityUniform);

    newLayer->Save();

    rootLayer->InsertSubLayerPath(filePath.string(), 0);
    m_stage->MuteLayer(filePath.string()); // starts inactive until explicitly activated

    Refresh();
    RefreshLayerManager();
    return newLayer->GetIdentifier();
}

bool RenderLayerManager::DeleteRenderLayer(const std::string& layerId) {
    if (!m_stage) return false;
    SdfLayerHandle rootLayer = m_stage->GetRootLayer();
    if (!rootLayer) return false;
    if (!FindRenderLayer(layerId)) return false;

    if (GetActiveRenderLayerId() == layerId)
        ActivateDefaultLayer();

    SdfSubLayerProxy subPaths = rootLayer->GetSubLayerPaths();
    for (size_t i = 0; i < subPaths.size(); ++i) {
        SdfLayerRefPtr sub = SdfLayer::FindOrOpenRelativeToLayer(rootLayer, subPaths[i]);
        if (sub && sub->GetIdentifier() == layerId) {
            subPaths.erase(subPaths.begin() + static_cast<long>(i));
            Refresh();
            RefreshLayerManager();
            return true;
        }
    }
    return false;
}

bool RenderLayerManager::RestoreDeletedRenderLayer(const std::string& layerIdentifier) {
    if (!m_stage) return false;
    SdfLayerHandle rootLayer = m_stage->GetRootLayer();
    if (!rootLayer) return false;
    rootLayer->InsertSubLayerPath(layerIdentifier, 0);
    Refresh();
    RefreshLayerManager();
    return true;
}

bool RenderLayerManager::RenameRenderLayer(const std::string& layerId, const std::string& newName) {
    SdfLayerHandle layer = FindRenderLayer(layerId);
    if (!layer || newName.empty()) return false;
    VtDictionary cld = layer->GetCustomLayerData();
    cld[kNameKey] = VtValue(newName);
    layer->SetCustomLayerData(cld);
    Refresh();
    RefreshLayerManager();
    return true;
}

// ---------------------------------------------------------------------------
// Switching
// ---------------------------------------------------------------------------

bool RenderLayerManager::ActivateRenderLayer(const std::string& layerId) {
    if (!m_stage) return false;
    SdfLayerHandle target = FindRenderLayer(layerId);
    if (!target) return false;

    // First entry into render-layer mode this "session" (Default -> a render
    // layer): remember the edit target so ActivateDefaultLayer() can restore
    // it later, no matter how many render-layer-to-render-layer switches
    // happen in between.
    if (IsDefaultActive()) {
        SdfLayerHandle currentEt = m_stage->GetEditTarget().GetLayer();
        m_savedEditTargetIdBeforeMode = currentEt ? currentEt->GetIdentifier() : std::string();
    }

    for (const auto& info : m_layers) {
        if (info.isDefault || info.layerIdentifier == layerId) continue;
        if (!m_stage->IsLayerMuted(info.layerIdentifier))
            m_stage->MuteLayer(info.layerIdentifier);
    }
    if (m_stage->IsLayerMuted(layerId))
        m_stage->UnmuteLayer(layerId);

    if (m_layerManager) m_layerManager->SetEditTarget(layerId);
    else                 m_stage->SetEditTarget(UsdEditTarget(target));

    RecomputeIsolation(target, layerId);
    PersistActiveLayerId(layerId);
    Refresh();
    RefreshLayerManager();
    return true;
}

void RenderLayerManager::ActivateDefaultLayer() {
    if (!m_stage) return;
    for (const auto& info : m_layers) {
        if (info.isDefault) continue;
        if (!m_stage->IsLayerMuted(info.layerIdentifier))
            m_stage->MuteLayer(info.layerIdentifier);
    }

    if (!m_savedEditTargetIdBeforeMode.empty()) {
        if (m_layerManager) {
            m_layerManager->SetEditTarget(m_savedEditTargetIdBeforeMode);
        } else if (SdfLayerHandle saved = SdfLayer::Find(m_savedEditTargetIdBeforeMode)) {
            m_stage->SetEditTarget(UsdEditTarget(saved));
        }
    }
    m_savedEditTargetIdBeforeMode.clear();

    PersistActiveLayerId("");
    Refresh();
    RefreshLayerManager();
}

std::string RenderLayerManager::GetActiveRenderLayerId() const {
    for (const auto& info : m_layers)
        if (!info.isDefault && info.isActive)
            return info.layerIdentifier;
    return {};
}

void RenderLayerManager::PersistActiveLayerId(const std::string& layerId) {
    if (!m_stage) return;
    SdfLayerHandle rootLayer = m_stage->GetRootLayer();
    if (!rootLayer) return;
    VtDictionary cld = rootLayer->GetCustomLayerData();
    if (layerId.empty()) cld.erase(kActiveLayerKey);
    else                 cld[kActiveLayerKey] = VtValue(layerId);
    rootLayer->SetCustomLayerData(cld);
}

void RenderLayerManager::RestorePersistedActiveLayer() {
    if (!m_stage) return;
    SdfLayerHandle rootLayer = m_stage->GetRootLayer();
    if (!rootLayer) return;
    VtDictionary cld = rootLayer->GetCustomLayerData();
    auto it = cld.find(kActiveLayerKey);
    if (it == cld.end() || !it->second.IsHolding<std::string>()) return;
    std::string layerId = it->second.Get<std::string>();
    if (FindRenderLayer(layerId))
        ActivateRenderLayer(layerId);
}

// ---------------------------------------------------------------------------
// Membership (Sdf-level — works regardless of mute state)
// ---------------------------------------------------------------------------

std::vector<SdfPath> RenderLayerManager::GetMembers(const std::string& layerId) const {
    std::vector<SdfPath> result;
    SdfLayerHandle layer = FindRenderLayer(layerId);
    if (!layer) return result;
    SdfRelationshipSpecHandle rel = layer->GetRelationshipAtPath(MembersIncludesPath());
    if (!rel) return result;
    for (const SdfPath& p : rel->GetTargetPathList().GetAddedOrExplicitItems())
        result.push_back(p);
    return result;
}

bool RenderLayerManager::AddMember(const std::string& layerId, const SdfPath& path) {
    SdfLayerHandle layer = FindRenderLayer(layerId);
    if (!layer) return false;

    SdfRelationshipSpecHandle rel = layer->GetRelationshipAtPath(MembersIncludesPath());
    if (!rel) {
        SdfPrimSpecHandle adminPrim = SdfCreatePrimInLayer(layer, kAdminPrimPath);
        if (!adminPrim) return false;
        rel = SdfRelationshipSpec::New(adminPrim, kMembersIncludesName.GetString(),
                                       /*custom=*/false, SdfVariabilityUniform);
    }
    if (!rel) return false;

    // Add() doesn't de-duplicate against existing explicit items — guard
    // here so repeated adds of the same path (e.g. re-running a command, or
    // a redundant drag) don't pile up duplicate targets.
    auto existing = rel->GetTargetPathList().GetAddedOrExplicitItems();
    if (std::find(existing.begin(), existing.end(), path) == existing.end())
        rel->GetTargetPathList().Add(path);

    if (GetActiveRenderLayerId() == layerId)
        RecomputeIsolation(layer, layerId);
    RefreshLayerManager(); // keeps StageEditorPanel's dirty-flag asterisk current
    return true;
}

bool RenderLayerManager::RemoveMember(const std::string& layerId, const SdfPath& path) {
    SdfLayerHandle layer = FindRenderLayer(layerId);
    if (!layer) return false;
    SdfRelationshipSpecHandle rel = layer->GetRelationshipAtPath(MembersIncludesPath());
    if (!rel) return false;

    rel->GetTargetPathList().Remove(path);

    if (GetActiveRenderLayerId() == layerId)
        RecomputeIsolation(layer, layerId);
    RefreshLayerManager();
    return true;
}

// ---------------------------------------------------------------------------
// Viewport isolation
// ---------------------------------------------------------------------------

void RenderLayerManager::ClearAutoHiddenPaths(const SdfLayerHandle& layer) {
    if (!layer) return;
    VtDictionary cld = layer->GetCustomLayerData();
    auto it = cld.find(kAutoHiddenKey);
    if (it == cld.end()) return;

    if (it->second.IsHolding<VtArray<std::string>>()) {
        for (const std::string& p : it->second.Get<VtArray<std::string>>()) {
            SdfPath primPath(p);
            SdfPrimSpecHandle owner = layer->GetPrimAtPath(primPath);
            SdfAttributeSpecHandle attrSpec =
                layer->GetAttributeAtPath(primPath.AppendProperty(TfToken("visibility")));
            if (owner && attrSpec)
                owner->RemoveProperty(attrSpec);
        }
    }
    cld.erase(it);
    layer->SetCustomLayerData(cld);
}

void RenderLayerManager::RecomputeIsolation(const SdfLayerHandle& layer, const std::string& layerId) {
    if (!m_stage || !layer) return;

    ClearAutoHiddenPaths(layer);

    std::vector<SdfPath> memberList = GetMembers(layerId);
    std::set<SdfPath> members(memberList.begin(), memberList.end());

    // Every strict prefix of every member path: a subtree rooted at one of
    // these paths might still contain a member somewhere below it, so it
    // must be recursed into rather than hidden outright.
    std::set<SdfPath> ancestors;
    for (const SdfPath& m : memberList) {
        SdfPath p = m.GetParentPath();
        while (!p.IsEmpty()) {
            if (!ancestors.insert(p).second) break; // already inserted by another member
            if (p == SdfPath::AbsoluteRootPath()) break;
            p = p.GetParentPath();
        }
    }

    std::vector<std::string> newHidden;

    std::function<void(const UsdPrim&)> visit = [&](const UsdPrim& prim) {
        for (const UsdPrim& child : prim.GetChildren()) {
            SdfPath cp = child.GetPath();
            if (cp == kAdminPrimPath) continue;

            if (members.count(cp)) {
                // Exact member: whole subtree stays default-visible (expandPrims
                // semantics — everything under a member is implicitly included),
                // nothing to author, no need to recurse.
                continue;
            }
            if (ancestors.count(cp)) {
                visit(child); // a member lies somewhere below — keep exploring
                continue;
            }

            UsdGeomImageable img(child);
            if (img) {
                UsdEditContext ec(m_stage, layer);
                img.GetVisibilityAttr().Set(UsdGeomTokens->invisible);
                newHidden.push_back(cp.GetString());
                // Minimal cut: stop here, children inherit the hidden state.
            } else {
                // Can't author visibility on a non-imageable prim (e.g. a pure
                // Scope holding unrelated data) — look deeper for an imageable
                // descendant to hide instead, rather than silently failing to
                // isolate this whole branch.
                visit(child);
            }
        }
    };
    visit(m_stage->GetPseudoRoot());

    if (!newHidden.empty()) {
        VtDictionary cld = layer->GetCustomLayerData();
        cld[kAutoHiddenKey] = VtValue(VtArray<std::string>(newHidden.begin(), newHidden.end()));
        layer->SetCustomLayerData(cld);
    }
}

} // namespace UsdLayerManager
