#!/usr/bin/env python

from dm.aps_bss.api.apsBssApi import ApsBssApi
from dm.ds_web_service.api.experimentDsApi import ExperimentDsApi
from dm.ds_web_service.api.userDsApi import UserDsApi
from dm.daq_web_service.api.experimentDaqApi import ExperimentDaqApi
from dm.common.constants import dmProcessingMode

from dm.common.exceptions.invalidRequest import InvalidRequest
from dm.common.exceptions.objectNotFound import ObjectNotFound
from dm.common.utility.configurationManager import ConfigurationManager
from dm.aps_beamline_tools.cli.apsBeamlineCli import ApsBeamlineCli

class UploadCli(ApsBeamlineCli):
    def __init__(self, validArgCount=ApsBeamlineCli.ANY_NUMBER_OF_POSITIONAL_ARGS):
        ApsBeamlineCli.__init__(self, validArgCount)
        configManager = ConfigurationManager.getInstance()
        self.allowedExperimentTypes = configManager.getAllowedExperimentTypes()
        allowedExperimentTypesHelp = ''
	self.defaultExperimentType = None
        if self.allowedExperimentTypes:
            allowedExperimentTypesHelp = ' Allowed types: %s' % self.allowedExperimentTypes
	    self.defaultExperimentType = self.allowedExperimentTypes.split(',')[0]

        self.addOption('', '--experiment', dest='experimentName', help='Experiment name.')
        self.addOption('', '--data-directory', dest='dataDirectory', help='Experiment data directory.')

        # Experiment options.
        expGroup = 'Add/Update Experiment Options'
        self.addOptionGroup(expGroup, prepend=True)
        self.addOptionToGroup(expGroup, '', '--type', dest='typeName', default=self.defaultExperimentType, help='Experiment type name.%s' % allowedExperimentTypesHelp)
        self.addOptionToGroup(expGroup, '', '--description', dest='description', help='Experiment description.')
        self.addOptionToGroup(expGroup, '', '--start-date', dest='startDate', help='Experiment start date in format DD-MMM-YY.')
        self.addOptionToGroup(expGroup, '', '--end-date', dest='endDate', help='Experiment end date in format DD-MMM-YY.')
        self.addOptionToGroup(expGroup, '', '--users', dest='users', help='Comma-separated list of DM usernames to be added to the new experiment as users.')
        self.addOptionToGroup(expGroup, '', '--proposal-id', dest='proposalId', help='Beamline proposal id. If specified, all users listed on the proposal will be added to the new experiment.')
        self.addOptionToGroup(expGroup, '', '--run', dest='runName', help='Run name. If not specified, current run name is assumed for beamline proposal.')

        # Upload Options
        uploadGroup = 'Upload Options'
        self.addOptionGroup(uploadGroup, prepend=True)
        self.addOptionToGroup(uploadGroup, '', '--dest-directory', dest='destDirectory', help='Destination directory relative to experiment root path.')
        self.addOptionToGroup(uploadGroup, '', '--process-hidden', dest='processHidden', action='store_true', default=False, help='Process hidden source files.')
        self.addOptionToGroup(uploadGroup, '', '--reprocess', dest='reprocess', action='store_true', default=False, help='Reprocess source files that are already in storage, even if they have not been modified.')
        self.addOptionToGroup(uploadGroup, '', '--processing-mode', dest='processingMode', default=dmProcessingMode.DM_PROCESSING_MODE_FILES, help='Processing mode can be one of %s (default: %s). In the "%s" mode files are processed individually, while in the "%s" mode processing plugins work on directories (if possible).' % (dmProcessingMode.DM_ALLOWED_PROCESSING_MODE_LIST, dmProcessingMode.DM_PROCESSING_MODE_FILES, dmProcessingMode.DM_PROCESSING_MODE_FILES, dmProcessingMode.DM_PROCESSING_MODE_DIRECTORY))
        self.addOptionToGroup(uploadGroup, '', '--skip-plugins', dest='skipPlugins', help='Comma-separated list of plugins which should not process files.')

    def checkArgs(self):
        if self.options.experimentName is None:
            raise InvalidRequest('Experiment name must be provided.')
        if self.options.dataDirectory is None:
            raise InvalidRequest('Experiment data directory must be provided.')
        if self.getTypeName() and self.allowedExperimentTypes:
            if self.getTypeName() not in self.allowedExperimentTypes.split(','):
                raise InvalidRequest('Experiment type %s is not allowed on this station. Allowed types are: %s.' % (self.getTypeName(), self.allowedExperimentTypes))
        if self.options.processingMode not in dmProcessingMode.DM_ALLOWED_PROCESSING_MODE_LIST:
            raise InvalidRequest('Processing mode must be one of %s.' % dmProcessingMode.DM_ALLOWED_PROCESSING_MODE_LIST)

    def updateDaqInfoFromOptions(self, daqInfo):
        if self.options.reprocess:
            daqInfo['reprocessFiles'] = True
        if self.options.processHidden:
            daqInfo['processHiddenFiles'] = True
        if self.options.skipPlugins:
            daqInfo['skipPlugins'] = self.options.skipPlugins
        daqInfo['processingMode'] = self.options.processingMode
        if self.options.destDirectory:
            daqInfo['destDirectory'] = self.options.destDirectory

    def getExperimentName(self):
        return self.options.experimentName

    def getTypeName(self):
        return self.options.typeName

    def getDescription(self):
        return self.options.description

    def getStartDate(self):
        return self.options.startDate

    def getEndDate(self):
        return self.options.endDate

    def getProposalId(self):
        proposalId = self.options.proposalId
        if proposalId:
            proposalId = int(proposalId)
        return proposalId

    def getUsers(self):
        # Return list of users and beamline managers that can access data
        users = self.options.users
        if users:
            users = users.split(',')
        else:
            users = []
        beamlineManagers = self.beamlineManagers
        if beamlineManagers:
            beamlineManagers = beamlineManagers.split(',')
        else:
            beamlineManagers = []
        # Remove duplicates by converting into set
        return list(set(users+beamlineManagers))

    def addOrUpdateExperiment(self):
        dsExperimentApi = ExperimentDsApi(self.loginUsername, self.loginPassword, self.dsServiceHost, self.dsServicePort, self.serviceProtocol)
        dsUserApi = UserDsApi(self.loginUsername, self.loginPassword, self.dsServiceHost, self.dsServicePort, self.serviceProtocol)

        description = self.getDescription()
        proposalId = self.getProposalId()
        experimenters = []
        if proposalId:
            bssApi = ApsBssApi(loginFile=self.options.bssLoginFile)
            proposal = bssApi.getBeamlineProposal(proposalId=proposalId, runName=self.options.runName)
            experimenters = proposal.get('experimenters', [])
            if not description:
                description = '%s (Proposal id: %s)' % (proposal['title'], proposalId)


        users = self.getUsers()
        pis = []
        for experimenter in experimenters:
            badge = int(experimenter['badge'])
            if not badge:
                #print 'Skipping user %s due to invalid badge.' % lastName 
                continue
            username = 'd%s' % badge

            # Clasify user
            if experimenter.get('piFlag') == 'Y':
                if not pis.count(username):
                    pis.append(username)
                if users.count(username):
                    users.remove(username)
            else:
                if not users.count(username):
                    users.append(username)

        for username in users+pis:
            # Check that user exists
            dsUserApi.getUserByUsername(username)
           

        # Everything looks good, add experiment
        try:
            experiment = dsExperimentApi.getExperimentByName(self.getExperimentName())
            experimentStation = experiment.get('experimentStation')
            stationName = ''
            if experimentStation:
                stationName = experimentStation.get('name')
            if stationName != self.getStationName():
                raise InvalidRequest('Experiment %s already exists for station %s.' % (self.getExperimentName(), stationName))
        except ObjectNotFound, ex:
            experiment = dsExperimentApi.addExperiment(self.getExperimentName(), self.getStationName(), self.getTypeName(), description, self.getStartDate(), self.getEndDate())

        # Add pis.
        experimentUsernameList = experiment.get('experimentUsernameList', [])
        experimentName = experiment['name']
        roleName = 'PI'
        for username in pis:
            if username not in experimentUsernameList: 
                dsUserApi.addUserExperimentRole(username, roleName, experimentName)
        roleName = 'User'
        for username in users:
            if username not in experimentUsernameList: 
                dsUserApi.addUserExperimentRole(username, roleName, experimentName)
        if len(users+pis):
            experiment = dsExperimentApi.getExperimentByName(experimentName)
        return experiment

    def startUpload(self):
        daqExperimentApi = ExperimentDaqApi(self.loginUsername, self.loginPassword, self.daqServiceHost, self.daqServicePort, self.serviceProtocol)
        daqInfo = self.splitArgsIntoDict()
        self.updateDaqInfoFromOptions(daqInfo)
        uploadInfo = daqExperimentApi.upload(self.getExperimentName(), self.getDataDirectory(), daqInfo=daqInfo)
        return uploadInfo

    def runCommand(self):
        self.parseArgs(usage="""
    dm-%s-daq --experiment=EXPERIMENTNAME --data-directory=DATADIRECTORY
        [--dest-directory=DESTDIRECTORY]
        [--reprocess]
        [--process-hidden]
        [--processing-mode=PROCESSINGMODE]
        [--skip-plugins=SKIPPLUGINS]
        [--type=TYPENAME]
        [--description=DESCRIPTION] 
        [--start-date=STARTDATE] 
        [--end-date=ENDDATE]
        [--users=USERS]
        [--proposal-id=PROPOSALID]
        [--run=RUNNAME]
        [key1:value1, key2:value2, ...]

Description:
    Run upload for experiment on station %s. If experiment does not exist, it will be added to the DM database. If list of users or proposal id is specified, this command will also add roles for all users listed on the proposal.
        """ % (self.getStationName().lower(), self.getStationName()))
        self.checkArgs()
        self.checkCredentials()
        experiment = self.addOrUpdateExperiment()
        print 'EXPERIMENT INFO'
        print experiment.getDisplayString(self.getDisplayKeys(), self.getDisplayFormat())

        print 
        uploadInfo = self.startUpload()
        print 'UPLOAD INFO'
        print uploadInfo.getDisplayString(self.getDisplayKeys(), self.getDisplayFormat())


#######################################################################
# Run command.
if __name__ == '__main__':
    cli = UploadCli()
    cli.run()