#include <Geant4/G4Box.hh>
#include <Geant4/G4LogicalVolume.hh>
#include <Geant4/G4Material.hh>
#include <Geant4/G4NistManager.hh>
#include <Geant4/G4PVPlacement.hh>
#include <Geant4/G4RunManager.hh>
#include <Geant4/G4RunManagerFactory.hh>
#include <Geant4/G4SystemOfUnits.hh>
#include <Geant4/G4VPhysicalVolume.hh>

#include "Core/Vector.h"
#include "HEP/Geant/DetectorConstruction.hh"
#include "Math/ContainerBox.h"
#include "Math/Dense.h"
#include "Solid.h"

#include "Scene.h"
#include "PhysicsList.hh"
#include "ActionInitialization.hh"
#include "SimulationContext.h"
#include "HEP/Detectors/DetectorChamber.h"

namespace uLib {
namespace Geant {

class SceneDetectorConstruction : public DetectorConstruction {
public:
  SceneDetectorConstruction(class SceneImpl *owner) : DetectorConstruction("Scene"), m_Owner(owner) {}
  G4VPhysicalVolume *Construct() override;
private:
  class SceneImpl *m_Owner;
};

static void CheckGeant4Environment() {
  static bool checked = false;
  if (checked) return;
  checked = true;
  if (!std::getenv("G4ENSDFSTATEDATA")) {
    std::cerr << "********************************************************" << std::endl;
    std::cerr << " WARNING: Geant4 environment variables are not set!" << std::endl;
    std::cerr << " Please activate the environment before running:" << std::endl;
    std::cerr << "   micromamba activate mutom" << std::endl;
    std::cerr << "********************************************************" << std::endl;
  }
}

class SceneImpl {
public:
  SceneImpl() : m_RunManager(G4RunManagerFactory::CreateRunManager(G4RunManagerType::Default)),
                m_Emitter(nullptr),
                m_InitCalled(false) {
    m_RunManager->SetUserInitialization(new PhysicsList);
  }

  ~SceneImpl() { 
    if (m_RunManager) delete m_RunManager; 
    // m_World deletion is handled in Scene destructor or here
  }

  void Initialize() {
    if (m_InitCalled) return;
    
    m_RunManager->SetUserInitialization(new SceneDetectorConstruction(this));
    m_RunManager->SetUserInitialization(new ActionInitialization(m_Emitter, &m_Context));
    
    m_RunManager->Initialize();
    m_InitCalled = true;
  }

  Vector<Solid *> m_Solids;
  Solid *m_World = nullptr;
  ContainerBox m_WorldBox;
  G4RunManager *m_RunManager;
  EmitterPrimary *m_Emitter;
  SimulationContext m_Context;
  bool m_InitCalled;
};

G4VPhysicalVolume *SceneDetectorConstruction::Construct() {
  return m_Owner->m_World->GetPhysical();
}

Scene::Scene() {
  CheckGeant4Environment();
  d = new SceneImpl();
}

Scene::~Scene() {
    // Delete solids
    for(auto s : d->m_Solids) delete s; 
    delete d;
}

void Scene::AddSolid(Solid *solid, Solid *parent) { 
  d->m_Solids.push_back(solid);
  if (!d->m_World) {
    d->m_World = solid;
  } else {
    solid->SetParent(parent ? parent : d->m_World);
  }
}

const Solid* Scene::GetWorld() const { return d->m_World; }
ContainerBox* Scene::GetWorldBox() const { return &d->m_WorldBox; }

void Scene::ConstructWorldBox(const Vector3f &size, const char *material) {
  d->m_WorldBox.Scale(size);
  d->m_WorldBox.SetPosition(-size/2.0f);

  if (!d->m_World) {
    d->m_World = new Solid("World");
    d->m_World->SetNistMaterial(material);
    AddSolid(d->m_World);
  }

  G4Box *solidWorld = new G4Box("World", 0.5 * size(0), 0.5 * size(1), 0.5 * size(2));
  G4LogicalVolume *logicWorld = new G4LogicalVolume(solidWorld, d->m_World->GetMaterial(), d->m_World->GetName());
  d->m_World->SetLogical(logicWorld);
  
  G4PVPlacement *physWorld = new G4PVPlacement(nullptr, G4ThreeVector(0, 0, 0), logicWorld, d->m_World->GetName(), 0, false, 0, true);
  d->m_World->SetPhysical(physWorld);
}

void Scene::SetEmitter(EmitterPrimary *emitter) { d->m_Emitter = emitter; }
void Scene::Initialize() { d->Initialize(); }
void Scene::SetVerbosity(int level) {
  d->m_Context.verbosity = level;
  if (d->m_RunManager) d->m_RunManager->SetVerboseLevel(level);
}

void Scene::RunSimulation(int nEvents, Vector<GeantEvent> &results) {
  d->Initialize(); // Ensure initialized
  d->m_Context.mode = SimulationMode::DETAILED;
  d->m_Context.outputGeant = &results;
  d->m_Context.outputMuon = nullptr;
  
  d->m_RunManager->BeamOn(nEvents);
}

void Scene::RunDetectorSimulation(int nEvents, Vector<MuonEvent> &results) {
  d->Initialize(); // Ensure initialized
  d->m_Context.mode = SimulationMode::DETECTOR;
  d->m_Context.outputGeant = nullptr;
  d->m_Context.outputMuon = &results;
  
  // Find detector planes
  d->m_Context.detectorPlanes.clear();
  for (Solid* s : d->m_Solids) {
      if (BoxSolid* bs = dynamic_cast<BoxSolid*>(s)) {
          if (DetectorChamber* dc = dynamic_cast<DetectorChamber*>(bs->GetObject())) {
              d->m_Context.detectorPlanes.push_back(dc->GetWorldProjectionPlane());
          }
      }
  }

  d->m_RunManager->BeamOn(nEvents);
}

} // namespace Geant
} // namespace uLib