#!/usr/bin/env python
# -*- coding: utf-8 -*-

from minicalo import GeometryDescriptor
from minicalo import G4System 
from minicalo_tools import TempFileManager
import uproot
import glob
import pandas as pd
import os

def run_batch(
        geometry: GeometryDescriptor,
        nEvents: int,
        particleSpec ,
        minEnergy_GeV: float,
        maxEnergy_GeV: float,
        temp_filename: str,
        seed: int
    ):

    if maxEnergy_GeV < 0:
        maxEnergy_GeV = minEnergy_GeV
    
    #check if particleSpec is a list of strings or a single string and assert
    if not isinstance(particleSpec, str):
        assert isinstance(particleSpec, list), "particleSpec must be a string or a list of strings"
        for p in particleSpec:
            assert isinstance(p, str), "particleSpec must be a string or a list of strings"
    if isinstance(particleSpec, str):
        particleSpec = [particleSpec]
        
    # filename without file ending(!)
    if not temp_filename.endswith(".root"):
        filename = temp_filename+ ".root"
    else:
        filename = temp_filename

    _G4System = G4System() # create instance of G4System only here
    _G4System.init(geometry, seed)

    try:
        _G4System.run_batch(nEvents, particleSpec, minEnergy_GeV, maxEnergy_GeV, filename)
        file = glob.glob(filename.replace(".root",  "*.root"))
        if file[0] != filename:
            os.replace(file[0], filename)
        print(f"G4Calo: output written to {filename}")
        if len(file)>1:
            raise RuntimeError(f"Expected exactly one ROOT file, got {len(file)} - check multithreading. Matches: {file}")
    except Exception as e:
        #clean up output file 
        os.remove(filename)

    
    # TO FIX: Geant4 adds "t<threadnumber>" to the filename, circumvent this for one thread, but this is not a good solution
    
    
    #rename to original filename
    

    root_path = filename
    
    return {'root_path':root_path, 'geometry':geometry}

if __name__ == '__main__':
    # do stuff
    # get the first two arguments that correspond to input and output file
    import sys
    assert len(sys.argv) == 3, "Please provide input and output file"

    io = TempFileManager.slave_init(sys.argv[1], sys.argv[2])
    input_data = io.read_input()
    output_data = run_batch(**input_data)
    io.dump_output(output_data)