#!/usr/bin/env python3

import sys
import os
import logging
import argparse
import datetime
import time

import pyrogue
import rogue

if '--local' in sys.argv:
    baseDir = os.path.dirname(os.path.realpath(__file__))
    pyrogue.addLibraryPath(f'{baseDir}/../../firmware/common/python')
    pyrogue.addLibraryPath(f'{baseDir}/../../firmware/submodules/surf/python')

import KpixDaq

#rogue.Logging.setFilter('pyrogue.SrpV3', rogue.Logging.Debug)

rootParser = KpixDaq.DesyTrackerRootArgparser()
runParser = argparse.ArgumentParser()


runParser.add_argument(
    '--config', '-c',
    type = str,
    required = True,
    help = 'Configuration yaml file')

runParser.add_argument(
    '--outfile', '-o',
    type = str,
    required = False,
    default = os.path.abspath(datetime.datetime.now().strftime("data/Run_%Y%m%d_%H%M%S.dat")),
    help = 'Output file name')


runParser.add_argument(
    '--runcount', '-r',
    type = int,
    required = False,
    default = 2**31-1,
    help = "Number of acquisitions to take")


if __name__ == "__main__":
    rootArgs, runArgs = rootParser.parse_known_args()
    print(runArgs)
    runArgs = runParser.parse_known_args(runArgs)[0]
    print(runArgs)
    

    with KpixDaq.DesyTrackerRoot(**vars(rootArgs)) as root:
        print('Reading all')
        # Read everything
        root.ReadAll()
        root.waitOnUpdate()

        # Print the version info
        root.DesyTracker.AxiVersion.printStatus()

        if os.path.isdir(runArgs.outfile):
            runArgs.outfile = os.path.abspath(datetime.datetime.now().strftime(f"{runArgs.outfile}/Run_%Y%m%d_%H%M%S.dat"))

        input(f'Data file will be {runArgs.outfile}. \n Hit any key to start run.')
            
        print(f'Opening data file: {runArgs.outfile}')
        root.DataWriter.DataFile.setDisp(runArgs.outfile)
        root.DataWriter.Open()

        print(f"Hard Reset")
        root.HardReset()
        
        print(f"Count Reset")        
        root.CountReset()
            
        print('Writing initial configuration')
        root.LoadConfig(runArgs.config)
        root.waitOnUpdate()
        
        root.ReadAll()
        root.waitOnUpdate()

        root.DesyTrackerRunControl.MaxRunCount.set(runArgs.runcount)

        try:
            root.DesyTrackerRunControl.runState.setDisp('Running')
            root.DesyTrackerRunControl.waitStopped()
            root.DataWriter.Close()            
        except (KeyboardInterrupt):
            root.DesyTrackerRunControl.runState.setDisp('Stopped')
            root.DataWriter.Close()                        
        
