#include "vtkObjectsContext.h"
#include "Vtk/Math/vtkContainerBox.h"
#include "Vtk/Math/vtkCylinder.h"
#include "Vtk/Math/vtkAssembly.h"
#include "Vtk/Math/vtkVoxImage.h"

#include "HEP/Detectors/vtkDetectorChamber.h"
#include "HEP/Geant/vtkBoxSolid.h"

#include <vtkAssembly.h>
#include <vtkPropCollection.h>
#include <iostream>
#include <cstring>

namespace uLib {
namespace Vtk {

vtkObjectsContext::vtkObjectsContext(uLib::ObjectsContext *context)
    : m_Context(context), m_Assembly(::vtkAssembly::New()) {
  this->SetProp(m_Assembly);
  if (m_Context) {
    Object::connect(m_Context, &uLib::ObjectsContext::ObjectAdded, this, &vtkObjectsContext::OnObjectAdded);
    Object::connect(m_Context, &uLib::ObjectsContext::ObjectRemoved, this, &vtkObjectsContext::OnObjectRemoved);
    this->Synchronize();
  }
}

vtkObjectsContext::~vtkObjectsContext() {
  for (auto const& [obj, puppet] : m_Puppets) {
    delete puppet;
  }
  m_Assembly->Delete();
}

void vtkObjectsContext::Synchronize() {
  if (!m_Context) return;

  // 1. Identify objects to add and remove
  const auto& objects = m_Context->GetObjects();
  std::map<uLib::Object*, bool> currentObjects;
  for (auto obj : objects) currentObjects[obj] = true;

  // Remove puppets for objects no longer in context
  for (auto it = m_Puppets.begin(); it != m_Puppets.end(); ) {
    if (currentObjects.find(it->first) == currentObjects.end()) {
      it->second->DisconnectRenderer(nullptr); // If we have a ref to a renderer we should disconnect but Puppet doesn't store it easily
      // Actually Puppet::DisconnectRenderer(vtkRenderer*) needs the renderer.
      // For now we just remove from assembly
      if (auto* p3d = vtkProp3D::SafeDownCast(it->second->GetProp()))
          m_Assembly->RemovePart(p3d);
      this->PuppetRemoved(it->second);
      delete it->second;
      it = m_Puppets.erase(it);
    } else {
      ++it;
    }
  }

  // Add puppets for new objects
  for (auto obj : objects) {
    if (m_Puppets.find(obj) == m_Puppets.end()) {
      Puppet* puppet = this->CreatePuppet(obj);
      if (puppet) {
        m_Puppets[obj] = puppet;
        if (auto* p3d = vtkProp3D::SafeDownCast(puppet->GetProp()))
            m_Assembly->AddPart(p3d);
        this->PuppetAdded(puppet);
      }
    }
  }
}

void vtkObjectsContext::OnObjectAdded(uLib::Object* obj) {
  if (!obj) return;
  if (m_Puppets.find(obj) == m_Puppets.end()) {
    Puppet* puppet = this->CreatePuppet(obj);
    if (puppet) {
      m_Puppets[obj] = puppet;
      if (auto* p3d = vtkProp3D::SafeDownCast(puppet->GetProp()))
          m_Assembly->AddPart(p3d);
      this->PuppetAdded(puppet);
    }
  }
}

void vtkObjectsContext::OnObjectRemoved(uLib::Object* obj) {
  if (!obj) return;
    auto it = m_Puppets.find(obj);
    if (it != m_Puppets.end()) {
        // For now we just remove from assembly. 
        // Puppet::DisconnectRenderer(vtkRenderer*) needs the renderer, but we don't have it here easily.
        if (auto* p3d = vtkProp3D::SafeDownCast(it->second->GetProp()))
            m_Assembly->RemovePart(p3d);
        this->PuppetRemoved(it->second);
        delete it->second;
        m_Puppets.erase(it);
    }
}

Puppet* vtkObjectsContext::GetPuppet(uLib::Object* obj) {
  auto it = m_Puppets.find(obj);
  if (it != m_Puppets.end()) return it->second;
  return nullptr;
}

void vtkObjectsContext::Update() {
  for (auto const& [obj, puppet] : m_Puppets) {
    puppet->Update();
  }
}

void vtkObjectsContext::SyncFromVtk() {
  for (auto const& [obj, puppet] : m_Puppets) {
    puppet->SyncFromVtk();
  }
}

Puppet* vtkObjectsContext::CreatePuppet(uLib::Object* obj) {
  if (!obj) return nullptr;

  if (auto* vox = dynamic_cast<uLib::Abstract::VoxImage*>(obj)) {
    return new vtkVoxImage(*vox);
  } else if (auto* box = dynamic_cast<uLib::ContainerBox*>(obj)) {
    return new vtkContainerBox(box);
  } else if (auto* chamber = dynamic_cast<uLib::DetectorChamber*>(obj)) {
    return new vtkDetectorChamber(chamber);
  } else if (auto* cylinder = dynamic_cast<uLib::Cylinder*>(obj)) {
    return new vtkCylinder(cylinder);
  } else if (auto* assembly = dynamic_cast<uLib::Assembly*>(obj)) {
    return new Assembly(assembly);
  } else if (auto* box = dynamic_cast<uLib::Geant::BoxSolid*>(obj)) {
    return new vtkBoxSolid(box);
  }
  
  // Fallback if we don't know the exact class but it might be a context itself
  if (auto subCtx = dynamic_cast<uLib::ObjectsContext*>(obj)) {
      return new vtkObjectsContext(subCtx);
  }

  return nullptr;
}

void vtkObjectsContext::PuppetAdded(Puppet* puppet) {
  ULIB_SIGNAL_EMIT(vtkObjectsContext::PuppetAdded, puppet);
}

void vtkObjectsContext::PuppetRemoved(Puppet* puppet) {
  ULIB_SIGNAL_EMIT(vtkObjectsContext::PuppetRemoved, puppet);
}

} // namespace Vtk
} // namespace uLib