#include "rive/data_bind/data_context.hpp"
#include "rive/data_bind/data_bind_path.hpp"
#include "rive/file.hpp"
#include "rive/viewmodel/viewmodel_instance.hpp"
#include "rive/viewmodel/viewmodel.hpp"
#include "rive/viewmodel/viewmodel_instance_viewmodel.hpp"
#include "rive/view_model_type.hpp"
#include <algorithm>

using namespace rive;

DataContext::DataContext(rcp<ViewModelInstance> viewModelInstance)
{
  if (viewModelInstance != nullptr)
  {
    m_ViewModelInstances.push_back(std::move(viewModelInstance));
  }
}

DataContext::DataContext(std::vector<rcp<ViewModelInstance>> viewModelInstances)
    : m_ViewModelInstances(std::move(viewModelInstances))
{}

DataContext::~DataContext()
{
  // Defensive: detach any still-registered containers so a surviving instance
  // (e.g. one a user still holds) can't keep a stale pointer to a container
  // that is going away with this context.
  for (auto& instance : m_ViewModelInstances)
  {
    detachContainers(instance.get());
  }
}

void DataContext::attachContainers(ViewModelInstance* instance)
{
  if (instance == nullptr)
  {
    return;
  }
  for (auto* container : m_dependentContainers)
  {
    instance->addDependent(container);
  }
}

void DataContext::detachContainers(ViewModelInstance* instance)
{
  if (instance == nullptr)
  {
    return;
  }
  for (auto* container : m_dependentContainers)
  {
    instance->removeDependent(container);
  }
}

void DataContext::addDependentContainer(DataBindContainer* container)
{
  if (container == nullptr)
  {
    return;
  }
  if (std::find(
          m_dependentContainers.begin(),
          m_dependentContainers.end(),
          container) != m_dependentContainers.end())
  {
    return;
  }
  m_dependentContainers.push_back(container);
  for (auto& instance : m_ViewModelInstances)
  {
    if (instance != nullptr)
    {
      instance->addDependent(container);
    }
  }
}

void DataContext::removeDependentContainer(DataBindContainer* container)
{
  for (auto& instance : m_ViewModelInstances)
  {
    if (instance != nullptr)
    {
      instance->removeDependent(container);
    }
  }
  m_dependentContainers.erase(
      std::remove(
          m_dependentContainers.begin(),
          m_dependentContainers.end(),
          container),
      m_dependentContainers.end());
}

void DataContext::ensureGlobalSlots()
{
  if (m_globalSlots != nullptr)
  {
    // Already tracking; kept index-aligned by the insert/remove primitives.
    return;
  }
  m_globalSlots = std::make_unique<GlobalSlots>();
  // Everything already in the context was a main before the first global was
  // slotted, so backfill each entry as unslotted.
  m_globalSlots->slotKeys.assign(m_ViewModelInstances.size(), kNoSlot);
}

uint32_t DataContext::slotKeyAt(size_t index) const
{
  if (m_globalSlots != nullptr && index < m_globalSlots->slotKeys.size())
  {
    return m_globalSlots->slotKeys[index];
  }
  // Untracked entries are mains (not on the slot keys).
  return kNoSlot;
}

void DataContext::insertInstanceAt(
    size_t index,
    rcp<ViewModelInstance> value,
    uint32_t slotKey)
{
  auto it = m_ViewModelInstances.insert(
      m_ViewModelInstances.begin() + index,
      std::move(value));
  if (m_globalSlots != nullptr)
  {
    m_globalSlots->slotKeys.insert(
        m_globalSlots->slotKeys.begin() + index,
        slotKey);
  }
  // Keep all attached containers registered as dependents of the new entry.
  attachContainers(it->get());
}

void DataContext::removeInstanceAt(size_t index)
{
  // Detach before erasing so no container stays registered on the removed
  // instance (which may outlive the context via a user-held reference).
  detachContainers(m_ViewModelInstances[index].get());
  m_ViewModelInstances.erase(m_ViewModelInstances.begin() + index);
  if (m_globalSlots != nullptr)
  {
    m_globalSlots->slotKeys.erase(m_globalSlots->slotKeys.begin() + index);
  }
}

void DataContext::viewModelInstance(rcp<ViewModelInstance> value)
{
  if (m_globalSlots != nullptr)
  {
    setMainViewModelInstance(std::move(value));
    return;
  }
  if (m_ViewModelInstances.empty())
  {
    m_ViewModelInstances.push_back(std::move(value));
    attachContainers(m_ViewModelInstances.back().get());
  }
  else
  {
    detachContainers(m_ViewModelInstances[0].get());
    m_ViewModelInstances[0] = std::move(value);
    attachContainers(m_ViewModelInstances[0].get());
  }
  notifyMainViewModelInstanceChanged();
}

void DataContext::notifyMainViewModelInstanceChanged()
{
  // Snapshot: a container may re-register while it resyncs.
  auto containers = m_dependentContainers;
  for (auto* container : containers)
  {
    container->mainViewModelInstanceChanged();
  }
}

void DataContext::setViewModelInstanceForSlot(
    uint32_t slotKey,
    rcp<ViewModelInstance> value)
{
  if (value == nullptr)
  {
    // A null instance empties the slot: drop any occupant so the slot reads
    // as unset. With no global slots tracked there is nothing to clear.
    if (m_globalSlots == nullptr)
    {
      return;
    }
    for (size_t i = 0; i < m_ViewModelInstances.size(); i++)
    {
      if (slotKeyAt(i) == slotKey)
      {
        removeInstanceAt(i);
        return;
      }
    }
    return;
  }
  // Holding a slot-addressed instance makes this a global-bearing context.
  ensureGlobalSlots();
  // Replace in place if the slot is already occupied (keeps its position).
  for (size_t i = 0; i < m_ViewModelInstances.size(); i++)
  {
    if (slotKeyAt(i) == slotKey)
    {
      detachContainers(m_ViewModelInstances[i].get());
      m_ViewModelInstances[i] = std::move(value);
      m_globalSlots->slotKeys[i] = slotKey;
      attachContainers(m_ViewModelInstances[i].get());
      return;
    }
  }
  // Insert at canonical position: keep any leading main (unslotted) instance
  // ahead of the globals, then order among globals by ascending slot key. The
  // instance's own viewModelId is irrelevant here — placement follows
  // slotKey.
  size_t index = 0;
  while (index < m_ViewModelInstances.size() && slotKeyAt(index) == kNoSlot)
  {
    index++;
  }
  while (index < m_ViewModelInstances.size() && slotKeyAt(index) != kNoSlot &&
         slotKeyAt(index) < slotKey)
  {
    index++;
  }
  insertInstanceAt(index, std::move(value), slotKey);
}

rcp<ViewModelInstance> DataContext::instanceForSlot(uint32_t slotKey) const
{
  for (size_t i = 0; i < m_ViewModelInstances.size(); i++)
  {
    if (slotKeyAt(i) == slotKey)
    {
      return m_ViewModelInstances[i];
    }
  }
  return nullptr;
}

void DataContext::removeMainViewModelInstance()
{
  removeMainViewModelInstanceSilently();
  notifyMainViewModelInstanceChanged();
}

void DataContext::removeMainViewModelInstanceSilently()
{
  for (size_t i = 0; i < m_ViewModelInstances.size();)
  {
    if (slotKeyAt(i) == kNoSlot)
    {
      removeInstanceAt(i);
    }
    else
    {
      i++;
    }
  }
}

void DataContext::setMainViewModelInstance(rcp<ViewModelInstance> value)
{
  // Drop any existing main (unslotted) entries, then place the new main at
  // the front so resolution order is [main, globals...]. One notification
  // for the whole swap.
  removeMainViewModelInstanceSilently();
  if (value != nullptr)
  {
    // Main is not on the slot keys.
    insertInstanceAt(0, std::move(value), kNoSlot);
  }
  notifyMainViewModelInstanceChanged();
}

rcp<ViewModelInstance> DataContext::mainViewModelInstance() const
{
  for (size_t i = 0; i < m_ViewModelInstances.size(); i++)
  {
    if (slotKeyAt(i) == kNoSlot)
    {
      return m_ViewModelInstances[i];
    }
  }
  return nullptr;
}

// --- Private helpers: try resolution against a single instance ---

ViewModelInstanceValue* DataContext::tryGetViewModelProperty(
    rcp<ViewModelInstance> instance,
    const std::vector<uint32_t>& path) const
{
  if (instance == nullptr || instance->viewModelId() != path[0])
  {
    return nullptr;
  }
  if (path.size() == 1)
  {
    return nullptr;
  }
  rcp<ViewModelInstance> current = instance;
  for (auto it = path.begin() + 1; it != path.end() - 1; it++)
  {
    auto viewModelInstanceValue = current->propertyValue(*it);
    if (viewModelInstanceValue != nullptr &&
        viewModelInstanceValue->is<ViewModelInstanceViewModel>())
    {
      current = viewModelInstanceValue->as<ViewModelInstanceViewModel>()
                    ->referenceViewModelInstance();
      if (current == nullptr)
      {
        return nullptr;
      }
    }
    else
    {
      return nullptr;
    }
  }
  return current->propertyValue(path.back());
}

ViewModelInstanceValue* DataContext::tryGetRelativeViewModelProperty(
    rcp<ViewModelInstance> instance,
    const std::vector<uint32_t>& path,
    DataResolver* resolver) const
{
  if (instance == nullptr)
  {
    return nullptr;
  }
  rcp<ViewModelInstance> current = instance;
  if (path.size() == 1)
  {
    return current->propertyValue(resolver->resolveName(path[0]));
  }
  for (auto it = path.begin(); it != path.end() - 1; it++)
  {
    auto viewModelInstanceValue = current->propertyValue(
        resolver->resolveName(*it));
    if (viewModelInstanceValue != nullptr &&
        viewModelInstanceValue->is<ViewModelInstanceViewModel>())
    {
      current = viewModelInstanceValue->as<ViewModelInstanceViewModel>()
                    ->referenceViewModelInstance();
      if (current == nullptr)
      {
        return nullptr;
      }
    }
    else
    {
      return nullptr;
    }
  }
  return current->propertyValue(resolver->resolveName(path.back()));
}

rcp<ViewModelInstance> DataContext::tryGetViewModelInstance(
    rcp<ViewModelInstance> instance,
    const std::vector<uint32_t>& path) const
{
  if (instance == nullptr || instance->viewModelId() != path[0])
  {
    return nullptr;
  }
  rcp<ViewModelInstance> current = instance;
  for (auto it = path.begin() + 1; it != path.end(); it++)
  {
    auto viewModelInstanceValue = current->propertyValue(*it);
    if (viewModelInstanceValue != nullptr &&
        viewModelInstanceValue->is<ViewModelInstanceViewModel>())
    {
      current = viewModelInstanceValue->as<ViewModelInstanceViewModel>()
                    ->referenceViewModelInstance();
      if (current == nullptr)
      {
        return nullptr;
      }
    }
    else
    {
      return nullptr;
    }
  }
  return current;
}

rcp<ViewModelInstance> DataContext::tryGetRelativeViewModelInstance(
    rcp<ViewModelInstance> instance,
    const std::vector<uint32_t>& path,
    DataResolver* resolver) const
{
  if (instance == nullptr)
  {
    return nullptr;
  }
  rcp<ViewModelInstance> current = instance;
  for (auto it = path.begin(); it != path.end(); it++)
  {
    auto viewModelInstanceValue = current->propertyValue(
        resolver->resolveName(*it));
    if (viewModelInstanceValue != nullptr &&
        viewModelInstanceValue->is<ViewModelInstanceViewModel>())
    {
      current = viewModelInstanceValue->as<ViewModelInstanceViewModel>()
                    ->referenceViewModelInstance();
      if (current == nullptr)
      {
        return nullptr;
      }
    }
    else
    {
      return nullptr;
    }
  }
  return current;
}

// --- Public resolution methods: iterate instances, then fall to parent ---

ViewModelInstanceValue* DataContext::getViewModelProperty(
    const std::vector<uint32_t> path) const
{
  if (path.size() == 0)
  {
    return nullptr;
  }
  for (auto& instance : m_ViewModelInstances)
  {
    auto result = tryGetViewModelProperty(instance, path);
    if (result != nullptr)
    {
      return result;
    }
  }
  if (m_Parent != nullptr)
  {
    return m_Parent->getViewModelProperty(path);
  }
  return nullptr;
}

ViewModelInstanceValue* DataContext::getRelativeViewModelProperty(
    const std::vector<uint32_t> path,
    DataResolver* resolver) const
{
  if (path.size() == 0 || !resolver)
  {
    return nullptr;
  }
  for (auto& instance : m_ViewModelInstances)
  {
    auto result = tryGetRelativeViewModelProperty(instance, path, resolver);
    if (result != nullptr)
    {
      return result;
    }
  }
  if (m_Parent != nullptr)
  {
    return m_Parent->getRelativeViewModelProperty(path, resolver);
  }
  return nullptr;
}

#ifdef WITH_RIVE_EDITOR
ViewModelInstanceValue* DataContext::getEditorViewModelProperty(
    const std::vector<Id>& path)
{
  if (path.size() < 2)
    return nullptr;
  for (auto* scope = this; scope; scope = scope->parent().get())
    for (const auto& instance : scope->viewModelInstances())
    {
      if (!instance || instance->viewModelId() != path.front())
        continue;
      auto* current = instance.get();
      for (size_t i = 1; current && i < path.size(); ++i)
      {
        ViewModelInstanceValue* value = nullptr;
        for (const auto& candidate : current->propertyValues())
          if (candidate->viewModelPropertyId() == path[i])
          {
            value = candidate.get();
            break;
          }
        if (i + 1 == path.size())
        {
          if (value)
            return value;
          break;
        }
        current = value && value->is<ViewModelInstanceViewModel>()
                      ? value->as<ViewModelInstanceViewModel>()
                            ->referenceViewModelInstance()
                            .get()
                      : nullptr;
      }
    }
  return nullptr;
}
#endif

ViewModelInstanceValue* DataContext::getViewModelProperty(
    DataBindPath* dataBindPath)
{
#ifdef WITH_RIVE_EDITOR
  if (dataBindPath->hasEditorPath())
    return getEditorViewModelProperty(dataBindPath->editorPath());
#endif
  if (dataBindPath->isRelative())
  {
    auto file = dataBindPath->file();
    if (file)
    {
      auto resolver = file->dataResolver();
      if (resolver)
      {
        return getRelativeViewModelProperty(
            dataBindPath->resolvedPath(),
            resolver);
      }
    }
  }
  return getViewModelProperty(dataBindPath->path());
}

rcp<ViewModelInstance> DataContext::getViewModelInstance(
    const std::vector<uint32_t> path) const
{
  if (path.size() == 0)
  {
    return nullptr;
  }
  for (auto& instance : m_ViewModelInstances)
  {
    auto result = tryGetViewModelInstance(instance, path);
    if (result != nullptr)
    {
      return result;
    }
  }
  if (m_Parent != nullptr)
  {
    return m_Parent->getViewModelInstance(path);
  }
  return nullptr;
}

rcp<ViewModelInstance> DataContext::getRelativeViewModelInstance(
    const std::vector<uint32_t> path,
    DataResolver* resolver) const
{
  if (path.size() == 0 || !resolver)
  {
    return nullptr;
  }
  for (auto& instance : m_ViewModelInstances)
  {
    auto result = tryGetRelativeViewModelInstance(instance, path, resolver);
    if (result != nullptr)
    {
      return result;
    }
  }
  if (m_Parent != nullptr)
  {
    return m_Parent->getRelativeViewModelInstance(path, resolver);
  }
  return nullptr;
}

rcp<ViewModelInstance> DataContext::getViewModelInstance(
    DataBindPath* dataBindPath) const
{
  if (!dataBindPath)
  {
    return nullptr;
  }
  if (dataBindPath->isRelative())
  {
    auto file = dataBindPath->file();
    if (file)
    {
      auto resolver = file->dataResolver();
      if (resolver)
      {
        return getRelativeViewModelInstance(
            dataBindPath->resolvedPath(),
            resolver);
      }
    }
    return nullptr;
  }
  else
  {
    return getViewModelInstance(dataBindPath->resolvedPath());
  }
}

rcp<ViewModelInstance> DataContext::resolveGlobalViewModel(
    File* file,
    const char* name)
{
  // Validate the name up front (mirrors
  // Artboard::setGlobalViewModelInstance): an unknown name makes viewModelId
  // return viewModelCount() (an invalid slot key), and a non-global name must
  // never resolve — even if something else populated that slot. Only after
  // confirming the id is in range and the referenced view model is global do
  // we attempt slot / identity resolution.
  uint32_t slotKey = file->viewModelId(name);
  if (slotKey >= file->viewModelCount())
  {
    return nullptr;
  }
  ViewModel* globalViewModel = file->viewModel(slotKey);
  if (globalViewModel == nullptr ||
      static_cast<ViewModelType>(globalViewModel->viewModelType()) !=
          ViewModelType::global)
  {
    return nullptr;
  }

  // Fast path: the pure runtime slots globals onto the root (top-most) data
  // context.
  DataContext* root = this;
  while (root->parent() != nullptr)
  {
    root = root->parent().get();
  }
  auto viewModelInstance = root->instanceForSlot(slotKey);
  if (viewModelInstance != nullptr)
  {
    return viewModelInstance;
  }

  // Fallback: the editor (and nested-artboard propagation) build unslotted
  // contexts, so match the global view model against the contexts' instances.
  for (DataContext* ctx = this; ctx != nullptr; ctx = ctx->parent().get())
  {
    for (const auto& instance : ctx->viewModelInstances())
    {
      if (instance != nullptr && instance->viewModel() == globalViewModel)
      {
        return instance;
      }
    }
  }
  return nullptr;
}