#include "Geant/Solid.h"
#include "Vtk/HEP/Geant/vtkGeantEvent.h"
#include "HEP/Geant/Scene.h"
#include "HEP/Geant/EmitterPrimary.hh"
#include "Vtk/Math/vtkContainerBox.h"
#include "Vtk/Math/vtkDense.h"
#include "Math/Units.h"
#include "Vtk/uLibVtkViewer.h"
#include "Vtk/HEP/Geant/vtkGeantEvent.h"
#include "Vtk/HEP/Geant/vtkEmitterPrimary.h"
#include "Vtk/Math/vtkContainerBox.h"
#include "Vtk/HEP/Detectors/vtkDetectorChamber.h"
#include "Vtk/HEP/Detectors/vtkDetectorChamber.h"

#include <vtkSmartPointer.h>
#include <vtkCallbackCommand.h>
#include <vtkRenderWindowInteractor.h>
#include <vtkProperty.h>
#include <vtkRenderer.h>
#include <vtkRenderWindow.h>

#include <Geant4/Randomize.hh>
#include <Geant4/G4ParticleGun.hh>
#include <Geant4/G4SystemOfUnits.hh>

#include <iostream>

using namespace uLib;

struct AppState {
    Geant::Scene* scene;
    Vtk::Viewer* viewer;
    Vector<Geant::GeantEvent> results;
};

void KeyPressCallbackFunction(vtkObject* caller, long unsigned int eventId, void* clientData, void* callData) {
    auto* interactor = static_cast<vtkRenderWindowInteractor*>(caller);
    auto* state = static_cast<AppState*>(clientData);
    
    std::string key = interactor->GetKeySym();
    if (key == "Return") {
        std::cout << "--> Firing muon from current emitter position..." << std::endl;
        
        // Run one event
        state->scene->RunSimulation(1, state->results);
        
        if (!state->results.empty()) {
            // Get the last event
            Geant::GeantEvent* lastEvent = &state->results.back();
            std::cout << "    Collected event " << lastEvent->GetEventID() 
                      << " with " << lastEvent->Path().size() << " steps." << std::endl;
                      
            // Wrap it for VTK
            Vtk::GeantEvent* vtkEvent = new Vtk::GeantEvent(lastEvent);
            state->viewer->AddProp3D(*vtkEvent);
            
            // Re-render
            state->viewer->GetRenderer()->Render();
            state->viewer->GetRenderWindow()->Render();
        }
        else {
            std::cout << "    No event collected." << std::endl;
        }
    }
}

int main(int argc, char** argv) {
    // 1. Setup Geant4 Scene
    Geant::Scene scene;
    scene.ConstructWorldBox(Vector3f(30_m, 30_m, 30_m), "G4_AIR");

    // Iron Cube
    SmartPointer<ContainerBox> iron_box(new ContainerBox());
    iron_box->Scale(Vector3f(10_m, 10_m, 10_m));
    iron_box->SetPosition(Vector3f(-5_m, -5_m, -5_m));
    
    SmartPointer<Geant::BoxSolid> iron_solid(new Geant::BoxSolid("IronSolid", iron_box));
    iron_solid->Update();

    SmartPointer<Geant::Material> iron_mat(new Geant::Material("IronMat"));
    iron_mat->SetFromNist("G4_Fe");

    SmartPointer<Geant::LogicalVolume> iron_lv(new Geant::LogicalVolume("IronLV"));
    iron_lv->SetSolid(iron_solid);
    iron_lv->SetMaterial(iron_mat);

    SmartPointer<Geant::PhysicalVolume> iron_pv(new Geant::PhysicalVolume("IronPV", iron_lv));
    scene.AddVolume(iron_pv);

    // Top Detector Chamber
    SmartPointer<DetectorChamber> top_chamber_box(new DetectorChamber());
    top_chamber_box->Scale(Vector3f(10_m, 10_m, 40_cm));
    top_chamber_box->SetPosition(Vector3f(-5_m, -5_m, 7_m));

    SmartPointer<Geant::BoxSolid> top_solid(new Geant::BoxSolid("TopSolid", top_chamber_box));
    top_solid->Update();

    SmartPointer<Geant::Material> air_mat(new Geant::Material("AirMat"));
    air_mat->SetFromNist("G4_AIR");

    SmartPointer<Geant::LogicalVolume> top_lv(new Geant::LogicalVolume("TopLV"));
    top_lv->SetSolid(top_solid);
    top_lv->SetMaterial(air_mat);

    SmartPointer<Geant::PhysicalVolume> top_pv(new Geant::PhysicalVolume("TopPV", top_lv));
    scene.AddVolume(top_pv);

    // Bottom Detector Chamber
    SmartPointer<DetectorChamber> bottom_chamber_box(new DetectorChamber());
    bottom_chamber_box->Scale(Vector3f(10_m, 10_m, 40_cm));
    bottom_chamber_box->SetPosition(Vector3f(-5_m, -5_m, -7.1_m));

    SmartPointer<Geant::BoxSolid> bottom_solid(new Geant::BoxSolid("BottomSolid", bottom_chamber_box));
    bottom_solid->Update();

    SmartPointer<Geant::LogicalVolume> bottom_lv(new Geant::LogicalVolume("BottomLV"));
    bottom_lv->SetSolid(bottom_solid);
    bottom_lv->SetMaterial(air_mat);

    SmartPointer<Geant::PhysicalVolume> bottom_pv(new Geant::PhysicalVolume("BottomPV", bottom_lv));
    scene.AddVolume(bottom_pv);

    Geant::EmitterPrimary* emitter = new Geant::EmitterPrimary();
    emitter->SetPosition(Vector3f(0, 0, 14_m));
    scene.SetEmitter(emitter);
    scene.Initialize();

    // 2. Setup VTK Viewer
    Vtk::Viewer viewer;
    viewer.GetRenderer()->SetBackground(0.05, 0.05, 0.1);
    
    // Visualize world box
    Vtk::ContainerBox* vtkWorld = new Vtk::ContainerBox(scene.GetWorldBox());
    vtkWorld->ShowScaleMeasures(true);
    vtkWorld->SetRepresentation(Vtk::Prop3D::Wireframe);
    vtkWorld->SetSelectable(false);
    viewer.AddProp3D(*vtkWorld);
    
    // Visualize iron cube
    Vtk::ContainerBox* vtkIron = new Vtk::ContainerBox(iron_box.Get());
    vtkIron->SetOpacity(0.2);
    vtkIron->SetRepresentation(Vtk::Prop3D::Surface);
    viewer.AddProp3D(*vtkIron);

    // Visualize Top Chamber
    Vtk::DetectorChamber* vtkTop = new Vtk::DetectorChamber(top_chamber_box.Get());
    vtkTop->SetOpacity(0.5);
    vtkTop->SetColor(0.2, 0.8, 0.2);
    vtkTop->SetRepresentation(Vtk::Prop3D::Surface);
    viewer.AddProp3D(*vtkTop);

    // Visualize Bottom Chamber
    Vtk::DetectorChamber* vtkBottom = new Vtk::DetectorChamber(bottom_chamber_box.Get());
    vtkBottom->SetOpacity(0.5);
    vtkBottom->SetColor(0.2, 0.8, 0.2);
    vtkBottom->SetRepresentation(Vtk::Prop3D::Surface);
    viewer.AddProp3D(*vtkBottom);

    // Visualize Emitter
    Vtk::EmitterPrimary* vtkEmitter = new Vtk::EmitterPrimary(*emitter);
    viewer.AddProp3D(*vtkEmitter);
    
    // 3. Event Handling
    AppState state = { &scene, &viewer, {} };
    
    vtkSmartPointer<vtkCallbackCommand> keyCallback = vtkSmartPointer<vtkCallbackCommand>::New();
    keyCallback->SetCallback(KeyPressCallbackFunction);
    keyCallback->SetClientData(&state);
    
    viewer.GetInteractor()->AddObserver(vtkCommand::KeyPressEvent, keyCallback);

    std::cout << "=================================================" << std::endl;
    std::cout << " Geant Muon Interactive Emitter Test" << std::endl;
    std::cout << " Use the Widget to move/rotate the Arrow (Emitter)" << std::endl;
    std::cout << " Press [ENTER] to fire a muon" << std::endl;
    std::cout << " Press [q] to exit" << std::endl;
    std::cout << "=================================================" << std::endl;

    viewer.ZoomAuto();
    viewer.Start();

    return 0;
}