#  Licensed to the Apache Software Foundation (ASF) under one or more
#  contributor license agreements.  See the NOTICE file distributed with
#  this work for additional information regarding copyright ownership.
#  The ASF licenses this file to You under the Apache License, Version 2.0
#  (the "License"); you may not use this file except in compliance with
#  the License.  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

import logging
import time
import logging
import airavata_sdk.samples.file_utils as fb

from airavata_sdk.clients.keycloak_token_fetcher import Authenticator

from airavata_sdk.clients.api_server_client import APIServerClient
from airavata_sdk.clients.utils.api_server_client_util import APIServerClientUtil

from airavata.model.experiment.ttypes import ExperimentModel, ExperimentType, UserConfigurationDataModel
from airavata.model.scheduling.ttypes import ComputationalResourceSchedulingModel
from airavata.model.data.replica.ttypes import DataProductModel, DataProductType, DataReplicaLocationModel, \
    ReplicaLocationCategory, ReplicaPersistentType

from airavata.model.application.io.ttypes import InputDataObjectType

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)


class DataModelCreationUtil(object):

    def __init__(self, configuration_file_location, username, password, gateway_id, access_token):
        self.authenticator = Authenticator(configuration_file_location)
        if access_token:
            self.token = self.authenticator.get_airavata_authz_token(username=username,
                                                                     token=access_token,
                                                                     gateway_id=gateway_id)
        else:
            self.token = self.authenticator.get_token_and_user_info_password_flow(username=username,
                                                                                  password=password,
                                                                                  gateway_id=gateway_id)
        self.gateway_id = gateway_id
        self.username = username
        self.password = password
        self.api_server_client = APIServerClient(configuration_file_location)
        self.airavata_util = APIServerClientUtil(configuration_file_location, self.username, self.password,
                                                 self.gateway_id, access_token)

    def get_experiment_data_model_for_single_application(self, project_name, application_name, experiment_name,
                                                         description):
        execution_id = self.airavata_util.get_execution_id(application_name)
        project_id = self.airavata_util.get_project_id(project_name)
        experiment = ExperimentModel()
        experiment.experimentName = experiment_name
        experiment.gatewayId = self.gateway_id
        experiment.userName = self.username
        experiment.description = description
        experiment.projectId = project_id
        experiment.experimentType = ExperimentType.SINGLE_APPLICATION
        experiment.executionId = execution_id
        return experiment

    def configure_computation_resource_scheduling(self,
                                                  experiment_model, computation_resource_name,
                                                  group_resource_profile_name,
                                                  storageId,
                                                  node_count, total_cpu_count, queue_name, wall_time_limit,
                                                  experiment_dir_path, auto_schedule=False):
        resource_host_id = self.airavata_util.get_resource_host_id(computation_resource_name)
        groupResourceProfileId = self.airavata_util.get_group_resource_profile_id(group_resource_profile_name)
        computRes = ComputationalResourceSchedulingModel()
        computRes.resourceHostId = resource_host_id
        computRes.nodeCount = node_count
        computRes.totalCPUCount = total_cpu_count
        computRes.queueName = queue_name
        computRes.wallTimeLimit = wall_time_limit

        userConfigData = UserConfigurationDataModel()
        userConfigData.computationalResourceScheduling = computRes

        userConfigData.groupResourceProfileId = groupResourceProfileId
        userConfigData.storageId = storageId

        userConfigData.experimentDataDir = experiment_dir_path
        userConfigData.airavataAutoSchedule = auto_schedule
        experiment_model.userConfigurationData = userConfigData

        return experiment_model

    def register_input_file(self, file_identifier, storage_name, storageId, input_file_name, uploaded_storage_path):
        dataProductModel = DataProductModel()
        dataProductModel.gatewayId = self.gateway_id
        dataProductModel.ownerName = self.username
        dataProductModel.productName = file_identifier
        dataProductModel.dataProductType = DataProductType.FILE

        replicaLocation = DataReplicaLocationModel()
        replicaLocation.storageResourceId = storageId
        replicaLocation.replicaName = "{} gateway data store copy".format(input_file_name)
        replicaLocation.replicaLocationCategory = ReplicaLocationCategory.GATEWAY_DATA_STORE
        replicaLocation.filePath = "file://{}:{}".format(storage_name,
                                                         uploaded_storage_path + input_file_name)
        dataProductModel.replicaLocations = [replicaLocation]

        return self.api_server_client.register_data_product(self.token, dataProductModel)

    def configure_input_and_outputs(self, experiment_model, input_files, application_name, file_mapping={}):
        execution_id = self.airavata_util.get_execution_id(application_name)

        inputs = self.api_server_client.get_application_inputs(self.token, execution_id)

        configured_inputs = []
        if (len(file_mapping.keys()) == 0):
            count = 0
            for obj in inputs:
                if isinstance(inputs[count], InputDataObjectType) and len(input_files) > count:
                    inputs[count].value = input_files[count]
                    count = count + 1
            configured_inputs = inputs
        else:
            for key in file_mapping.keys():
                for input in inputs:
                    if key == input.name:
                        if input.type == 3:
                            input.value = file_mapping[key]
                            configured_inputs.append(input)
                        elif input.type == 4:
                            val = ','.join(file_mapping[key])
                            input.value = val
                            configured_inputs.append(input)

        experiment_model.experimentInputs = configured_inputs

        outputs = self.api_server_client.get_application_outputs(self.token, execution_id)

        experiment_model.experimentOutputs = outputs

        return experiment_model