/*
 * Copyright 2013 the original author or authors.
 *
 * Licensed 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.
 */
package org.springframework.yarn.am;

import java.io.IOException;
import java.net.InetSocketAddress;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.security.SecurityUtil;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.security.token.Token;
import org.apache.hadoop.security.token.TokenIdentifier;
import org.apache.hadoop.yarn.api.ApplicationMasterProtocol;
import org.apache.hadoop.yarn.api.protocolrecords.AllocateRequest;
import org.apache.hadoop.yarn.api.protocolrecords.AllocateResponse;
import org.apache.hadoop.yarn.api.protocolrecords.FinishApplicationMasterRequest;
import org.apache.hadoop.yarn.api.protocolrecords.FinishApplicationMasterResponse;
import org.apache.hadoop.yarn.api.protocolrecords.RegisterApplicationMasterRequest;
import org.apache.hadoop.yarn.api.protocolrecords.RegisterApplicationMasterResponse;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.hadoop.yarn.exceptions.YarnException;
import org.apache.hadoop.yarn.security.AMRMTokenIdentifier;
import org.apache.hadoop.yarn.util.Records;
import org.springframework.yarn.rpc.YarnRpcAccessor;
import org.springframework.yarn.rpc.YarnRpcCallback;

/**
 * Template implementation for {@link AppmasterRmOperations} wrapping
 * communication using {@link ApplicationMasterProtocol}. Methods for this
 * template wraps possible exceptions into Spring Dao exception hierarchy.
 *
 * @author Janne Valkealahti
 *
 */
public class AppmasterRmTemplate extends YarnRpcAccessor<ApplicationMasterProtocol>
		implements AppmasterRmOperations {

	private static final Log log = LogFactory.getLog(AppmasterCmTemplate.class);

	/**
	 * Instantiates a new AppmasterRmTemplate.
	 *
	 * @param config the hadoop configuration
	 */
	public AppmasterRmTemplate(Configuration config) {
		super(ApplicationMasterProtocol.class, config);
	}

	@Override
	public RegisterApplicationMasterResponse registerApplicationMaster(final String host,
			final Integer rpcPort, final String trackUrl) {
		return execute(new YarnRpcCallback<RegisterApplicationMasterResponse, ApplicationMasterProtocol>() {
			@Override
			public RegisterApplicationMasterResponse doInYarn(ApplicationMasterProtocol proxy) throws YarnException,
					IOException {
				RegisterApplicationMasterRequest appMasterRequest = Records
						.newRecord(RegisterApplicationMasterRequest.class);
				appMasterRequest.setHost(host != null ? host : "");
				appMasterRequest.setRpcPort(rpcPort != null ? rpcPort : 0);
				appMasterRequest.setTrackingUrl(trackUrl != null ? trackUrl : "");
				return proxy.registerApplicationMaster(appMasterRequest);
			}
		});
	}

	@Override
	public AllocateResponse allocate(final AllocateRequest request) {
		return execute(new YarnRpcCallback<AllocateResponse, ApplicationMasterProtocol>() {
			@Override
			public AllocateResponse doInYarn(ApplicationMasterProtocol proxy) throws YarnException, IOException {
				return proxy.allocate(request);
			}
		});
	}

	@Override
	public FinishApplicationMasterResponse finish(final FinishApplicationMasterRequest request) {
		return execute(new YarnRpcCallback<FinishApplicationMasterResponse, ApplicationMasterProtocol>() {
			@Override
			public FinishApplicationMasterResponse doInYarn(ApplicationMasterProtocol proxy) throws YarnException,
					IOException {
				return proxy.finishApplicationMaster(request);
			}
		});
	}

	@Override
	protected InetSocketAddress getRpcAddress(Configuration config) {
		InetSocketAddress addr = config.getSocketAddr(YarnConfiguration.RM_SCHEDULER_ADDRESS,
				YarnConfiguration.DEFAULT_RM_SCHEDULER_ADDRESS, YarnConfiguration.DEFAULT_RM_SCHEDULER_PORT);
		try {
			setupTokens(addr);
		} catch (IOException e) {
			log.error("Error setting up tokens", e);
		}
		return addr;
	}

	private static void setupTokens(InetSocketAddress resourceManagerAddress) throws IOException {
		// It is assumed for now that the only AMRMToken in AM's UGI is for this
		// cluster/RM. TODO: Fix later when we have some kind of cluster-ID as
		// default service-address, see YARN-986.
		for (Token<? extends TokenIdentifier> token : UserGroupInformation.getCurrentUser().getTokens()) {
			if (token.getKind().equals(AMRMTokenIdentifier.KIND_NAME)) {
				// This token needs to be directly provided to the AMs, so set
				// the appropriate service-name. We'll need more infrastructure when
				// we need to set it in HA case.
				SecurityUtil.setTokenService(token, resourceManagerAddress);
			}
		}
	}

}