#include "GeometryDescriptor.hh"
#include "G4System.hh"

#include <cstdlib>
#include <iostream>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

namespace {

GeometryDescriptor buildPbScint() {
    GeometryDescriptor gd;
    for (int i = 0; i < 60; ++i) {
        gd.addLayer(0.2, "G4_Pb", false);
        gd.addLayer(0.3, "G4_PLASTIC_SC_VINYLTOLUENE", true, 10, 10);
    }
    return gd;
}

GeometryDescriptor buildFeScint() {
    GeometryDescriptor gd;
    for (int i = 0; i < 40; ++i) {
        gd.addLayer(1.0, "G4_Fe", false);
        gd.addLayer(0.5, "G4_PLASTIC_SC_VINYLTOLUENE", true, 10, 10);
    }
    return gd;
}

GeometryDescriptor buildWScintEcal() {
    GeometryDescriptor gd;
    // Thin homogeneous preshower: ~0.05 X0, sees MIPs/shower-start, not containment.
    gd.addLayer(2.0, "G4_PLASTIC_SC_VINYLTOLUENE", true, 10, 10);
    for (int i = 0; i < 75; ++i) {
        gd.addLayer(0.1, "G4_W", false);
        gd.addLayer(0.2, "G4_PLASTIC_SC_VINYLTOLUENE", true, 10, 10);
    }
    return gd;
}

GeometryDescriptor buildPbLAr() {
    GeometryDescriptor gd;
    for (int i = 0; i < 70; ++i) {
        gd.addLayer(0.2, "G4_Pb", false);
        gd.addLayer(0.4, "G4_lAr", true, 10, 10);
    }
    return gd;
}

// Ordered (not alphabetical) so configs can also be selected by 1-based index.
const std::vector<std::pair<std::string, GeometryDescriptor (*)()>> kConfigs = {
    {"pb_scint", buildPbScint},
    {"fe_scint", buildFeScint},
    {"w_scint_ecal", buildWScintEcal},
    {"pb_lar", buildPbLAr},
};

void printUsage() {
    std::cerr << "Usage: run_sampling [configName|configIndex] [nEvents] [energy_GeV]" << std::endl;
    std::cerr << "Available configs:" << std::endl;
    for (std::size_t i = 0; i < kConfigs.size(); ++i) {
        std::cerr << "  " << (i + 1) << ": " << kConfigs[i].first << std::endl;
    }
}

}  // namespace

int main(int argc, char** argv) {
    std::string configName = kConfigs.front().first;
    int nEvents = 10;
    double energy_GeV = 1.0;

    if (argc > 4) {
        printUsage();
        return 1;
    }
    if (argc >= 2) {
        configName = argv[1];
    }
    if (argc >= 3) {
        try {
            nEvents = std::stoi(argv[2]);
            if (nEvents <= 0) throw std::invalid_argument("must be positive");
        } catch (const std::exception& e) {
            std::cerr << "Invalid nEvents '" << argv[2] << "': " << e.what() << std::endl;
            return 1;
        }
    }
    if (argc == 4) {
        try {
            energy_GeV = std::stod(argv[3]);
            if (energy_GeV <= 0) throw std::invalid_argument("must be positive");
        } catch (const std::exception& e) {
            std::cerr << "Invalid energy_GeV '" << argv[3] << "': " << e.what() << std::endl;
            return 1;
        }
    }

    GeometryDescriptor (*builder)() = nullptr;
    try {
        std::size_t pos = 0;
        int index = std::stoi(configName, &pos);
        if (pos == configName.size() && index >= 1 && index <= static_cast<int>(kConfigs.size())) {
            configName = kConfigs[index - 1].first;
            builder = kConfigs[index - 1].second;
        }
    } catch (const std::exception&) {
        // Not a valid index; fall through to name lookup below.
    }
    if (builder == nullptr) {
        for (const auto& kv : kConfigs) {
            if (kv.first == configName) {
                builder = kv.second;
                break;
            }
        }
    }
    if (builder == nullptr) {
        std::cerr << "Unknown config '" << configName << "'." << std::endl;
        printUsage();
        return 1;
    }

    int seed = -1;
    if (const char* seedEnv = std::getenv("MINICALOSIM_SEED")) {
        try {
            std::size_t pos = 0;
            std::string seedStr(seedEnv);
            int parsed = std::stoi(seedStr, &pos);
            if (pos == seedStr.size()) seed = parsed;
        } catch (const std::exception&) {
            // Not a valid integer; fall back to time-based seed.
        }
    }

    GeometryDescriptor gd = builder();

    std::string outfile = "sampling_" + configName + "_" + std::to_string(nEvents) + "events_hits.root";

    G4System g4;
    g4.init(gd, seed);
    g4.run_batch(nEvents, {"e-"}, energy_GeV, energy_GeV, outfile);

    std::cout << "Saved " << nEvents << " events to " << outfile << std::endl;
    return 0;
}