/*
 * Copyright 2014 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.integration.websocket.support;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import org.springframework.messaging.Message;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.messaging.SubProtocolHandler;

/**
 * The utility class to encapsulate search algorithms for a set of provided {@link SubProtocolHandler}s.
 * <p>
 * For internal use only.
 *
 * @author Andy Wilkinson
 * @author Artem Bilan
 * @since 4.1
 * @see org.springframework.integration.websocket.inbound.WebSocketInboundChannelAdapter
 * @see org.springframework.integration.websocket.outbound.WebSocketOutboundMessageHandler
 */
public final class SubProtocolHandlerRegistry {

	private final static Log logger = LogFactory.getLog(SubProtocolHandlerRegistry.class);

	private final Map<String, SubProtocolHandler> protocolHandlers =
			new TreeMap<String, SubProtocolHandler>(String.CASE_INSENSITIVE_ORDER);

	private final SubProtocolHandler defaultProtocolHandler;

	public SubProtocolHandlerRegistry(List<SubProtocolHandler> protocolHandlers) {
		this(protocolHandlers, null);
	}

	public SubProtocolHandlerRegistry(SubProtocolHandler defaultProtocolHandler) {
		this(null, defaultProtocolHandler);
	}

	public SubProtocolHandlerRegistry(List<SubProtocolHandler> protocolHandlers,
			SubProtocolHandler defaultProtocolHandler) {
		Assert.state(!CollectionUtils.isEmpty(protocolHandlers) || defaultProtocolHandler != null,
				"One of 'protocolHandlers' or 'defaultProtocolHandler' must be provided");

		if (!CollectionUtils.isEmpty(protocolHandlers)) {
			for (SubProtocolHandler handler : protocolHandlers) {
				List<String> protocols = handler.getSupportedProtocols();
				if (CollectionUtils.isEmpty(protocols)) {
					logger.warn("No sub-protocols, ignoring handler " + handler);
					continue;
				}
				for (String protocol : protocols) {
					SubProtocolHandler replaced = this.protocolHandlers.put(protocol, handler);
					if (replaced != null) {
						throw new IllegalStateException("Failed to map handler " + handler
								+ " to protocol '" + protocol + "', it is already mapped to handler " + replaced);
					}
				}
			}
		}

		if (this.protocolHandlers.size() == 1 && defaultProtocolHandler == null) {
			this.defaultProtocolHandler = this.protocolHandlers.values().iterator().next();
		}
		else {
			this.defaultProtocolHandler = defaultProtocolHandler;
			if (this.protocolHandlers.isEmpty() && this.defaultProtocolHandler != null) {
				List<String> protocols = this.defaultProtocolHandler.getSupportedProtocols();
				for (String protocol : protocols) {
					SubProtocolHandler replaced = this.protocolHandlers.put(protocol, this.defaultProtocolHandler);
					if (replaced != null) {
						throw new IllegalStateException("Failed to map handler " + this.defaultProtocolHandler
								+ " to protocol '" + protocol + "', it is already mapped to handler " + replaced);
					}
				}
			}
		}
	}

	/**
	 * Resolves the {@link SubProtocolHandler} for the given {@code session} using
	 * its {@link WebSocketSession#getAcceptedProtocol() accepted sub-protocol}.
	 * @param session The session to resolve the sub-protocol handler for
	 * @return The sub-protocol handler
	 * @throws IllegalStateException if a protocol handler cannot be resolved
	 */
	public SubProtocolHandler findProtocolHandler(WebSocketSession session) {
		SubProtocolHandler handler;
		String protocol = session.getAcceptedProtocol();
		if (StringUtils.hasText(protocol)) {
			handler = this.protocolHandlers.get(protocol);
			Assert.state(handler != null,
					"No handler for sub-protocol '" + protocol + "', handlers = " + this.protocolHandlers);
		}
		else {
			handler = this.defaultProtocolHandler;
			Assert.state(handler != null,
					"No sub-protocol was requested and a default sub-protocol handler was not configured");
		}
		return handler;
	}

	/**
	 * Resolves the {@code sessionId} for the given {@code message} using
	 * the {@link SubProtocolHandler#resolveSessionId} algorithm.
	 * @param message The message to resolve the {@code sessionId} from.
	 * @return The sessionId or {@code null}, if no one {@link SubProtocolHandler}
	 * can't resolve it against provided {@code message}.
	 */
	public String resolveSessionId(Message<?> message) {
		for (SubProtocolHandler handler : this.protocolHandlers.values()) {
			String sessionId = handler.resolveSessionId(message);
			if (sessionId != null) {
				return sessionId;
			}
		}
		if (this.defaultProtocolHandler != null) {
			String sessionId = this.defaultProtocolHandler.resolveSessionId(message);
			if (sessionId != null) {
				return sessionId;
			}
		}
		return null;
	}

	/**
	 * Return the {@link List} of sub-protocols from provided {@link SubProtocolHandler}.
	 * @return The the {@link List} of supported sub-protocols.
	 */
	public List<String> getSubProtocols() {
		return new ArrayList<String>(this.protocolHandlers.keySet());
	}

}